"""Reusable Flax model classes for HiPPO-based sequence models."""
from __future__ import annotations

import jax
import jax.numpy as jnp
import flax.linen as nn

__all__ = ["AssocMemHiPPO", "MLPHiPPO", "SalienceHiPPO", "VanillaHiPPO", "legendre_orthonormal_basis01"]


def legendre_orthonormal_basis01(x: jnp.ndarray, n: int) -> jnp.ndarray:
    """
    Orthonormal Legendre basis on [0,1] with:
      p_k(x) = sqrt(2k+1) * (-1)^k * P_k(2x-1)

    x: scalar or array in [0,1]
    returns: array with shape x.shape + (n,)
    """
    x = jnp.asarray(x)
    t = 2.0 * x - 1.0

    out_shape = t.shape + (n,)
    P = jnp.zeros(out_shape, dtype=t.dtype)

    P = P.at[..., 0].set(jnp.ones_like(t))

    if n > 1:
        P = P.at[..., 1].set(t)

        def body(k, P):
            kk = jnp.asarray(k, dtype=t.dtype)
            Pkm1 = P[..., k - 1]  # P_{k-1}
            Pk = P[..., k]        # P_k
            Pkp1 = ((2.0 * kk + 1.0) * t * Pk - kk * Pkm1) / (kk + 1.0)
            return P.at[..., k + 1].set(Pkp1)

        P = jax.lax.fori_loop(1, n - 1, body, P)

    # Orthonormalization/sign convention
    idx = jnp.arange(n, dtype=t.dtype)
    scale = jnp.sqrt(2.0 * idx + 1.0) * ((-1.0) ** idx)  # (n,)
    return P * scale  # broadcast over leading dims


class VanillaHiPPO(nn.Module):
    """
    Minimal HiPPO step model with fixed ZOH discretization.

    Each step:
      x_proj = W_in @ x_in                          # (d_in,) -> (d_model,)
      S_next[j] = A_d @ S[j] + b_d * x_proj[j]     # per channel, vectorized
      y_hat  = W_out @ S_next.reshape(-1)            # (d_model * n,) -> (d_in,)

    State layout:
      S: (d_model, n)  — one HiPPO state per model dimension

    Args:
      d_in:    input / output token dimension
      d_model: number of parallel HiPPO channels
      n:       HiPPO state order (polynomial degree)
      A_d:     (n, n) pre-discretized state-transition matrix
      b_d:     (n,)   pre-discretized input vector

    Construct A_d and b_d via ``get_system_params`` + a ZOH step, e.g.::

        (A, b), _, _ = get_system_params("legs", n)
        A_d = jnp.array(scipy.linalg.expm(A * dt))
        b_d = jnp.linalg.solve(A, (A_d - I) @ b)
    """
    d_in: int
    d_model: int
    n: int

    A_d: jnp.ndarray  # (n, n)
    b_d: jnp.ndarray  # (n,)

    @nn.compact
    def __call__(self, x_in: jnp.ndarray, S: jnp.ndarray):
        """
        x_in:   (d_in,)
        S:      (d_model, n)
        returns:
          S_next: (d_model, n)
          y_hat:  (d_in,)
        """
        x_proj = nn.Dense(self.d_model, use_bias=False, name="W_in")(x_in)  # (d_model,)

        def update_channel(s_j, u_j):
            return self.A_d @ s_j + self.b_d * u_j

        S_next = jax.vmap(update_channel, in_axes=(0, 0), out_axes=0)(S, x_proj)  # (d_model, n)

        y_hat = nn.Dense(self.d_in, use_bias=True, name="W_out")(
            S_next.reshape((self.d_model * self.n,))
        )  # (d_in,)
        return S_next, y_hat


class MLPHiPPO(nn.Module):
    """
    HiPPO step model identical to VanillaHiPPO but with a single-hidden-layer
    MLP readout instead of a linear map.

    Each step:
      x_proj = W_in @ x_in                          # (d_in,) -> (d_model,)
      S_next[j] = A_d @ S[j] + b_d * x_proj[j]     # per channel, vectorized
      h      = activation(W_h @ S_next.reshape(-1)) # (d_model * n,) -> (mlp_hidden,)
      y_hat  = W_out @ h                             # (mlp_hidden,) -> (d_in,)

    Args:
      d_in:       input / output token dimension
      d_model:    number of parallel HiPPO channels
      n:          HiPPO state order (polynomial degree)
      A_d:        (n, n) pre-discretized state-transition matrix
      b_d:        (n,)   pre-discretized input vector
      mlp_hidden: hidden layer width (default: d_model * n)
    """
    d_in: int
    d_model: int
    n: int

    A_d: jnp.ndarray  # (n, n)
    b_d: jnp.ndarray  # (n,)

    mlp_hidden: int = 0  # 0 means use d_model * n

    @nn.compact
    def __call__(self, x_in: jnp.ndarray, S: jnp.ndarray):
        """
        x_in:   (d_in,)
        S:      (d_model, n)
        returns:
          S_next: (d_model, n)
          y_hat:  (d_in,)
        """
        x_proj = nn.Dense(self.d_model, use_bias=False, name="W_in")(x_in)  # (d_model,)

        def update_channel(s_j, u_j):
            return self.A_d @ s_j + self.b_d * u_j

        S_next = jax.vmap(update_channel, in_axes=(0, 0), out_axes=0)(S, x_proj)  # (d_model, n)

        hidden = self.mlp_hidden if self.mlp_hidden > 0 else self.d_model * self.n
        h = nn.Dense(hidden, use_bias=True, name="W_h")(
            S_next.reshape((self.d_model * self.n,))
        )  # (hidden,)
        h = nn.softplus(h)
        y_hat = nn.Dense(self.d_in, use_bias=True, name="W_out")(h)  # (d_in,)
        return S_next, y_hat


class SalienceHiPPO(nn.Module):
    """
    Single-example step model with ZOH discretization using:
      expm(gA) ~= expm(g0 A) @ expm(r A)

    State layout:
      S: (d_model, n)  (one HiPPO state per model dimension)

    Notes:
      A and b should be passed as JAX arrays (not numpy arrays) at construction
      time, as Flax treats them as static module attributes.
    """
    d_in: int
    d_model: int
    n: int

    A: jnp.ndarray  # (n, n)
    b: jnp.ndarray  # (n,)

    g_max: float = 5.0
    sal_hidden: int = 128
    mem_dim: int = 32

    outvec_hidden: int = 128

    num_grid: int = 256
    taylor_order: int = 4

    def setup(self):
        self.I_n = jnp.eye(self.n, dtype=self.A.dtype)

        self.dg = jnp.asarray(self.g_max / self.num_grid, dtype=self.A.dtype)
        g_grid = jnp.linspace(0.0, self.g_max, self.num_grid + 1, dtype=self.A.dtype)

        A_grid = jax.vmap(lambda gg: jax.scipy.linalg.expm(gg * self.A))(g_grid)

        def B_from_A_d(A_d):
            rhs = (A_d - self.I_n) @ self.b
            return jnp.linalg.solve(self.A, rhs)

        B_grid = jax.vmap(B_from_A_d)(A_grid)

        self.A_grid = A_grid
        self.B_grid = B_grid

    def _expm_taylor(self, r: jnp.ndarray) -> jnp.ndarray:
        rA = r * self.A
        I = self.I_n

        def body(k, carry):
            E, term = carry
            term = (term @ rA) / jnp.asarray(k, dtype=I.dtype)
            E = E + term
            return (E, term)

        E0 = I
        term0 = I
        E, _ = jax.lax.fori_loop(1, self.taylor_order + 1, body, (E0, term0))
        return E

    def _Brem_taylor(self, r: jnp.ndarray) -> jnp.ndarray:
        dtype = self.b.dtype
        b_rem = jnp.zeros((self.n,), dtype=dtype)
        Akb = self.b

        def body(k, carry):
            b_rem, Akb = carry
            kk = jnp.asarray(k, dtype=dtype)
            # coeff = r^{k+1}/(k+1)!
            coeff = (r ** (kk + 1.0)) / jnp.exp(jax.scipy.special.gammaln(kk + 2.0))
            b_rem = b_rem + coeff * Akb
            Akb = self.A @ Akb
            return (b_rem, Akb)

        b_rem, _ = jax.lax.fori_loop(0, self.taylor_order + 1, body, (b_rem, Akb))
        return b_rem

    def _zoh_discretize(self, g: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
        g = jnp.clip(g, 0.0, self.g_max)
        idx = jnp.minimum(jnp.floor(g / self.dg).astype(jnp.int32), jnp.int32(self.num_grid))
        g0 = self.dg * idx.astype(self.A.dtype)
        r = g - g0

        A0 = self.A_grid[idx]
        B0 = self.B_grid[idx]

        E = self._expm_taylor(r)
        B_r = self._Brem_taylor(r)

        A_d = A0 @ E
        b_d = B0 + A0 @ B_r
        return A_d, b_d

    @nn.compact
    def __call__(self, x_in: jnp.ndarray, S: jnp.ndarray):
        """
        x_in: (d_in,)
        S:    (d_model, n)
        returns:
          S_next: (d_model, n)
          y_hat:  (d_in,)
          g:      scalar
          out_vec: (n,)
        """
        # token projection
        x_proj = nn.Dense(self.d_model, use_bias=False, name="W_in")(x_in)

        # salience
        ms = nn.Dense(self.mem_dim, name="sal_pool")(S)
        ms = nn.tanh(ms)
        ms = jnp.mean(ms, axis=0)

        sal_inp = jnp.concatenate([x_proj, ms], axis=0)
        h = nn.Dense(self.sal_hidden, name="sal_fc1")(sal_inp)
        h = nn.softplus(h)
        h = nn.Dense(1, name="sal_fc3")(h)
        g = self.g_max * nn.sigmoid(h[0])

        A_d, b_d = self._zoh_discretize(g)

        def update_channel(s_n, u_j):
            return (A_d @ s_n) + (b_d * u_j)

        S_next = jax.vmap(update_channel, in_axes=(0, 0), out_axes=0)(S, x_proj)

        S_flat = S_next.reshape((self.d_model * self.n,))
        h = nn.Dense(self.outvec_hidden, use_bias=True, name="W_z_1")(S_flat)
        h = nn.softplus(h)
        out_vec = nn.Dense(self.n, use_bias=True, name="W_z_2")(h)

        y_mid = jnp.einsum("ij,j->i", S_next, out_vec)  # (d_model,)
        y_hat = nn.Dense(self.d_in, use_bias=True, name="out_proj")(y_mid)
        return S_next, y_hat, g, out_vec


class AssocMemHiPPO(nn.Module):
    """
    Fixed-ZOH HiPPO + banked continuous-time associative memory with:
      1) Key/query vectors from the SAME linear map of per-channel HiPPO state
      2) Explicit write gate g_write in [0,1] (interpolates between no write and full write)
      3) Query-based readout: q -> r = C @ q (per bank) -> linear to d_in

    State:
      S: (d_model, n_hippo)
      C: (d_model, n_assoc)

    HiPPO step (fixed dt=1 ZOH):
      S_next[j] = A_d @ S[j] + b_d * x_proj[j]

    Key/query maps (shared across channels):
      k_j = W_k @ S_next[j]         in R^{n_assoc}
      q_j = W_q @ S_next[j]         in R^{n_assoc}

    Memory write (per bank j), with dt=1 exact update scaled by gate:
      proj_j = <k_j, C_j>
      err_j  = y_j - proj_j
      C_j <- C_j + g_write / ||k_j||^2 * err_j * k_j

    where:
      - g_write is scalar in [0,1] from a small MLP on global state
      - y_j is value signal per bank from a small MLP on global state

    Readout:
      r_j = <q_j, C_j>     (scalar per bank)  -> r in R^{d_model}
      y_hat = W_out @ r    in R^{d_in}
    """
    d_in: int
    d_model: int
    n_hippo: int
    n_assoc: int

    A_d: jnp.ndarray  # (n_hippo, n_hippo)
    b_d: jnp.ndarray  # (n_hippo,)

    write_hidden: int
    out_hidden: int

    @nn.compact
    def __call__(self, x_in: jnp.ndarray, S: jnp.ndarray, C: jnp.ndarray):
        """
        x_in: (d_in,)
        S:    (d_model, n_hippo)
        C:    (d_model, n_assoc)  # OP coefficients per bank

        Returns:
        S_next: (d_model, n_hippo)
        C_next: (d_model, n_assoc)
        y_hat:  (d_in,)
        aux: dict(g_write, g_out, y_vec, x_key, x_query)
        """
        # ---- (1) token projection to channels
        x_proj = nn.Dense(self.d_model, use_bias=False, name="W_in")(x_in)  # (d_model,)

        # ---- (2) fixed-ZOH HiPPO evolution
        def update_channel(s_n, u_j):
            return self.A_d @ s_n + self.b_d * u_j

        S_next = jax.vmap(update_channel, in_axes=(0, 0), out_axes=0)(S, x_proj)  # (d_model,n_hippo)

        # ---- (3) write and out gate MLPs
        S_flat = S_next.reshape((self.d_model * self.n_hippo,))

        g_write = nn.Dense(self.write_hidden)(S_flat)
        g_write = nn.tanh(g_write)
        g_write = nn.sigmoid(nn.Dense(1)(g_write) + nn.Dense(1)(S_flat))[0]

        g_out = nn.Dense(self.out_hidden)(S_flat)
        g_out = nn.tanh(g_out)
        g_out = nn.sigmoid(nn.Dense(1)(g_out) + nn.Dense(1)(S_flat))[0]

        # ---- (4) value per bank y_vec
        # (d_model,) "value" to store in each bank
        y_vec = nn.Dense(self.d_model, name="val_out")(x_proj)

        # ---- (5) Learn OP key/query *locations* in [0,1]
        # These are the actual "x values" for orthogonal polynomial evaluation.
        x_key   = nn.sigmoid(nn.Dense(1, name="x_key")(S_flat)[0])       # scalar in (0,1)
        x_query = nn.sigmoid(nn.Dense(1, name="x_query")(S_flat)[0])     # scalar in (0,1)

        # Build orthonormal Legendre basis vectors
        K = legendre_orthonormal_basis01(x_key, self.n_assoc)       # (n_assoc,)
        Q = legendre_orthonormal_basis01(x_query, self.n_assoc)     # (n_assoc,)

        # ---- (6) exact associative memory write, scaled by gate
        # proj_j = <C_j, K>
        proj = C @ K                                                # (d_model,)
        err  = y_vec - proj                                         # (d_model,)

        # ||K||^2 (scalar)
        K_norm2 = jnp.sum(K * K) + 1e-8
        gain = g_write / K_norm2

        C_next = C + gain * err[:, None] * K[None, :]               # (d_model, n_assoc)

        # ---- (7) query-based readout (OP evaluation)
        r = C_next @ Q                                              # (d_model,)
        y_hat = nn.Dense(self.d_in, use_bias=True, name="out_proj")(r)  # (d_in,)
        y_hat = y_hat * g_out

        aux = dict(
            g_write=g_write,
            g_out=g_out,
            y_vec=y_vec,
            x_key=x_key,
            x_query=x_query,
            K_norm2=K_norm2,
        )
        return S_next, C_next, y_hat, aux