"""Sequence Mixer for S5 (Simplified State Space Layers).

S5 implementation modified from: https://github.com/lindermanlab/S5/blob/main/s5/ssm_init.py

See: https://openreview.net/pdf?id=Ai8Hw3AXqks
"""

from __future__ import annotations

from typing import Literal

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
from jax.nn.initializers import lecun_normal, normal
from jax.scipy.linalg import block_diag
from jaxtyping import Array, PRNGKeyArray

from discretax.sequence_mixers.base import AbstractSequenceMixer


class S5SequenceMixer(AbstractSequenceMixer):
    """S5 sequence mixer layer.

    This layer implements the Simplified State Space Layers (S5) sequence mixer,
    which uses structured state space models with HiPPO initialization and efficient
    parallel scan operations.

    Attributes:
        Lambda_re: Real part of diagonal state matrix eigenvalues.
        Lambda_im: Imaginary part of diagonal state matrix eigenvalues.
        B: Input projection matrix (parameterized as V^{-1}B).
        C: Output projection matrix (parameterized as CV).
        D: Skip connection weights.
        log_step: Log of discretization step sizes.
        in_features: Number of hidden channels (input features).
        P: Effective state dimensionality.
        conj_sym: Whether conjugate symmetry is enforced.
        clip_eigs: Whether to clip eigenvalues for stability.
        discretization: Discretization method being used.
        step_rescale: Rescaling factor for step sizes.
    """

    Lambda_re: jax.Array
    Lambda_im: jax.Array
    B: jax.Array
    C: jax.Array
    D: jax.Array
    log_step: jax.Array

    in_features: int
    P: int
    conj_sym: bool = eqx.field(static=True)
    clip_eigs: bool = eqx.field(static=True)
    discretization: str = eqx.field(static=True)
    step_rescale: float = eqx.field(static=True)

    def __init__(
        self,
        in_features: int,
        key: PRNGKeyArray,
        *args,
        state_dim: int = 64,
        ssm_blocks: int = 1,
        C_init: Literal[
            "trunc_standard_normal", "lecun_normal", "complex_normal"
        ] = "lecun_normal",
        conj_sym: bool = True,
        clip_eigs: bool = True,
        discretization: Literal["zoh", "bilinear"] = "zoh",
        dt_min: float = 0.001,
        dt_max: float = 1.0,
        step_rescale: float = 1.0,
        **kwargs,
    ):
        """Initialize the S5 sequence mixer layer.

        Args:
            in_features: dimension of the input features.
            key: JAX random key for initialization.
            state_dim: dimension of the state space (total SSM size).
            ssm_blocks: number of blocks in the block-diagonal HiPPO initialization.
            C_init: initialization method for output matrix C.
            conj_sym: whether to enforce conjugate symmetry (reduces parameters by half).
            clip_eigs: whether to clip eigenvalues to ensure stability.
            discretization: discretization method to use.
            dt_min: minimum discretization step size.
            dt_max: maximum discretization step size.
            step_rescale: rescaling factor for the discretization step.
            *args: Additional positional arguments (ignored).
            **kwargs: Additional keyword arguments (ignored).
        """
        B_key, C_key, D_key, step_key, key = jr.split(key, 5)

        ssm_size = state_dim
        blocks = ssm_blocks

        block_size = int(ssm_size / blocks)
        # Initialize state matrix A using approximation to HiPPO-LegS matrix
        Lambda, _, B, V, B_orig = _make_dplr_hippo(block_size)

        if conj_sym:
            block_size = block_size // 2
            P = ssm_size // 2
        else:
            P = ssm_size

        Lambda = Lambda[:block_size]
        V = V[:, :block_size]
        Vc = V.conj().T

        # If initializing state matrix A as block-diagonal, put HiPPO approximation
        # on each block
        Lambda = (Lambda * jnp.ones((blocks, block_size))).ravel()
        V = block_diag(*([V] * blocks))
        Vinv = block_diag(*([Vc] * blocks))

        self.in_features = in_features
        self.P = P
        if conj_sym:
            local_P = 2 * P
        else:
            local_P = P

        self.Lambda_re = Lambda.real
        self.Lambda_im = Lambda.imag

        self.conj_sym = conj_sym

        self.clip_eigs = clip_eigs

        self.B = _init_vinb(lecun_normal(), B_key, (local_P, self.in_features), Vinv)

        # Initialize state to output (C) matrix
        if C_init in ["trunc_standard_normal"]:
            C_init_fn = _trunc_standard_normal
        elif C_init in ["lecun_normal"]:
            C_init_fn = lecun_normal()
        elif C_init in ["complex_normal"]:
            C_init_fn = normal(stddev=0.5**0.5)
        else:
            raise NotImplementedError(f"C_init method {C_init} not implemented")

        if C_init in ["complex_normal"]:
            self.C = C_init_fn(C_key, (self.in_features, 2 * self.P, 2))
        else:
            self.C = _init_cv(C_init_fn, C_key, (self.in_features, local_P, 2), V)

        self.D = normal(stddev=1.0)(D_key, (self.in_features,))

        # Initialize learnable discretization timescale value
        self.log_step = _init_log_steps(step_key, (self.P, dt_min, dt_max))

        self.step_rescale = step_rescale
        self.discretization = discretization

    def __call__(self, x: Array, key: PRNGKeyArray) -> Array:
        """Forward pass of the S5 sequence mixer layer.

        Args:
            x: Input sequence of features.
            key: JAX random key (unused, for compatibility).

        Returns:
            The output of the S5 sequence mixer.
        """
        if self.clip_eigs:
            Lambda = jnp.clip(self.Lambda_re, None, -1e-4) + 1j * self.Lambda_im
        else:
            Lambda = self.Lambda_re + 1j * self.Lambda_im

        B_tilde = self.B[..., 0] + 1j * self.B[..., 1]
        C_tilde = self.C[..., 0] + 1j * self.C[..., 1]

        step = self.step_rescale * jnp.exp(self.log_step[:, 0])

        # Discretize
        if self.discretization in ["zoh"]:
            Lambda_bar, B_bar = _discretize_zoh(Lambda, B_tilde, step)
        elif self.discretization in ["bilinear"]:
            Lambda_bar, B_bar = _discretize_bilinear(Lambda, B_tilde, step)
        else:
            raise NotImplementedError(
                f"Discretization method {self.discretization} not implemented"
            )

        ys = _apply_ssm(Lambda_bar, B_bar, C_tilde, x, self.conj_sym)

        # Add feedthrough matrix output Du
        Du = jax.vmap(lambda u: self.D * u)(x)
        return ys + Du


def _make_hippo(N: int) -> Array:
    """Create a HiPPO-LegS matrix.

    From https://github.com/srush/annotated-s4/blob/main/s4/s4.py

    Args:
        N: State size.

    Returns:
        N x N HiPPO LegS matrix.
    """
    P = jnp.sqrt(1 + 2 * jnp.arange(N))
    A = P[:, jnp.newaxis] * P[jnp.newaxis, :]
    A = jnp.tril(A) - jnp.diag(jnp.arange(N))
    return -A


def _make_nplr_hippo(N: int) -> tuple[Array, Array, Array]:
    """Make components needed for NPLR representation of HiPPO-LegS.

    From https://github.com/srush/annotated-s4/blob/main/s4/s4.py

    Args:
        N: State size.

    Returns:
        Tuple of (HiPPO matrix, low-rank factor P, input matrix B).
    """
    # Make -HiPPO
    hippo = _make_hippo(N)

    # Add in a rank 1 term. Makes it Normal.
    P = jnp.sqrt(jnp.arange(N) + 0.5)

    # HiPPO also specifies the B matrix
    B = jnp.sqrt(2 * jnp.arange(N) + 1.0)
    return hippo, P, B


def _make_dplr_hippo(N: int) -> tuple[Array, Array, Array, Array, Array]:
    """Make components needed for DPLR representation of HiPPO-LegS.

    From https://github.com/srush/annotated-s4/blob/main/s4/s4.py
    Note: we will only use the diagonal part.

    Args:
        N: State size.

    Returns:
        Tuple of (eigenvalues Lambda, low-rank term P, conjugated input matrix B,
        eigenvectors V, original B matrix).
    """
    A, P, B = _make_nplr_hippo(N)

    S = A + P[:, jnp.newaxis] * P[jnp.newaxis, :]

    S_diag = jnp.diagonal(S)
    Lambda_real = jnp.mean(S_diag) * jnp.ones_like(S_diag)

    # Diagonalize S to V \Lambda V^*
    Lambda_imag, V = jnp.linalg.eigh(S * -1j)

    P = V.conj().T @ P
    B_orig = B
    B = V.conj().T @ B
    return Lambda_real + 1j * Lambda_imag, P, B, V, B_orig


def _log_step_initializer(dt_min: float = 0.001, dt_max: float = 0.1):
    """Initialize the learnable timescale Delta by sampling uniformly.

    Args:
        dt_min: Minimum discretization step value.
        dt_max: Maximum discretization step value.

    Returns:
        Initialization function that samples log_step uniformly.
    """

    def init(key: PRNGKeyArray, shape: tuple) -> Array:
        """Initialize log_step values.

        Args:
            key: JAX random key.
            shape: Desired shape.

        Returns:
            Sampled log_step values.
        """
        return jr.uniform(key, shape) * (jnp.log(dt_max) - jnp.log(dt_min)) + jnp.log(dt_min)

    return init


def _init_log_steps(key: PRNGKeyArray, input: tuple) -> Array:
    """Initialize an array of learnable timescale parameters.

    Args:
        key: JAX random key.
        input: Tuple containing (dim, dt_min, dt_max) where dim is the shape.

    Returns:
        Initialized array of timescales of shape (dim,).
    """
    dim, dt_min, dt_max = input
    log_steps = []
    for i in range(dim):
        key, skey = jr.split(key)
        log_step = _log_step_initializer(dt_min=dt_min, dt_max=dt_max)(skey, shape=(1,))
        log_steps.append(log_step)

    return jnp.array(log_steps)


def _init_vinb(init_fun, rng: PRNGKeyArray, shape: tuple, Vinv: Array) -> Array:
    """Initialize B_tilde = V^{-1}B.

    First samples B, then computes V^{-1}B. Parameterized with two different
    matrices for complex numbers.

    Args:
        init_fun: Initialization function to use (e.g. lecun_normal()).
        rng: JAX random key.
        shape: Desired shape (P, H).
        Vinv: Inverse eigenvectors for initialization.

    Returns:
        B_tilde of shape (P, H, 2) for complex parameterization.
    """
    B = init_fun(rng, shape)
    VinvB = Vinv @ B
    VinvB_real = VinvB.real
    VinvB_imag = VinvB.imag
    return jnp.concatenate((VinvB_real[..., None], VinvB_imag[..., None]), axis=-1)


def _trunc_standard_normal(key: PRNGKeyArray, shape: tuple) -> Array:
    """Sample C with a truncated normal distribution.

    Args:
        key: JAX random key.
        shape: Desired shape of length 3, (H, P, _).

    Returns:
        Sampled C matrix of shape (H, P, 2) for complex parameterization.
    """
    H, P, _ = shape
    Cs = []
    for i in range(H):
        key, skey = jr.split(key)
        C = lecun_normal()(skey, shape=(1, P, 2))
        Cs.append(C)
    return jnp.array(Cs)[:, 0]


def _init_cv(init_fun, rng: PRNGKeyArray, shape: tuple, V: Array) -> Array:
    """Initialize C_tilde = CV.

    First samples C, then computes CV. Parameterized with two different
    matrices for complex numbers.

    Args:
        init_fun: Initialization function to use (e.g. lecun_normal()).
        rng: JAX random key.
        shape: Desired shape (H, P).
        V: Eigenvectors for initialization.

    Returns:
        C_tilde of shape (H, P, 2) for complex parameterization.
    """
    C_ = init_fun(rng, shape)
    C = C_[..., 0] + 1j * C_[..., 1]
    CV = C @ V
    CV_real = CV.real
    CV_imag = CV.imag
    return jnp.concatenate((CV_real[..., None], CV_imag[..., None]), axis=-1)


def _discretize_bilinear(Lambda: Array, B_tilde: Array, Delta: Array) -> tuple[Array, Array]:
    """Discretize a diagonalized continuous-time linear SSM using bilinear transform.

    Args:
        Lambda: Diagonal state matrix (P,).
        B_tilde: Input matrix (P, H).
        Delta: Discretization step sizes (P,).

    Returns:
        Tuple of (discretized Lambda_bar, discretized B_bar).
    """
    Identity = jnp.ones(Lambda.shape[0])

    BL = 1 / (Identity - (Delta / 2.0) * Lambda)
    Lambda_bar = BL * (Identity + (Delta / 2.0) * Lambda)
    B_bar = (BL * Delta)[..., None] * B_tilde
    return Lambda_bar, B_bar


def _discretize_zoh(Lambda: Array, B_tilde: Array, Delta: Array) -> tuple[Array, Array]:
    """Discretize a diagonalized continuous-time linear SSM using zero-order hold.

    Args:
        Lambda: Diagonal state matrix (P,).
        B_tilde: Input matrix (P, H).
        Delta: Discretization step sizes (P,).

    Returns:
        Tuple of (discretized Lambda_bar, discretized B_bar).
    """
    Identity = jnp.ones(Lambda.shape[0])
    Lambda_bar = jnp.exp(Lambda * Delta)
    B_bar = (1 / Lambda * (Lambda_bar - Identity))[..., None] * B_tilde
    return Lambda_bar, B_bar


@jax.vmap
def _binary_operator(q_i, q_j):
    """Binary operator for parallel scan of linear recurrence.

    Assumes a diagonal matrix A.

    Args:
        q_i: Tuple containing (A_i, Bu_i) at position i.
        q_j: Tuple containing (A_j, Bu_j) at position j.

    Returns:
        New element (A_out, Bu_out).
    """
    A_i, b_i = q_i
    A_j, b_j = q_j
    return A_j * A_i, A_j * b_i + b_j


def _apply_ssm(
    Lambda_bar: Array,
    B_bar: Array,
    C_tilde: Array,
    input_sequence: Array,
    conj_sym: bool,
) -> Array:
    """Compute the L x H output of discretized SSM given an L x H input.

    Args:
        Lambda_bar: Discretized diagonal state matrix (P,).
        B_bar: Discretized input matrix (P, H).
        C_tilde: Output matrix (H, P).
        input_sequence: Input sequence of features (L, H).
        conj_sym: Whether conjugate symmetry is enforced.

    Returns:
        The SSM outputs (S5 layer preactivations) of shape (L, H).
    """
    Lambda_elements = Lambda_bar * jnp.ones((input_sequence.shape[0], Lambda_bar.shape[0]))
    Bu_elements = jax.vmap(lambda u: B_bar @ u)(input_sequence)

    _, xs = jax.lax.associative_scan(_binary_operator, (Lambda_elements, Bu_elements))

    if conj_sym:
        return jax.vmap(lambda x: 2 * (C_tilde @ x).real)(xs)
    else:
        return jax.vmap(lambda x: (C_tilde @ x).real)(xs)
