from __future__ import annotations

import jax
import jax.numpy as jnp
from dataclasses import dataclass
from typing import TYPE_CHECKING, TypeAlias

from .basis import Basis3D

# FormContext/ScalarFormField/VectorFormField were dataclasses with the default __eq__. JAX ended up calling that
#   __eq__ during the vmap over residuals, which tries to compare array fields element‑wise and then coerce to a
#   bool, triggering “truth value of an array is ambiguous.” Setting eq=False (and for ElementVector for consistency)
#   removes the autogenerated __eq__, so vmap no longer evaluates array equality and the residual/Jacobian assembly
#   succeeds.


@dataclass(eq=False)
class ElementVector:
    """
    Simple vector-valued element wrapper (scikit-fem ElementVector style).
    dim: dofs per node (e.g., 3 for displacement)
    """
    dim: int

    def dof_map(self, conn: jnp.ndarray) -> jnp.ndarray:
        """
        Expand scalar connectivity (n_elems, n_nodes_per_elem) to vector dofs.
        Returns shape (n_elems, n_nodes_per_elem * dim).
        """
        base = conn[..., None] * self.dim                     # (n_elems, n_nodes, 1)
        offsets = jnp.arange(self.dim, dtype=conn.dtype)      # (dim,)
        dofs = base + offsets                                 # (n_elems, n_nodes, dim)
        return dofs.reshape(conn.shape[0], -1)


@jax.tree_util.register_pytree_node_class
@dataclass(eq=False)
class ScalarFormField:
    """Scalar FE field evaluated on one element."""
    N: jnp.ndarray              # (n_q, n_nodes)
    elem_coords: jnp.ndarray    # (n_nodes, 3)
    basis: Basis3D
    _gradN: jnp.ndarray | None = None
    _detJ: jnp.ndarray | None = None

    @property
    def gradN(self):
        if self._gradN is None:
            self._gradN, self._detJ = self.basis.spatial_grads_and_detJ(self.elem_coords)
        return self._gradN

    @property
    def detJ(self):
        if self._detJ is None:
            self._gradN, self._detJ = self.basis.spatial_grads_and_detJ(self.elem_coords)
        return self._detJ

    def eval(self, u_elem: jnp.ndarray) -> jnp.ndarray:
        # u_elem: (n_nodes,)
        return jnp.einsum("qa,a->q", self.N, u_elem)

    def grad(self, u_elem: jnp.ndarray) -> jnp.ndarray:
        # returns (n_q, 3)
        return jnp.einsum("qaj,a->qj", self.gradN, u_elem)

    def tree_flatten(self):
        children = (self.N, self.elem_coords, self._gradN, self._detJ)
        aux = {"basis": self.basis}
        return children, aux

    @classmethod
    def tree_unflatten(cls, aux, chirdren):
        N, elem_coords, gradN, detJ = chirdren
        return cls(N, elem_coords, aux["basis"], gradN, detJ)


@jax.tree_util.register_pytree_node_class
@dataclass(eq=False)
class PrecomputedScalarFormField:
    """
    Scalar FE field with precomputed gradients/Jacobian determinant.
    Lightweight alternative that avoids storing element coordinates and
    element-wise duplicated shape-function tables.
    """
    basis: Basis3D
    _gradN: jnp.ndarray         # (n_q, n_nodes, 3)
    _detJ: jnp.ndarray          # (n_q,)

    @property
    def N(self):
        return self.basis.shape_functions()

    @property
    def gradN(self):
        return self._gradN

    @property
    def detJ(self):
        return self._detJ

    def eval(self, u_elem: jnp.ndarray) -> jnp.ndarray:
        return jnp.einsum("qa,a->q", self.N, u_elem)

    def grad(self, u_elem: jnp.ndarray) -> jnp.ndarray:
        return jnp.einsum("qaj,a->qj", self.gradN, u_elem)

    def tree_flatten(self):
        children = (self._gradN, self._detJ)
        aux = {"basis": self.basis}
        return children, aux

    @classmethod
    def tree_unflatten(cls, aux, children):
        gradN, detJ = children
        return cls(aux["basis"], gradN, detJ)


@jax.tree_util.register_pytree_node_class
@dataclass(eq=False)
class VectorFormField:
    """Vector-valued FE field evaluated on one element."""
    N: jnp.ndarray
    elem_coords: jnp.ndarray
    basis: Basis3D
    value_dim: int  # ★Python int (static)
    _gradN: jnp.ndarray | None = None
    _detJ: jnp.ndarray | None = None

    @property
    def gradN(self):
        if self._gradN is None:
            self._gradN, self._detJ = self.basis.spatial_grads_and_detJ(self.elem_coords)
        return self._gradN

    @property
    def detJ(self):
        if self._detJ is None:
            self._gradN, self._detJ = self.basis.spatial_grads_and_detJ(self.elem_coords)
        return self._detJ

    def eval(self, u_elem: jnp.ndarray) -> jnp.ndarray:
        # u_elem: (value_dim*n_nodes,) expected
        u_nodes = u_elem.reshape((-1, self.value_dim))  # (n_nodes, vd); vd is a Python int
        return jnp.einsum("qa,ai->qi", self.N, u_nodes)  # (n_q, vd)

    def grad(self, u_elem: jnp.ndarray) -> jnp.ndarray:
        u_nodes = u_elem.reshape((-1, self.value_dim))
        return jnp.einsum("qaj,ai->qij", self.gradN, u_nodes)  # (n_q, vd, 3)

    def tree_flatten(self):
        children = (self.N, self.elem_coords, self._gradN, self._detJ)
        aux = {
            "basis": self.basis,
            "value_dim": int(self.value_dim)
        }
        return children, aux

    @classmethod
    def tree_unflatten(cls, aux, children):
        N, elem_coords, gradN, detJ = children
        return cls(N, elem_coords, aux["basis"], aux["value_dim"], gradN, detJ)


@jax.tree_util.register_pytree_node_class
@dataclass(eq=False)
class PrecomputedVectorFormField:
    """
    Vector FE field with precomputed gradients/Jacobian determinant.
    Lightweight alternative that avoids storing element coordinates/basis.
    """
    basis: Basis3D
    value_dim: int
    _gradN: jnp.ndarray
    _detJ: jnp.ndarray

    @property
    def N(self):
        return self.basis.shape_functions()

    @property
    def gradN(self):
        return self._gradN

    @property
    def detJ(self):
        return self._detJ

    def eval(self, u_elem: jnp.ndarray) -> jnp.ndarray:
        u_nodes = u_elem.reshape((-1, int(self.value_dim)))
        return jnp.einsum("qa,ai->qi", self.N, u_nodes)

    def grad(self, u_elem: jnp.ndarray) -> jnp.ndarray:
        u_nodes = u_elem.reshape((-1, int(self.value_dim)))
        return jnp.einsum("qaj,ai->qij", self.gradN, u_nodes)

    def tree_flatten(self):
        children = (self._gradN, self._detJ)
        aux = {"basis": self.basis, "value_dim": int(self.value_dim)}
        return children, aux

    @classmethod
    def tree_unflatten(cls, aux, children):
        gradN, detJ = children
        return cls(aux["basis"], aux["value_dim"], gradN, detJ)


if TYPE_CHECKING:
    FormFieldLike: TypeAlias = (
        ScalarFormField
        | VectorFormField
        | PrecomputedScalarFormField
        | PrecomputedVectorFormField
    )
else:
    FormFieldLike = object


def vector_load_form(field: FormFieldLike, load_vec: jnp.ndarray) -> jnp.ndarray:
    """
    Build vector linear form values from a FormField and a load vector.
    """
    lv = jnp.asarray(load_vec)
    if lv.ndim == 1:
        lv = lv[None, :]
    elif lv.ndim != 2:
        raise ValueError("load_vec must be shape (dim,) or (n_q, dim)")
    if lv.shape[0] == 1:
        lv = jnp.broadcast_to(lv, (field.N.shape[0], lv.shape[1]))
    elif lv.shape[0] != field.N.shape[0]:
        raise ValueError("load_vec must be shape (dim,) or (n_q, dim)")
    load = field.N[..., None] * lv[:, None, :]
    return load.reshape(load.shape[0], -1)


@jax.tree_util.register_pytree_node_class
@dataclass(eq=False)
class FormContext:
    """Bundle test/trial fields and quadrature data for element assembly."""
    test: FormFieldLike
    trial: FormFieldLike
    x_q: jnp.ndarray       # (n_q, 3)
    w: jnp.ndarray         # (n_q,)
    elem_id: jnp.ndarray | int = 0

    @property
    def u(self) -> FormFieldLike:
        return self.trial

    @property
    def v(self) -> FormFieldLike:
        return self.test

    def tree_flatten(self):
        children = (
            self.test,
            self.trial,
            self.x_q,
            self.w,
            self.elem_id,
        )
        return children, {}

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        (
            test,
            trial,
            x_q,
            w,
            elem_id,
        ) = children
        return cls(
            test,
            trial,
            x_q,
            w,
            elem_id,
        )


@jax.tree_util.register_pytree_node_class
@dataclass(eq=False)
class FieldPair:
    """Named test/trial/unknown grouping for mixed formulations."""
    test: FormFieldLike
    trial: FormFieldLike
    unknown: FormFieldLike | None = None

    def tree_flatten(self):
        children = (self.test, self.trial, self.unknown)
        return children, {}

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        test, trial, unknown = children
        return cls(test=test, trial=trial, unknown=unknown)


@jax.tree_util.register_pytree_node_class
@dataclass(eq=False)
class MixedFormContext:
    """FormContext for mixed formulations keyed by field name."""
    fields: dict[str, FieldPair]
    x_q: jnp.ndarray       # (n_q, 3)
    w: jnp.ndarray         # (n_q,)
    elem_id: jnp.ndarray | int = 0
    unknown: FormFieldLike | None = None
    trial_fields: dict[str, FormFieldLike] | None = None
    test_fields: dict[str, FormFieldLike] | None = None
    unknown_fields: dict[str, FormFieldLike] | None = None

    def tree_flatten(self):
        children = (
            self.fields,
            self.x_q,
            self.w,
            self.elem_id,
            self.unknown,
            self.trial_fields,
            self.test_fields,
            self.unknown_fields,
        )
        return children, {}

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        (
            fields,
            x_q,
            w,
            elem_id,
            unknown,
            trial_fields,
            test_fields,
            unknown_fields,
        ) = children
        return cls(fields, x_q, w, elem_id, unknown, trial_fields, test_fields, unknown_fields)
