"""State-space model representation (West & Harrison notation)."""

from __future__ import annotations

from dataclasses import dataclass

import jax
import jax.numpy as jnp
from jax import Array

from dynaris.core.types import GaussianState


@dataclass(frozen=True)
class StateSpaceModel:
    """Linear-Gaussian Dynamic Linear Model.

    Following West & Harrison (1997) notation:

    System equation:    theta_t = G @ theta_{t-1} + omega_t,  omega_t ~ N(0, W)
    Observation eq:     Y_t = F' @ theta_t + nu_t,            nu_t ~ N(0, V)

    Attributes:
        observation_matrix: F, shape (obs_dim, state_dim).
        system_matrix: G, shape (state_dim, state_dim).
        obs_cov: V, shape (obs_dim, obs_dim).
        evolution_cov: W, shape (state_dim, state_dim).
        input_matrix: B, shape (state_dim, input_dim) or None.
    """

    observation_matrix: Array  # F: (m, n)
    system_matrix: Array  # G: (n, n)
    obs_cov: Array  # V: (m, m)
    evolution_cov: Array  # W: (n, n)
    input_matrix: Array | None = None  # B: (n, p) or None

    # --- Dimension properties ---

    @property
    def state_dim(self) -> int:
        """Dimension of the state (parameter) vector."""
        return int(self.system_matrix.shape[-1])

    @property
    def obs_dim(self) -> int:
        """Dimension of the observation vector."""
        return int(self.observation_matrix.shape[-2])

    # --- Short aliases (West & Harrison notation) ---

    @property
    def F(self) -> Array:  # noqa: N802
        """Observation/regression matrix (F in W&H)."""
        return self.observation_matrix

    @property
    def G(self) -> Array:  # noqa: N802
        """System/evolution matrix (G in W&H)."""
        return self.system_matrix

    @property
    def V(self) -> Array:  # noqa: N802
        """Observational variance/covariance (V in W&H)."""
        return self.obs_cov

    @property
    def W(self) -> Array:  # noqa: N802
        """Evolution covariance (W in W&H)."""
        return self.evolution_cov

    @property
    def B(self) -> Array | None:  # noqa: N802
        """Alias for input_matrix."""
        return self.input_matrix

    # --- Factory methods ---

    def initial_state(
        self,
        mean: Array | None = None,
        cov: Array | None = None,
    ) -> GaussianState:
        """Create a default initial GaussianState for this model.

        Args:
            mean: Initial state mean (m_0). Defaults to zeros.
            cov: Initial state covariance (C_0). Defaults to 1e6 * I (diffuse prior).

        Returns:
            GaussianState with the specified or default initial conditions.
        """
        n = self.state_dim
        if mean is None:
            mean = jnp.zeros(n)
        if cov is None:
            cov = jnp.eye(n) * 1e6
        return GaussianState(mean=mean, cov=cov)

    # --- Composition (superposition principle) ---

    def __add__(self, other: StateSpaceModel) -> StateSpaceModel:
        """Compose two models via superposition (West & Harrison).

        The resulting model has:
        - G_new = block_diag(G1, G2)
        - F_new = [F1, F2]  (horizontal concatenation)
        - W_new = block_diag(W1, W2)
        - V_new = V1 + V2  (shared observation noise adds)
        """
        n1, n2 = self.state_dim, other.state_dim

        system = jnp.block(
            [
                [self.G, jnp.zeros((n1, n2))],
                [jnp.zeros((n2, n1)), other.G],
            ]
        )
        observation = jnp.concatenate([self.F, other.F], axis=-1)
        evolution = jnp.block(
            [
                [self.W, jnp.zeros((n1, n2))],
                [jnp.zeros((n2, n1)), other.W],
            ]
        )
        obs = self.V + other.V

        input_mat: Array | None = None
        if self.input_matrix is not None and other.input_matrix is not None:
            p1 = self.input_matrix.shape[-1]
            p2 = other.input_matrix.shape[-1]
            input_mat = jnp.block(
                [
                    [self.input_matrix, jnp.zeros((n1, p2))],
                    [jnp.zeros((n2, p1)), other.input_matrix],
                ]
            )
        elif self.input_matrix is not None:
            input_mat = jnp.concatenate(
                [
                    self.input_matrix,
                    jnp.zeros((n2, self.input_matrix.shape[-1])),
                ],
                axis=-2,
            )
        elif other.input_matrix is not None:
            input_mat = jnp.concatenate(
                [
                    jnp.zeros((n1, other.input_matrix.shape[-1])),
                    other.input_matrix,
                ],
                axis=-2,
            )

        return StateSpaceModel(
            observation_matrix=observation,
            system_matrix=system,
            obs_cov=obs,
            evolution_cov=evolution,
            input_matrix=input_mat,
        )

    def __repr__(self) -> str:
        b_info = (
            f", input_dim={self.input_matrix.shape[-1]}" if self.input_matrix is not None else ""
        )
        return f"StateSpaceModel(state_dim={self.state_dim}, obs_dim={self.obs_dim}{b_info})"

    # --- JAX pytree registration ---

    def tree_flatten(self) -> tuple[list[Array], dict[str, bool]]:
        """Flatten into JAX pytree leaves and auxiliary data."""
        has_input = self.input_matrix is not None
        leaves: list[Array] = [
            self.observation_matrix,
            self.system_matrix,
            self.obs_cov,
            self.evolution_cov,
        ]
        if has_input:
            leaves.append(self.input_matrix)  # type: ignore[arg-type]
        return leaves, {"has_input": has_input}

    @classmethod
    def tree_unflatten(cls, aux_data: dict[str, bool], children: list[Array]) -> StateSpaceModel:
        """Reconstruct from JAX pytree leaves."""
        if aux_data["has_input"]:
            return cls(
                observation_matrix=children[0],
                system_matrix=children[1],
                obs_cov=children[2],
                evolution_cov=children[3],
                input_matrix=children[4],
            )
        return cls(
            observation_matrix=children[0],
            system_matrix=children[1],
            obs_cov=children[2],
            evolution_cov=children[3],
        )


jax.tree_util.register_pytree_node_class(StateSpaceModel)
