"""ODE infrastructure for reachability analysis.

Provides:
    - ODE solver wrapper around ``scipy.integrate.solve_ivp``
    - State transition matrix (STM) computation (LTI via expm, LTV via ODE)
    - Center trajectory integration
    - Direction vector evolution
    - Precomputation helpers (BPB, BPBsr, GQG, GQGsr)

Reference
---------
MATLAB Ellipsoidal Toolbox v3,
``control/auxiliary/ell_ode_solver.m``,
``control/auxiliary/ell_stm_ode.m``,
``control/auxiliary/ell_center_ode.m``.
"""

from __future__ import annotations

from typing import Callable, Optional

import numpy as np
import scipy.linalg
import scipy.integrate

from ..core.config import ABS_TOL, REL_TOL
from ..core.helpers import regularize
from .linsys import LinSys

__all__ = [
    "ode_solve",
    "compute_stm",
    "compute_center",
    "evolve_directions",
    "precompute_bounds",
    "PrecomputedData",
]


# =====================================================================
# ODE solver wrapper
# =====================================================================

def ode_solve(
    fun: Callable[[float, np.ndarray], np.ndarray],
    t_span: tuple[float, float],
    y0: np.ndarray,
    t_eval: Optional[np.ndarray] = None,
    method: str = "RK45",
    rtol: float = REL_TOL,
    atol: float = ABS_TOL,
    **kwargs,
) -> scipy.integrate.OdeResult:
    """Thin wrapper around ``scipy.integrate.solve_ivp``.

    Parameters
    ----------
    fun : callable(t, y) -> dy
        Right-hand side of the ODE.
    t_span : (t0, tf)
        Integration interval.
    y0 : ndarray
        Initial condition.
    t_eval : ndarray or None
        Time points at which to store the solution.
    method : str
        Integration method (default ``'RK45'``, MATLAB uses ode45).
    rtol, atol : float
        Relative and absolute tolerances.

    Returns
    -------
    OdeResult
        Solution object with ``.t`` and ``.y`` attributes.
    """
    return scipy.integrate.solve_ivp(
        fun,
        t_span,
        y0,
        method=method,
        t_eval=t_eval,
        rtol=rtol,
        atol=atol,
        **kwargs,
    )


# =====================================================================
# State transition matrix
# =====================================================================

def compute_stm(
    sys: LinSys,
    t_span: tuple[float, float],
    t_eval: np.ndarray,
    backward: bool = False,
) -> np.ndarray:
    """Compute the state transition matrix Φ(t, t₀) at each t in *t_eval*.

    For LTI systems, uses ``scipy.linalg.expm``.
    For LTV systems, solves dΦ/dt = A(t) Φ numerically.

    Parameters
    ----------
    sys : LinSys
    t_span : (t0, tf)
    t_eval : ndarray of shape (N,)
        Time points.
    backward : bool
        If True, compute backward STM (time-reversed).

    Returns
    -------
    Phi : ndarray of shape (N, n, n)
        ``Phi[k]`` is Φ(t_eval[k], t0).
    """
    if sys.is_empty():
        raise ValueError("compute_stm: system must be non-empty.")

    n = sys._n
    N = len(t_eval)

    # ── LTI shortcut: Φ(t) = expm(A * (t - t0)) ────────────────────
    if sys.is_lti and not callable(sys.A):
        A = sys.A
        t0 = t_span[0]
        Phi = np.zeros((N, n, n))
        for k, tk in enumerate(t_eval):
            dt = tk - t0
            if backward:
                dt = -dt
            Phi[k] = scipy.linalg.expm(A * dt)
        return Phi

    # ── LTV: solve dΦ/dt = A(t) Φ via ODE ──────────────────────────
    Phi0 = np.eye(n).ravel()  # vec(I_n)

    def _stm_rhs(t: float, x: np.ndarray) -> np.ndarray:
        if backward:
            A_t = sys.A_at(-t)
            dPhidt = -A_t @ x.reshape(n, n)
        else:
            A_t = sys.A_at(t)
            dPhidt = A_t @ x.reshape(n, n)
        return dPhidt.ravel()

    if backward:
        t_span_ode = (-t_span[0], -t_span[1])
        t_eval_ode = -t_eval
    else:
        t_span_ode = t_span
        t_eval_ode = t_eval

    sol = ode_solve(_stm_rhs, t_span_ode, Phi0, t_eval=t_eval_ode)

    if sol.status != 0:
        raise RuntimeError(f"compute_stm: ODE solver failed: {sol.message}")

    Phi = np.zeros((N, n, n))
    for k in range(sol.y.shape[1]):
        Phi[k] = sol.y[:, k].reshape(n, n)

    return Phi


def compute_stm_inv(
    Phi: np.ndarray,
) -> np.ndarray:
    """Compute Φ^{-1}(t) for each time step.

    Parameters
    ----------
    Phi : ndarray (N, n, n)

    Returns
    -------
    Phi_inv : ndarray (N, n, n)
    """
    N = Phi.shape[0]
    Phi_inv = np.zeros_like(Phi)
    for k in range(N):
        Phi_inv[k] = np.linalg.inv(Phi[k])
    return Phi_inv


# =====================================================================
# Center trajectory
# =====================================================================

def compute_center(
    sys: LinSys,
    x0: np.ndarray,
    t_span: tuple[float, float],
    t_eval: np.ndarray,
    backward: bool = False,
) -> np.ndarray:
    """Integrate the center trajectory.

    Solves::

        dx/dt = A(t) x + B(t) p(t) + G(t) q(t)

    where p(t) = center of U(t), q(t) = center of V(t).

    Parameters
    ----------
    sys : LinSys
    x0 : ndarray (n,)
        Initial center.
    t_span : (t0, tf)
    t_eval : ndarray (N,)
    backward : bool

    Returns
    -------
    centers : ndarray (N, n)
        Center at each time point.
    """
    if sys.is_empty():
        raise ValueError("compute_center: system must be non-empty.")

    n = sys._n
    x0 = np.asarray(x0, dtype=float).ravel()

    def _center_rhs(t: float, x: np.ndarray) -> np.ndarray:
        t_actual = -t if backward else t
        s = -1.0 if backward else 1.0

        A_t = sys.A_at(t_actual)
        dx = s * A_t @ x

        B_t = sys.B_at(t_actual)
        if B_t is not None:
            U_t = sys.U_at(t_actual)
            if U_t is not None and not U_t.is_empty():
                p = U_t.center
            else:
                p = np.zeros(B_t.shape[1])
            dx += s * B_t @ p

        G_t = sys.G_at(t_actual)
        if G_t is not None and sys.has_disturbance():
            V_t = sys.V_at(t_actual)
            if V_t is not None and not V_t.is_empty():
                q = V_t.center
            else:
                q = np.zeros(G_t.shape[1])
            dx += s * G_t @ q

        return dx

    if backward:
        t_span_ode = (-t_span[0], -t_span[1])
        t_eval_ode = -t_eval
    else:
        t_span_ode = t_span
        t_eval_ode = t_eval

    sol = ode_solve(_center_rhs, t_span_ode, x0, t_eval=t_eval_ode)

    if sol.status != 0:
        raise RuntimeError(f"compute_center: ODE solver failed: {sol.message}")

    centers = sol.y.T  # (N, n)
    return centers


# =====================================================================
# Direction evolution
# =====================================================================

def evolve_directions(
    Phi: np.ndarray,
    L0: np.ndarray,
    backward: bool = False,
) -> np.ndarray:
    """Evolve direction vectors through the adjoint equation.

    For forward reachability::

        l(t) = Φ(t, t₀)^{-T}  l₀

    For backward::

        l(t) = Φ(t, t₀)^T  l₀

    Parameters
    ----------
    Phi : ndarray (N, n, n)
        State transition matrices.
    L0 : ndarray (n, m)
        Initial direction columns.
    backward : bool

    Returns
    -------
    L : ndarray (N, n, m)
        Direction vectors at each time step.
    """
    N, n, _ = Phi.shape
    m = L0.shape[1]
    L = np.zeros((N, n, m))

    for k in range(N):
        if backward:
            L[k] = Phi[k].T @ L0
        else:
            Phi_inv_T = np.linalg.inv(Phi[k]).T
            L[k] = Phi_inv_T @ L0

    return L


# =====================================================================
# Precomputation helpers
# =====================================================================

class PrecomputedData:
    """Container for precomputed quantities at each time step.

    Attributes
    ----------
    t_eval : ndarray (N,)
    Phi : ndarray (N, n, n)
        State transition matrices.
    Phi_inv : ndarray (N, n, n)
    centers : ndarray (N, n)
        Center trajectory.
    L : ndarray (N, n, m)
        Direction vectors.
    BPB : ndarray (N, n, n)
        B(t) P(t) B(t)^T at each time step.
    BPBsr : ndarray (N, n, n)
        sqrt(BPB) at each time step.
    GQG : ndarray (N, n, n) or None
        G(t) Q_v(t) G(t)^T (None if no disturbance).
    GQGsr : ndarray (N, n, n) or None
        sqrt(GQG) (None if no disturbance).
    """

    def __init__(self) -> None:
        self.t_eval: np.ndarray = np.array([])
        self.Phi: np.ndarray = np.array([])
        self.Phi_inv: np.ndarray = np.array([])
        self.centers: np.ndarray = np.array([])
        self.L: np.ndarray = np.array([])
        self.BPB: np.ndarray = np.array([])
        self.BPBsr: np.ndarray = np.array([])
        self.GQG: Optional[np.ndarray] = None
        self.GQGsr: Optional[np.ndarray] = None


def precompute_bounds(
    sys: LinSys,
    t_eval: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
    """Precompute BPB, BPBsr and optionally GQG, GQGsr at each time step.

    Parameters
    ----------
    sys : LinSys
    t_eval : ndarray (N,)

    Returns
    -------
    BPB : ndarray (N, n, n)
    BPBsr : ndarray (N, n, n)
    GQG : ndarray (N, n, n) or None
    GQGsr : ndarray (N, n, n) or None
    """
    n = sys._n
    N = len(t_eval)

    BPB = np.zeros((N, n, n))
    BPBsr = np.zeros((N, n, n))

    has_dist = sys.has_disturbance()
    GQG = np.zeros((N, n, n)) if has_dist else None
    GQGsr = np.zeros((N, n, n)) if has_dist else None

    for k, t in enumerate(t_eval):
        B_t = sys.B_at(t)
        if B_t is not None:
            U_t = sys.U_at(t)
            if U_t is not None and not U_t.is_empty():
                P = U_t.shape
                if P is not None:
                    bpb = B_t @ P @ B_t.T
                    bpb = 0.5 * (bpb + bpb.T)
                    BPB[k] = bpb
                    BPBsr[k] = _safe_sqrtm(bpb)

        if has_dist:
            G_t = sys.G_at(t)
            if G_t is not None:
                V_t = sys.V_at(t)
                if V_t is not None and not V_t.is_empty():
                    Q_v = V_t.shape
                    if Q_v is not None:
                        gqg = G_t @ Q_v @ G_t.T
                        gqg = 0.5 * (gqg + gqg.T)
                        GQG[k] = gqg
                        GQGsr[k] = _safe_sqrtm(gqg)

    return BPB, BPBsr, GQG, GQGsr


def precompute_all(
    sys: LinSys,
    x0: np.ndarray,
    L0: np.ndarray,
    t_span: tuple[float, float],
    N_steps: int = 100,
    backward: bool = False,
) -> PrecomputedData:
    """Run all precomputations for reach set analysis.

    Parameters
    ----------
    sys : LinSys
    x0 : ndarray (n,)
        Initial center of the ellipsoidal set.
    L0 : ndarray (n, m)
        Initial direction matrix.
    t_span : (t0, tf)
    N_steps : int
        Number of time grid points.
    backward : bool

    Returns
    -------
    PrecomputedData
    """
    t0, tf = t_span
    t_eval = np.linspace(t0, tf, N_steps)

    data = PrecomputedData()
    data.t_eval = t_eval

    data.Phi = compute_stm(sys, t_span, t_eval, backward=backward)
    data.Phi_inv = compute_stm_inv(data.Phi)

    data.centers = compute_center(sys, x0, t_span, t_eval, backward=backward)

    if L0.ndim == 1:
        L0 = L0.reshape(-1, 1)
    data.L = evolve_directions(data.Phi, L0, backward=backward)

    data.BPB, data.BPBsr, data.GQG, data.GQGsr = precompute_bounds(sys, t_eval)

    return data


# =====================================================================
# Internal helpers
# =====================================================================

def _safe_sqrtm(M: np.ndarray) -> np.ndarray:
    """Compute matrix square root, regularizing if needed."""
    n = M.shape[0]
    if np.linalg.matrix_rank(M, tol=ABS_TOL) < n:
        M = regularize(M)
    S = np.real(scipy.linalg.sqrtm(M))
    return 0.5 * (S + S.T)
