import jax
import jax.numpy as jnp
import quax
import jax.core as core

class ArrayWrapper(quax.ArrayValue):
    array: jax.Array

    def aval(self):
        return core.ShapedArray(self.array.shape, self.array.dtype)

    def materialise(self) -> jax.Array:
        raise ValueError("Refusing to materialise")
        
@quax.register(jax.lax.while_p)
def _(*args, cond_nconsts: int, cond_jaxpr, body_nconsts: int, body_jaxpr):
    new_args = [a.array if isinstance(a, ArrayWrapper) else a for a in args]
    is_quaxed = [isinstance(a, ArrayWrapper) for a in args]
    out = jax.lax.while_p.bind(
        *new_args, 
        cond_nconsts=cond_nconsts,
        cond_jaxpr=cond_jaxpr,
        body_nconsts=body_nconsts,
        body_jaxpr=body_jaxpr,
    )
    return [
        ArrayWrapper(a) if quaxed else a
        for a, quaxed in zip(out, is_quaxed[cond_nconsts+body_nconsts:])
    ]

def body_fn(a: jax.Array):
    return a + 1

def cond_fn(a: jax.Array):
    return a[1] < 10

def loop_fn(a: jax.Array):
    res = jax.lax.while_loop(
        body_fun=body_fn,
        cond_fun=cond_fn,
        init_val=a,
    )
    return res

a = ArrayWrapper(jnp.arange(10))
a = jax.jit(quax.quaxify(loop_fn))(a)
print(a.array)