"""Linear system class for the Ellipsoidal Toolbox.

Models continuous-time or discrete-time linear systems of the form

    Continuous:   dx/dt = A(t) x(t) + B(t) u(t) + G(t) v(t)
                   y(t) = C(t) x(t) + w(t)

    Discrete:    x[k+1] = A[k] x[k] + B[k] u[k] + G[k] v[k]
                   y[k] = C[k] x[k] + w[k]

where control u is bounded by an ellipsoid U, disturbance v by V,
and observation noise w by W.

Time-varying matrices are represented as callables ``f(t) -> ndarray``.
Constant matrices are plain numpy arrays.  The system is time-invariant
(LTI) when all matrices are constant.

Reference
---------
MATLAB Ellipsoidal Toolbox v3, ``control/@linsys/linsys.m``
"""

from __future__ import annotations

from typing import Callable, Optional, Union

import numpy as np
from numpy.typing import ArrayLike

from ..core.ellipsoid import Ellipsoid
from ..core.config import ABS_TOL

__all__ = ["LinSys"]

MatrixLike = Union[ArrayLike, Callable[[float], np.ndarray]]
BoundLike = Union[Ellipsoid, ArrayLike, Callable[[float], Ellipsoid], None]


class LinSys:
    """A linear dynamical system with ellipsoidal uncertainty bounds.

    Parameters
    ----------
    A : array_like or callable
        State dynamics matrix (n x n).  Pass a callable ``f(t) -> ndarray``
        for time-varying systems.
    B : array_like or callable or None
        Control input matrix (n x k).
    U : Ellipsoid, array_like, callable, or None
        Control bounds.  ``Ellipsoid`` for constant bounds, ``ndarray``
        for a deterministic control signal, ``callable`` for time-varying
        bounds, or ``None`` for no bounds.
    G : array_like or callable or None
        Disturbance input matrix (n x l).
    V : Ellipsoid, array_like, callable, or None
        Disturbance bounds (same convention as *U*).
    C : array_like or callable or None
        Output matrix (m x n).  Defaults to ``I_n``.
    W : Ellipsoid, array_like, callable, or None
        Observation noise bounds (same convention as *U*).
    discrete : bool
        If *True*, the system is discrete-time.

    Examples
    --------
    >>> A = np.array([[0, 1], [-1, 0]])
    >>> B = np.array([[0], [1]])
    >>> U = Ellipsoid([0.0], np.eye(1))
    >>> sys = LinSys(A, B, U)
    """

    def __init__(
        self,
        A: Optional[MatrixLike] = None,
        B: Optional[MatrixLike] = None,
        U: BoundLike = None,
        G: Optional[MatrixLike] = None,
        V: BoundLike = None,
        C: Optional[MatrixLike] = None,
        W: BoundLike = None,
        discrete: bool = False,
    ) -> None:
        # ── Empty system ─────────────────────────────────────────────
        if A is None:
            self._A: Optional[np.ndarray | Callable] = None
            self._B: Optional[np.ndarray | Callable] = None
            self._G: Optional[np.ndarray | Callable] = None
            self._C: Optional[np.ndarray | Callable] = None
            self._U: BoundLike = None
            self._V: BoundLike = None
            self._W: BoundLike = None
            self._discrete: bool = False
            self._lti: bool = True
            self._n: int = 0
            return

        # ── State matrix A ───────────────────────────────────────────
        lti = True
        if callable(A):
            self._A = A
            lti = False
            n = A(0.0).shape[0]
        else:
            A_arr = np.asarray(A, dtype=float)
            if A_arr.ndim != 2 or A_arr.shape[0] != A_arr.shape[1]:
                raise ValueError("LinSys: A must be a square matrix.")
            self._A = A_arr
            n = A_arr.shape[0]
        self._n = n

        # ── Input matrix B ───────────────────────────────────────────
        if B is None:
            self._B = None
            n_inputs = 0
        elif callable(B):
            self._B = B
            lti = False
            B0 = B(0.0)
            if B0.shape[0] != n:
                raise ValueError(
                    f"LinSys: B row count ({B0.shape[0]}) must match "
                    f"A dimension ({n})."
                )
            n_inputs = B0.shape[1]
        else:
            B_arr = np.asarray(B, dtype=float)
            if B_arr.ndim == 1:
                B_arr = B_arr.reshape(-1, 1)
            if B_arr.shape[0] != n:
                raise ValueError(
                    f"LinSys: B row count ({B_arr.shape[0]}) must match "
                    f"A dimension ({n})."
                )
            self._B = B_arr
            n_inputs = B_arr.shape[1]

        # ── Control bounds U ─────────────────────────────────────────
        self._U = _validate_bound(U, n_inputs, "U", "B")

        # ── Disturbance input matrix G ───────────────────────────────
        if G is None:
            self._G = None
            n_disturb = 0
        elif callable(G):
            self._G = G
            lti = False
            G0 = G(0.0)
            if G0.shape[0] != n:
                raise ValueError(
                    f"LinSys: G row count ({G0.shape[0]}) must match "
                    f"A dimension ({n})."
                )
            n_disturb = G0.shape[1]
        else:
            G_arr = np.asarray(G, dtype=float)
            if G_arr.ndim == 1:
                G_arr = G_arr.reshape(-1, 1)
            if G_arr.shape[0] != n:
                raise ValueError(
                    f"LinSys: G row count ({G_arr.shape[0]}) must match "
                    f"A dimension ({n})."
                )
            self._G = G_arr
            n_disturb = G_arr.shape[1]

        # ── Disturbance bounds V ─────────────────────────────────────
        if G is None or V is None:
            self._V = None
            if G is not None and V is None:
                pass  # G exists but no bounds — unbounded disturbance
        else:
            self._V = _validate_bound(V, n_disturb, "V", "G")

        # ── Output matrix C ──────────────────────────────────────────
        if C is None:
            self._C = np.eye(n)
            n_outputs = n
        elif callable(C):
            self._C = C
            lti = False
            C0 = C(0.0)
            if C0.shape[1] != n:
                raise ValueError(
                    f"LinSys: C column count ({C0.shape[1]}) must match "
                    f"A dimension ({n})."
                )
            n_outputs = C0.shape[0]
        else:
            C_arr = np.asarray(C, dtype=float)
            if C_arr.ndim == 1:
                C_arr = C_arr.reshape(1, -1)
            if C_arr.shape[1] != n:
                raise ValueError(
                    f"LinSys: C column count ({C_arr.shape[1]}) must match "
                    f"A dimension ({n})."
                )
            self._C = C_arr
            n_outputs = C_arr.shape[0]

        # ── Noise bounds W ───────────────────────────────────────────
        self._W = _validate_bound(W, n_outputs, "W", "C")

        self._discrete = discrete
        self._lti = lti

    # ------------------------------------------------------------------
    # Matrix evaluation
    # ------------------------------------------------------------------

    def A_at(self, t: float = 0.0) -> np.ndarray:
        """Evaluate A at time *t*.  Returns the constant matrix for LTI."""
        if self._A is None:
            raise ValueError("LinSys is empty.")
        if callable(self._A):
            return np.asarray(self._A(t), dtype=float)
        return self._A

    def B_at(self, t: float = 0.0) -> Optional[np.ndarray]:
        """Evaluate B at time *t*."""
        if self._B is None:
            return None
        if callable(self._B):
            return np.asarray(self._B(t), dtype=float)
        return self._B

    def G_at(self, t: float = 0.0) -> Optional[np.ndarray]:
        """Evaluate G at time *t*."""
        if self._G is None:
            return None
        if callable(self._G):
            return np.asarray(self._G(t), dtype=float)
        return self._G

    def C_at(self, t: float = 0.0) -> Optional[np.ndarray]:
        """Evaluate C at time *t*."""
        if self._C is None:
            return None
        if callable(self._C):
            return np.asarray(self._C(t), dtype=float)
        return self._C

    def U_at(self, t: float = 0.0) -> Optional[Ellipsoid]:
        """Evaluate control bounds U at time *t*.

        Returns an ``Ellipsoid`` (possibly degenerate for a fixed vector),
        or ``None`` if unbounded.
        """
        return _eval_bound(self._U, t)

    def V_at(self, t: float = 0.0) -> Optional[Ellipsoid]:
        """Evaluate disturbance bounds V at time *t*."""
        return _eval_bound(self._V, t)

    def W_at(self, t: float = 0.0) -> Optional[Ellipsoid]:
        """Evaluate noise bounds W at time *t*."""
        return _eval_bound(self._W, t)

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def A(self) -> Optional[np.ndarray | Callable]:
        """System dynamics matrix (or callable)."""
        return self._A

    @property
    def B(self) -> Optional[np.ndarray | Callable]:
        """Control input matrix (or callable)."""
        return self._B

    @property
    def G(self) -> Optional[np.ndarray | Callable]:
        """Disturbance input matrix (or callable)."""
        return self._G

    @property
    def C(self) -> Optional[np.ndarray | Callable]:
        """Output matrix (or callable)."""
        return self._C

    @property
    def control(self) -> BoundLike:
        """Control bounds U."""
        return self._U

    @property
    def disturbance(self) -> BoundLike:
        """Disturbance bounds V."""
        return self._V

    @property
    def noise(self) -> BoundLike:
        """Noise bounds W."""
        return self._W

    @property
    def is_discrete(self) -> bool:
        """True if the system is discrete-time."""
        return self._discrete

    @property
    def is_lti(self) -> bool:
        """True if all matrices are constant (Linear Time-Invariant)."""
        return self._lti

    # ------------------------------------------------------------------
    # Basic queries
    # ------------------------------------------------------------------

    def is_empty(self) -> bool:
        """Return *True* if the system has no dynamics matrix."""
        return self._A is None

    def dimension(self) -> tuple[int, int, int, int]:
        """Return ``(state_dim, input_dim, output_dim, disturbance_dim)``.

        Returns
        -------
        state_dim : int
        input_dim : int
        output_dim : int
        disturbance_dim : int
        """
        if self.is_empty():
            return 0, 0, 0, 0

        n = self._n
        B = self.B_at(0.0)
        n_in = B.shape[1] if B is not None else 0
        C = self.C_at(0.0)
        n_out = C.shape[0] if C is not None else n
        G = self.G_at(0.0)
        n_dist = G.shape[1] if G is not None else 0

        return n, n_in, n_out, n_dist

    def has_disturbance(self) -> bool:
        """Return *True* if the system has a disturbance channel (G and V)."""
        return self._G is not None and self._V is not None

    def has_noise(self) -> bool:
        """Return *True* if the system has observation noise (W)."""
        return self._W is not None

    # ------------------------------------------------------------------
    # Display
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        if self.is_empty():
            return "LinSys(empty)"
        n, n_in, n_out, n_dist = self.dimension()
        kind = "discrete" if self._discrete else "continuous"
        lti_tag = ", LTI" if self._lti else ", LTV"
        return (
            f"LinSys({kind}{lti_tag}, state={n}, input={n_in}, "
            f"output={n_out}, disturbance={n_dist})"
        )

    def __str__(self) -> str:
        if self.is_empty():
            return "Empty linear system."
        n, n_in, n_out, n_dist = self.dimension()
        kind = "discrete-time" if self._discrete else "continuous-time"
        variant = "LTI" if self._lti else "LTV"
        lines = [
            f"Linear {kind} {variant} system:",
            f"  state dim       = {n}",
            f"  input dim       = {n_in}",
            f"  output dim      = {n_out}",
            f"  disturbance dim = {n_dist}",
            f"  control bounds  = {_bound_summary(self._U)}",
            f"  disturbance bounds = {_bound_summary(self._V)}",
            f"  noise bounds    = {_bound_summary(self._W)}",
        ]
        return "\n".join(lines)


# =====================================================================
# Module-level helpers
# =====================================================================

def _validate_bound(
    bound: BoundLike,
    expected_dim: int,
    name: str,
    matrix_name: str,
) -> BoundLike:
    """Validate a control/disturbance/noise bound."""
    if bound is None:
        return None

    if isinstance(bound, Ellipsoid):
        if bound.is_empty():
            return None
        bd, _ = bound.dimension()
        if bd != expected_dim:
            raise ValueError(
                f"LinSys: {name} dimension ({bd}) must match "
                f"{matrix_name} columns ({expected_dim})."
            )
        return bound

    if callable(bound):
        b0 = bound(0.0)
        if isinstance(b0, Ellipsoid):
            if not b0.is_empty():
                bd, _ = b0.dimension()
                if bd != expected_dim:
                    raise ValueError(
                        f"LinSys: {name}(0) dimension ({bd}) must match "
                        f"{matrix_name} columns ({expected_dim})."
                    )
        return bound

    arr = np.asarray(bound, dtype=float).ravel()
    if arr.shape[0] != expected_dim:
        raise ValueError(
            f"LinSys: {name} vector length ({arr.shape[0]}) must match "
            f"{matrix_name} columns ({expected_dim})."
        )
    return arr


def _eval_bound(bound: BoundLike, t: float) -> Optional[Ellipsoid]:
    """Evaluate a bound at time *t*, returning an Ellipsoid or None."""
    if bound is None:
        return None
    if isinstance(bound, Ellipsoid):
        return bound
    if callable(bound):
        result = bound(t)
        if isinstance(result, Ellipsoid):
            return result
        return None
    if isinstance(bound, np.ndarray):
        n = bound.shape[0]
        return Ellipsoid(bound, np.zeros((n, n)))
    return None


def _bound_summary(bound: BoundLike) -> str:
    """One-line description of a bound for __str__."""
    if bound is None:
        return "none"
    if isinstance(bound, Ellipsoid):
        d, r = bound.dimension()
        return f"Ellipsoid(dim={d}, rank={r})"
    if callable(bound):
        return "time-varying (callable)"
    if isinstance(bound, np.ndarray):
        return f"fixed vector ({bound.ravel()})"
    return str(type(bound))
