import quax
import jax
import jax.tree as jt
import jax._src.core as core
import jax.extend.core as jexc
import jax.numpy as jnp
import numpy as np
import equinox as eqx
from jaxtyping import Float, Array, ArrayLike, DTypeLike, PyTree
from typing import Union, Sequence, TypeGuard

from jax.extend.core.primitives import dot_general_p, conv_general_dilated_p


def _make_default_autocast_config(dtype):
    return {
        dot_general_p: dtype,
        conv_general_dilated_p: dtype,
    }


def grampify(
    tree: PyTree,
    /,
    dtype: DTypeLike = jnp.bfloat16,
    config: dict[jexc.Primitive, DTypeLike] | None = None,
) -> PyTree:
    config = config or _make_default_autocast_config(dtype)
    autocast, rest = eqx.partition(tree, filter_spec=_is_nonquax_float_array)
    autocast = jt.map(
        eqx.Partial(AutocastArray, config=_AutocastConfig(config)), autocast
    )

    return eqx.combine(autocast, rest, is_leaf=_is_autocast_array)


def ungrampify(tree: PyTree) -> PyTree:
    return jt.map(
        lambda x: x.array if _is_autocast_array(x) else x,
        tree,
        is_leaf=_is_autocast_array,
    )


def _is_nonquax_float_array(x):
    return (
        eqx.is_array(x)
        and jnp.issubdtype(jnp.dtype(x), np.floating)
        and not isinstance(x, quax._core._QuaxTracer)
    )


# not registered as a pytree node. JAX expects dict keys to be sortable and Primitive isn't
class _AutocastConfig(dict[jexc.Primitive, DTypeLike]):
    def __or__(self, other: "_AutocastConfig") -> "_AutocastConfig":
        merged = self.copy()
        merged.update(
            {
                prim: jnp.result_type(self.get(prim, dtype), dtype)
                for prim, dtype in other.items()
            }
        )
        return _AutocastConfig(merged)

    def __hash__(self):
        return hash(frozenset(self.items()))


class AutocastArray(quax.ArrayValue):
    array: Float[Array, "..."]
    config: _AutocastConfig = eqx.field(static=True)

    def __init__(self, array, config):
        assert not isinstance(array, AutocastArray)
        self.array = array
        self.config = _AutocastConfig(config)

    def __check_init__(self):
        assert all(jnp.issubdtype(dtype, np.floating) for dtype in self.config.values())

    @staticmethod
    def merge_configs(*values: Sequence[ArrayLike | quax.Value]) -> _AutocastConfig:
        assert len(values) >= 1

        conf = None
        for val in values:
            if isinstance(val, AutocastArray):
                if conf is None:
                    conf = val.config
                else:
                    conf = conf | val.config

        assert conf is not None
        return conf

    @staticmethod
    def default(
        primitive: jexc.Primitive,
        values: Sequence[Union[ArrayLike | quax.Value]],
        params: dict,
    ) -> Union[ArrayLike, quax.Value, Sequence[Union[ArrayLike, quax.Value]]]:
        conf = AutocastArray.merge_configs(*values)
        print(values[0])

        values = ungrampify(values)

        def maybe_cast(v, dtype):
            if _is_nonquax_float_array(v):
                return v.astype(dtype)
            return v

        if primitive in conf:
            values = [
                v.materialise() if isinstance(v, quax.Value) else v for v in values
            ]
            _, out_struct, _ = eqx.filter_make_jaxpr(primitive.bind)(*values, **params)

            if "preferred_element_type" in params:
                params["preferred_element_type"] = conf[primitive]

            values = [maybe_cast(v, conf[primitive]) for v in values]
            out = primitive.bind(*values, **params)

            if primitive.multiple_results:
                out = [maybe_cast(v, s.dtype) for v, s in zip(out, out_struct)]
            else:
                out = maybe_cast(out, out_struct.dtype)

        else:
            out = quax.quaxify(primitive.bind)(*values, **params)

        return grampify(out, config=conf)

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

    def materialise(self) -> Float[Array, "..."]:
        raise ValueError("Refusing to materialise AutocastArray")


def _is_autocast_array(x) -> TypeGuard[AutocastArray]:
    return isinstance(x, AutocastArray)


def partition_autocast(tree):
    return eqx.partition(
        tree, filter_spec=_is_autocast_array, is_leaf=_is_autocast_array
    )


@quax.register(jax.lax.scan_p)
def _(
    *args: AutocastArray | ArrayLike,
    reverse: bool,
    length: int,
    jaxpr,
    num_consts: int,
    num_carry: int,
    linear,
    unroll: int = 1,
    _split_transpose: bool | None = None,
):
    in_flat, inputs_struct = jt.flatten(args)

    jax_fn = core.jaxpr_as_fun(jaxpr)

    def quax_fn(*args):
        return quax.quaxify(jax_fn)(*jt.unflatten(inputs_struct, args))

    trace_in = (
        *in_flat[:num_consts],
        *in_flat[num_consts : num_consts + num_carry],
        *[x[0] for x in in_flat[num_consts + num_carry :]],
    )

    quax_jaxpr = jax.make_jaxpr(lambda *x: jt.leaves(quax_fn(*x)))(*trace_in)

    out_flat = jax.lax.scan_p.bind(
        *in_flat,
        reverse=reverse,
        length=length,
        jaxpr=quax_jaxpr,
        num_consts=num_consts,
        num_carry=num_carry,
        linear=linear,
        unroll=unroll,
        _split_transpose=_split_transpose,
    )

    out_struct = jt.structure(quax_fn(*trace_in))
    return jt.unflatten(out_struct, out_flat)
