"""Reach set computation for linear systems with ellipsoidal bounds.

Phase 3C: External approximation (EA)
Phase 3D: Internal approximation (IA)
Phase 3E: Discrete-time
Phase 3F: Reach tube operations

Reference
---------
MATLAB Ellipsoidal Toolbox v3,
``control/@reach/reach.m``,
``control/auxiliary/ell_eesm_ode.m``,
``control/auxiliary/ell_eedist_ode.m``,
``control/auxiliary/ell_iesm_ode.m``,
``control/auxiliary/ell_iedist_ode.m``.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional, Union

import numpy as np
import scipy.linalg

from ..core.config import ABS_TOL
from ..core.ellipsoid import Ellipsoid
from ..core.hyperplane import Hyperplane
from ..core.helpers import ell_valign, regularize
from .linsys import LinSys
from .ode_utils import (
    ode_solve,
    compute_stm,
    compute_stm_inv,
    compute_center,
    precompute_bounds,
    _safe_sqrtm,
)

__all__ = [
    "ReachSet",
    "reach_ea",
    "reach_ia",
    "reach_ea_d",
    "reach_ia_d",
]


# =====================================================================
# Data structure
# =====================================================================

@dataclass
class ReachSet:
    """Container for reach set approximation results.

    Attributes
    ----------
    t_eval : ndarray (N,)
        Time grid.
    centers : ndarray (N, n)
        Center trajectory.
    ea_shapes : ndarray (N, m, n, n) or None
        External approximation shape matrices.
        ``ea_shapes[k, j]`` is the shape matrix at time ``t_eval[k]``
        for direction ``j``.
    ia_shapes : ndarray (N, m, n, n) or None
        Internal approximation shape matrices.
    n_dirs : int
        Number of directions.
    n : int
        State dimension.
    sys : LinSys or None
        Reference to the originating system.
    X0 : Ellipsoid or None
        Initial set.
    L0 : ndarray or None
        Initial direction matrix.
    _is_cut : bool
        True if this is a cut of another reach set.
    _is_projection : bool
        True if this is a projection of another reach set.
    """
    t_eval: np.ndarray = field(default_factory=lambda: np.array([]))
    centers: np.ndarray = field(default_factory=lambda: np.array([]))
    ea_shapes: Optional[np.ndarray] = None
    ia_shapes: Optional[np.ndarray] = None
    n_dirs: int = 0
    n: int = 0
    sys: Optional[LinSys] = field(default=None, repr=False)
    X0: Optional[Ellipsoid] = field(default=None, repr=False)
    L0: Optional[np.ndarray] = field(default=None, repr=False)
    _is_cut: bool = False
    _is_projection: bool = False

    # ── Ellipsoid accessors ──────────────────────────────────────────

    def ea_ellipsoid(self, k: int, j: int) -> Ellipsoid:
        """Return the EA ellipsoid at time step *k*, direction *j*."""
        if self.ea_shapes is None:
            raise ValueError("No external approximation computed.")
        return Ellipsoid(self.centers[k], self.ea_shapes[k, j])

    def ea_ellipsoids_at(self, k: int) -> list[Ellipsoid]:
        """Return all EA ellipsoids at time step *k*."""
        return [self.ea_ellipsoid(k, j) for j in range(self.n_dirs)]

    def ia_ellipsoid(self, k: int, j: int) -> Ellipsoid:
        """Return the IA ellipsoid at time step *k*, direction *j*."""
        if self.ia_shapes is None:
            raise ValueError("No internal approximation computed.")
        return Ellipsoid(self.centers[k], self.ia_shapes[k, j])

    def ia_ellipsoids_at(self, k: int) -> list[Ellipsoid]:
        """Return all IA ellipsoids at time step *k*."""
        return [self.ia_ellipsoid(k, j) for j in range(self.n_dirs)]

    # ── Accessors (MATLAB-compatible names) ──────────────────────────

    def get_ea(self) -> list[list[Ellipsoid]]:
        """Return EA ellipsoids for all time steps.

        Returns
        -------
        list of list of Ellipsoid
            ``result[k][j]`` is the EA ellipsoid at time step *k*, direction *j*.
        """
        if self.ea_shapes is None:
            raise ValueError("No external approximation computed.")
        N = len(self.t_eval)
        return [self.ea_ellipsoids_at(k) for k in range(N)]

    def get_ia(self) -> list[list[Ellipsoid]]:
        """Return IA ellipsoids for all time steps."""
        if self.ia_shapes is None:
            raise ValueError("No internal approximation computed.")
        N = len(self.t_eval)
        return [self.ia_ellipsoids_at(k) for k in range(N)]

    def get_center(self) -> np.ndarray:
        """Return the center trajectory.

        Returns
        -------
        ndarray (N, n)
        """
        return self.centers.copy()

    def get_directions(self) -> Optional[np.ndarray]:
        """Return the initial direction matrix L0.

        Returns
        -------
        ndarray (n, m) or None
        """
        if self.L0 is not None:
            return self.L0.copy()
        return None

    def get_system(self) -> Optional[LinSys]:
        """Return the originating linear system."""
        return self.sys

    def dimension(self) -> int:
        """Return the state dimension of the reach set."""
        return self.n

    def iscut(self) -> bool:
        """Return True if this is a cut of another reach set."""
        return self._is_cut

    def isprojection(self) -> bool:
        """Return True if this is a projection of another reach set."""
        return self._is_projection

    # ── cut ───────────────────────────────────────────────────────────

    def cut(
        self,
        t: Union[float, tuple[float, float]],
    ) -> "ReachSet":
        """Extract the reach set at a specific time or time interval.

        Parameters
        ----------
        t : float or (t_a, t_b)
            If a float, return the single cross-section closest to *t*.
            If a tuple, return the sub-tube over the interval [t_a, t_b].

        Returns
        -------
        ReachSet
            A new ReachSet marked as a cut.
        """
        if isinstance(t, (int, float)):
            idx = int(np.argmin(np.abs(self.t_eval - float(t))))
            slc = slice(idx, idx + 1)
        else:
            t_a, t_b = float(t[0]), float(t[1])
            mask = (self.t_eval >= t_a - ABS_TOL) & (self.t_eval <= t_b + ABS_TOL)
            indices = np.where(mask)[0]
            if len(indices) == 0:
                raise ValueError(
                    f"cut: no time points in [{t_a}, {t_b}]. "
                    f"Range is [{self.t_eval[0]}, {self.t_eval[-1]}]."
                )
            slc = slice(indices[0], indices[-1] + 1)

        return ReachSet(
            t_eval=self.t_eval[slc].copy(),
            centers=self.centers[slc].copy(),
            ea_shapes=self.ea_shapes[slc].copy() if self.ea_shapes is not None else None,
            ia_shapes=self.ia_shapes[slc].copy() if self.ia_shapes is not None else None,
            n_dirs=self.n_dirs,
            n=self.n,
            sys=self.sys,
            X0=self.X0,
            L0=self.L0,
            _is_cut=True,
        )

    # ── projection ────────────────────────────────────────────────────

    def projection(self, basis: np.ndarray) -> "ReachSet":
        """Project the reach tube onto a lower-dimensional subspace.

        For E(q, Q), the projection onto the column space of *basis* B is
        E(B^T q, B^T Q B).

        Parameters
        ----------
        basis : ndarray (n, p)
            Orthonormal basis columns defining the subspace. Must have
            ``basis.shape[0] == self.n``.

        Returns
        -------
        ReachSet
            A new ReachSet in the projected space, marked as a projection.
        """
        B = np.asarray(basis, dtype=float)
        if B.ndim == 1:
            B = B.reshape(-1, 1)
        if B.shape[0] != self.n:
            raise ValueError(
                f"projection: basis rows ({B.shape[0]}) must match "
                f"state dimension ({self.n})."
            )
        p = B.shape[1]
        N = len(self.t_eval)

        new_centers = self.centers @ B  # (N, p)

        new_ea = None
        if self.ea_shapes is not None:
            new_ea = np.zeros((N, self.n_dirs, p, p))
            for k in range(N):
                for j in range(self.n_dirs):
                    new_ea[k, j] = B.T @ self.ea_shapes[k, j] @ B

        new_ia = None
        if self.ia_shapes is not None:
            new_ia = np.zeros((N, self.n_dirs, p, p))
            for k in range(N):
                for j in range(self.n_dirs):
                    new_ia[k, j] = B.T @ self.ia_shapes[k, j] @ B

        return ReachSet(
            t_eval=self.t_eval.copy(),
            centers=new_centers,
            ea_shapes=new_ea,
            ia_shapes=new_ia,
            n_dirs=self.n_dirs,
            n=p,
            sys=self.sys,
            X0=self.X0,
            L0=self.L0,
            _is_projection=True,
        )

    # ── intersect ─────────────────────────────────────────────────────

    def intersect(
        self,
        X: Union[Ellipsoid, Hyperplane],
    ) -> tuple[bool, Optional[int]]:
        """Check whether the reach tube intersects *X* at any time step.

        Uses distance-based checks at each cross-section.

        Parameters
        ----------
        X : Ellipsoid or Hyperplane

        Returns
        -------
        hit : bool
            True if any cross-section intersects *X*.
        k_first : int or None
            Index of the first time step where intersection occurs
            (None if no intersection).
        """
        from ..calculus.distance import distance_to_ellipsoid, distance_to_hyperplane

        N = len(self.t_eval)

        for k in range(N):
            ells: list[Ellipsoid] = []
            if self.ea_shapes is not None:
                ells.extend(self.ea_ellipsoids_at(k))
            elif self.ia_shapes is not None:
                ells.extend(self.ia_ellipsoids_at(k))
            else:
                raise ValueError("intersect: no approximation data available.")

            for E in ells:
                if isinstance(X, Ellipsoid):
                    d = distance_to_ellipsoid(E, X)
                elif isinstance(X, Hyperplane):
                    d = distance_to_hyperplane(E, X)
                else:
                    raise TypeError(f"intersect: unsupported type {type(X)}.")

                if d <= ABS_TOL:
                    return True, k

        return False, None

    # ── evolve ────────────────────────────────────────────────────────

    def evolve(self, t_new: float, N_extra: int = 50) -> "ReachSet":
        """Extend the reach set to a new final time.

        Continues the computation from the last time step to *t_new*.
        Requires that the reach set was created with system metadata
        (via ``reach_ea``, ``reach_ia``, etc.).

        Parameters
        ----------
        t_new : float
            New final time (must be > current final time).
        N_extra : int
            Number of extra grid points for the extension.

        Returns
        -------
        ReachSet
            A new ReachSet covering the full time span [t0, t_new].
        """
        if self.sys is None or self.X0 is None or self.L0 is None:
            raise ValueError(
                "evolve: system metadata not available. "
                "Use reach_ea/reach_ia to create the reach set."
            )

        t_last = float(self.t_eval[-1])
        if t_new <= t_last + ABS_TOL:
            return self

        is_discrete = getattr(self.sys, '_discrete', False)
        has_ea = self.ea_shapes is not None
        has_ia = self.ia_shapes is not None

        if is_discrete:
            N_ext = int(round(t_new - t_last))
            if has_ea:
                last_ea = ReachSet(
                    t_eval=np.array([t_last]),
                    centers=self.centers[-1:],
                    ea_shapes=self.ea_shapes[-1:],
                    n_dirs=self.n_dirs, n=self.n,
                )
                X_last_ea = Ellipsoid(self.centers[-1], self.ea_shapes[-1, 0])
                ext_ea = reach_ea_d(self.sys, X_last_ea, self.L0, N_steps=N_ext)
            if has_ia:
                X_last_ia = Ellipsoid(self.centers[-1], self.ia_shapes[-1, 0])
                ext_ia = reach_ia_d(self.sys, X_last_ia, self.L0, N_steps=N_ext)

            new_t = np.concatenate([self.t_eval, self.t_eval[-1] + ext_ea.t_eval[1:] if has_ea else ext_ia.t_eval[1:]])
            new_centers = np.concatenate([self.centers, (ext_ea if has_ea else ext_ia).centers[1:]])
            new_ea = np.concatenate([self.ea_shapes, ext_ea.ea_shapes[1:]]) if has_ea else None
            new_ia = np.concatenate([self.ia_shapes, ext_ia.ia_shapes[1:]]) if has_ia else None
        else:
            t0_ext = t_last
            N_total = len(self.t_eval) + N_extra

            if has_ea:
                X_last = Ellipsoid(self.centers[-1], self.ea_shapes[-1, 0])
                ext = reach_ea(self.sys, X_last, self.L0,
                               (t0_ext, t_new), N_steps=N_extra)
            if has_ia:
                X_last_ia = Ellipsoid(self.centers[-1], self.ia_shapes[-1, 0])
                ext_ia_rs = reach_ia(self.sys, X_last_ia, self.L0,
                                     (t0_ext, t_new), N_steps=N_extra)

            new_t = np.concatenate([self.t_eval, ext.t_eval[1:] if has_ea else ext_ia_rs.t_eval[1:]])
            new_centers = np.concatenate([self.centers, (ext if has_ea else ext_ia_rs).centers[1:]])
            new_ea = np.concatenate([self.ea_shapes, ext.ea_shapes[1:]]) if has_ea else None
            new_ia = np.concatenate([self.ia_shapes, ext_ia_rs.ia_shapes[1:]]) if has_ia else None

        return ReachSet(
            t_eval=new_t,
            centers=new_centers,
            ea_shapes=new_ea,
            ia_shapes=new_ia,
            n_dirs=self.n_dirs,
            n=self.n,
            sys=self.sys,
            X0=self.X0,
            L0=self.L0,
        )

    # ── refine ────────────────────────────────────────────────────────

    def refine(self, L_new: np.ndarray) -> "ReachSet":
        """Add new approximation directions to the reach set.

        Recomputes EA/IA for the new directions and merges with existing
        data.

        Parameters
        ----------
        L_new : ndarray (n, p)
            New direction columns.

        Returns
        -------
        ReachSet
            A new ReachSet with the original + new directions.
        """
        if self.sys is None or self.X0 is None:
            raise ValueError(
                "refine: system metadata not available. "
                "Use reach_ea/reach_ia to create the reach set."
            )

        L_new = np.asarray(L_new, dtype=float)
        if L_new.ndim == 1:
            L_new = L_new.reshape(-1, 1)

        t0, tf = float(self.t_eval[0]), float(self.t_eval[-1])
        N = len(self.t_eval)
        is_discrete = getattr(self.sys, '_discrete', False)

        has_ea = self.ea_shapes is not None
        has_ia = self.ia_shapes is not None

        if is_discrete:
            N_steps = N - 1
            if has_ea:
                ext_ea = reach_ea_d(self.sys, self.X0, L_new, N_steps=N_steps)
            if has_ia:
                ext_ia = reach_ia_d(self.sys, self.X0, L_new, N_steps=N_steps)
        else:
            if has_ea:
                ext_ea = reach_ea(self.sys, self.X0, L_new, (t0, tf), N_steps=N)
            if has_ia:
                ext_ia = reach_ia(self.sys, self.X0, L_new, (t0, tf), N_steps=N)

        L0_merged = self.L0
        if L0_merged is not None:
            L0_merged = np.hstack([L0_merged, L_new])
        else:
            L0_merged = L_new

        new_ea = None
        if has_ea:
            new_ea = np.concatenate([self.ea_shapes, ext_ea.ea_shapes], axis=1)

        new_ia = None
        if has_ia:
            new_ia = np.concatenate([self.ia_shapes, ext_ia.ia_shapes], axis=1)

        return ReachSet(
            t_eval=self.t_eval.copy(),
            centers=self.centers.copy(),
            ea_shapes=new_ea,
            ia_shapes=new_ia,
            n_dirs=self.n_dirs + L_new.shape[1],
            n=self.n,
            sys=self.sys,
            X0=self.X0,
            L0=L0_merged,
        )

    # ── display ───────────────────────────────────────────────────────

    def __repr__(self) -> str:
        N = len(self.t_eval) if self.t_eval.size > 0 else 0
        parts = [f"ReachSet(n={self.n}, dirs={self.n_dirs}, steps={N}"]
        if N > 0:
            parts.append(f", t=[{self.t_eval[0]:.4g}, {self.t_eval[-1]:.4g}]")
        if self.ea_shapes is not None:
            parts.append(", EA")
        if self.ia_shapes is not None:
            parts.append(", IA")
        if self._is_cut:
            parts.append(", cut")
        if self._is_projection:
            parts.append(", proj")
        parts.append(")")
        return "".join(parts)

    def __str__(self) -> str:
        N = len(self.t_eval) if self.t_eval.size > 0 else 0
        lines = [f"Reach Set (dimension {self.n})"]
        lines.append(f"  Directions: {self.n_dirs}")
        lines.append(f"  Time steps: {N}")
        if N > 0:
            lines.append(f"  Time range: [{self.t_eval[0]:.6g}, {self.t_eval[-1]:.6g}]")
        lines.append(f"  EA computed: {self.ea_shapes is not None}")
        lines.append(f"  IA computed: {self.ia_shapes is not None}")
        if self._is_cut:
            lines.append("  (cut of parent reach set)")
        if self._is_projection:
            lines.append("  (projected)")
        if self.sys is not None:
            lines.append(f"  System: {self.sys}")
        return "\n".join(lines)


# =====================================================================
# External approximation
# =====================================================================

def reach_ea(
    sys: LinSys,
    X0: Ellipsoid,
    L0: np.ndarray,
    t_span: tuple[float, float],
    N_steps: int = 100,
) -> ReachSet:
    """Compute external ellipsoidal approximation of the reach set.

    For each direction column in *L0*, integrates the external shape
    matrix ODE to produce a tight outer-bounding ellipsoid tube.

    Parameters
    ----------
    sys : LinSys
        The linear system (continuous-time).
    X0 : Ellipsoid
        Initial ellipsoidal set E(q0, Q0).
    L0 : ndarray (n, m)
        Direction matrix — each column is one approximation direction.
    t_span : (t0, tf)
        Time interval.
    N_steps : int
        Number of time grid points.

    Returns
    -------
    ReachSet
    """
    if sys.is_empty():
        raise ValueError("reach_ea: system must be non-empty.")
    if X0.is_empty():
        raise ValueError("reach_ea: initial set must be non-empty.")

    q0, Q0 = X0.parameters()
    assert q0 is not None and Q0 is not None
    n = q0.shape[0]

    L0 = np.asarray(L0, dtype=float)
    if L0.ndim == 1:
        L0 = L0.reshape(-1, 1)
    m = L0.shape[1]

    t0, tf = t_span
    t_eval = np.linspace(t0, tf, N_steps)

    # ── Precompute STM, center, bounds ───────────────────────────────
    Phi = compute_stm(sys, t_span, t_eval)
    Phi_inv = compute_stm_inv(Phi)
    centers = compute_center(sys, q0, t_span, t_eval)
    BPB, BPBsr, GQG, GQGsr = precompute_bounds(sys, t_eval)

    has_dist = GQG is not None

    # ── Integrate shape ODE for each direction ───────────────────────
    ea_shapes = np.zeros((N_steps, m, n, n))
    ea_shapes[0, :, :, :] = Q0[np.newaxis, :, :]

    for j in range(m):
        l0 = L0[:, j]
        Q_traj = _integrate_eesm(
            sys, Q0, l0, t_eval, Phi, Phi_inv,
            BPB, BPBsr, GQG, GQGsr, has_dist, n,
        )
        ea_shapes[:, j, :, :] = Q_traj

    rs = ReachSet(
        t_eval=t_eval,
        centers=centers,
        ea_shapes=ea_shapes,
        n_dirs=m,
        n=n,
        sys=sys,
        X0=X0,
        L0=L0,
    )
    return rs


# =====================================================================
# Shape matrix ODE integration
# =====================================================================

def _integrate_eesm(
    sys: LinSys,
    Q0: np.ndarray,
    l0: np.ndarray,
    t_eval: np.ndarray,
    Phi: np.ndarray,
    Phi_inv: np.ndarray,
    BPB: np.ndarray,
    BPBsr: np.ndarray,
    GQG: Optional[np.ndarray],
    GQGsr: Optional[np.ndarray],
    has_dist: bool,
    n: int,
) -> np.ndarray:
    """Integrate the external shape matrix ODE for one direction.

    Returns
    -------
    Q_traj : ndarray (N, n, n)
    """
    N = len(t_eval)
    t0, tf = t_eval[0], t_eval[-1]

    # For interpolation of precomputed data during ODE integration
    def _get_F(t: float) -> np.ndarray:
        """Get Φ^{-1}(t) by interpolating from precomputed grid."""
        idx = np.searchsorted(t_eval, t, side="right") - 1
        idx = np.clip(idx, 0, N - 2)
        return Phi_inv[idx]

    def _get_bpb(t: float) -> np.ndarray:
        idx = np.searchsorted(t_eval, t, side="right") - 1
        idx = np.clip(idx, 0, N - 2)
        return BPB[idx]

    def _get_gqg(t: float) -> np.ndarray:
        idx = np.searchsorted(t_eval, t, side="right") - 1
        idx = np.clip(idx, 0, N - 2)
        return GQG[idx] if GQG is not None else np.zeros((n, n))

    def _get_gqgsr(t: float) -> np.ndarray:
        idx = np.searchsorted(t_eval, t, side="right") - 1
        idx = np.clip(idx, 0, N - 2)
        return GQGsr[idx] if GQGsr is not None else np.zeros((n, n))

    if has_dist:
        rhs = lambda t, x: _eedist_ode_rhs(
            t, x, l0, sys, n, _get_F, _get_bpb, _get_gqg, _get_gqgsr,
        )
    else:
        rhs = lambda t, x: _eesm_ode_rhs(
            t, x, l0, sys, n, _get_F, _get_bpb,
        )

    sol = ode_solve(rhs, (t0, tf), Q0.ravel(), t_eval=t_eval)

    Q_traj = np.zeros((N, n, n))
    for k in range(sol.y.shape[1]):
        Qk = sol.y[:, k].reshape(n, n)
        Qk = 0.5 * (Qk + Qk.T)
        # Clamp tiny negative eigenvalues
        eigvals = np.linalg.eigvalsh(Qk)
        if np.min(eigvals) < -ABS_TOL:
            Qk = Qk - np.min(eigvals) * np.eye(n)
        Q_traj[k] = Qk

    return Q_traj


def _eesm_ode_rhs(
    t: float,
    x: np.ndarray,
    l0: np.ndarray,
    sys: LinSys,
    n: int,
    get_F,
    get_bpb,
) -> np.ndarray:
    """RHS of external shape matrix ODE (no disturbance).

    dQ/dt = A Q + Q A' + (p1/p2) Q + (p2/p1) BPB

    p1 = sqrt(l0' F BPB F' l0)
    p2 = sqrt(l0' F Q F' l0)
    """
    Q = x.reshape(n, n)
    Q = 0.5 * (Q + Q.T)

    A = sys.A_at(t)
    F = get_F(t)
    bpb = get_bpb(t)

    Fl = F.T @ l0
    p1_sq = max(float(Fl @ bpb @ Fl), 0.0)
    p2_sq = max(float(Fl @ Q @ Fl), 0.0)
    p1 = np.sqrt(p1_sq)
    p2 = np.sqrt(p2_sq)

    p1 = max(p1, ABS_TOL)
    p2 = max(p2, ABS_TOL)

    dQ = A @ Q + Q @ A.T + (p1 / p2) * Q + (p2 / p1) * bpb
    dQ = 0.5 * (dQ + dQ.T)
    return dQ.ravel()


def _eedist_ode_rhs(
    t: float,
    x: np.ndarray,
    l0: np.ndarray,
    sys: LinSys,
    n: int,
    get_F,
    get_bpb,
    get_gqg,
    get_gqgsr,
) -> np.ndarray:
    """RHS of external shape matrix ODE (with disturbance).

    dQ/dt = A Q + Q A' + (p1/p2) Q + (p2/p1) BPB - Z - Z'

    Z = Y S GQGsr,  Y = sqrtm(Q)
    S = ell_valign(Y F' l0, GQGsr F' l0)
    p1 = sqrt(l0' F BPB F' l0)
    p2 = sqrt(l0' F Q F' l0)
    """
    Q = x.reshape(n, n)
    Q = 0.5 * (Q + Q.T)

    A = sys.A_at(t)
    F = get_F(t)
    bpb = get_bpb(t)
    gqg = get_gqg(t)
    gqgsr = get_gqgsr(t)

    Fl = F.T @ l0

    p1_sq = max(float(Fl @ bpb @ Fl), 0.0)
    p2_sq = max(float(Fl @ Q @ Fl), 0.0)
    p1 = np.sqrt(p1_sq)
    p2 = np.sqrt(p2_sq)
    p1 = max(p1, ABS_TOL)
    p2 = max(p2, ABS_TOL)

    Y = _safe_sqrtm(Q)

    l1 = Y @ Fl
    l2 = gqgsr @ Fl

    if np.linalg.norm(l1) < ABS_TOL or np.linalg.norm(l2) < ABS_TOL:
        S = np.eye(n)
    else:
        S = ell_valign(l1, l2)

    Z = Y @ S @ gqgsr

    dQ = (A @ Q + Q @ A.T
          + (p1 / p2) * Q
          + (p2 / p1) * bpb
          - Z - Z.T)
    dQ = 0.5 * (dQ + dQ.T)
    return dQ.ravel()


# =====================================================================
# Internal approximation
# =====================================================================

def reach_ia(
    sys: LinSys,
    X0: Ellipsoid,
    L0: np.ndarray,
    t_span: tuple[float, float],
    N_steps: int = 100,
) -> ReachSet:
    """Compute internal ellipsoidal approximation of the reach set.

    For each direction column in *L0*, integrates the internal shape
    matrix ODE to produce a tight inner-bounding ellipsoid tube.

    Parameters
    ----------
    sys : LinSys
        The linear system (continuous-time).
    X0 : Ellipsoid
        Initial ellipsoidal set E(q0, Q0).
    L0 : ndarray (n, m)
        Direction matrix — each column is one approximation direction.
    t_span : (t0, tf)
        Time interval.
    N_steps : int
        Number of time grid points.

    Returns
    -------
    ReachSet
    """
    if sys.is_empty():
        raise ValueError("reach_ia: system must be non-empty.")
    if X0.is_empty():
        raise ValueError("reach_ia: initial set must be non-empty.")

    q0, Q0 = X0.parameters()
    assert q0 is not None and Q0 is not None
    n = q0.shape[0]

    L0 = np.asarray(L0, dtype=float)
    if L0.ndim == 1:
        L0 = L0.reshape(-1, 1)
    m = L0.shape[1]

    t0, tf = t_span
    t_eval = np.linspace(t0, tf, N_steps)

    # ── Precompute STM, center, bounds ───────────────────────────────
    Phi = compute_stm(sys, t_span, t_eval)
    Phi_inv = compute_stm_inv(Phi)
    centers = compute_center(sys, q0, t_span, t_eval)
    BPB, BPBsr, GQG, GQGsr = precompute_bounds(sys, t_eval)

    has_dist = GQG is not None

    # ── Integrate shape ODE for each direction ───────────────────────
    ia_shapes = np.zeros((N_steps, m, n, n))
    ia_shapes[0, :, :, :] = Q0[np.newaxis, :, :]

    for j in range(m):
        l0 = L0[:, j]
        Q_traj = _integrate_iesm(
            sys, Q0, l0, t_eval, Phi, Phi_inv,
            BPB, BPBsr, GQG, GQGsr, has_dist, n,
        )
        ia_shapes[:, j, :, :] = Q_traj

    rs = ReachSet(
        t_eval=t_eval,
        centers=centers,
        ia_shapes=ia_shapes,
        n_dirs=m,
        n=n,
        sys=sys,
        X0=X0,
        L0=L0,
    )
    return rs


# =====================================================================
# Internal shape matrix ODE integration
# =====================================================================

def _integrate_iesm(
    sys: LinSys,
    Q0: np.ndarray,
    l0: np.ndarray,
    t_eval: np.ndarray,
    Phi: np.ndarray,
    Phi_inv: np.ndarray,
    BPB: np.ndarray,
    BPBsr: np.ndarray,
    GQG: Optional[np.ndarray],
    GQGsr: Optional[np.ndarray],
    has_dist: bool,
    n: int,
) -> np.ndarray:
    """Integrate the internal shape matrix ODE for one direction.

    Returns
    -------
    Q_traj : ndarray (N, n, n)
    """
    N = len(t_eval)
    t0, tf = t_eval[0], t_eval[-1]

    def _get_F(t: float) -> np.ndarray:
        idx = np.searchsorted(t_eval, t, side="right") - 1
        idx = np.clip(idx, 0, N - 2)
        return Phi_inv[idx]

    def _get_bpbsr(t: float) -> np.ndarray:
        idx = np.searchsorted(t_eval, t, side="right") - 1
        idx = np.clip(idx, 0, N - 2)
        return BPBsr[idx]

    def _get_gqg(t: float) -> np.ndarray:
        idx = np.searchsorted(t_eval, t, side="right") - 1
        idx = np.clip(idx, 0, N - 2)
        return GQG[idx] if GQG is not None else np.zeros((n, n))

    if has_dist:
        rhs = lambda t, x: _iedist_ode_rhs(
            t, x, l0, sys, n, _get_F, _get_bpbsr, _get_gqg,
        )
    else:
        rhs = lambda t, x: _iesm_ode_rhs(
            t, x, l0, sys, n, _get_F, _get_bpbsr,
        )

    sol = ode_solve(rhs, (t0, tf), Q0.ravel(), t_eval=t_eval)

    Q_traj = np.zeros((N, n, n))
    for k in range(sol.y.shape[1]):
        Qk = sol.y[:, k].reshape(n, n)
        Qk = 0.5 * (Qk + Qk.T)
        eigvals = np.linalg.eigvalsh(Qk)
        if np.min(eigvals) < -ABS_TOL:
            Qk = Qk - np.min(eigvals) * np.eye(n)
        Q_traj[k] = Qk

    return Q_traj


def _iesm_ode_rhs(
    t: float,
    x: np.ndarray,
    l0: np.ndarray,
    sys: LinSys,
    n: int,
    get_F,
    get_bpbsr,
) -> np.ndarray:
    """RHS of internal shape matrix ODE (no disturbance).

    dQ/dt = A Q + Q A' + Z + Z'

    Z = sqrtm(Q) · S · BPBsr
    S = ell_valign(BPBsr F' l0,  sqrtm(Q) F' l0)

    This is the Minkowski-sum alignment formula: the IA grows by
    aligning the control ellipsoid's square root with Q's square root.
    """
    Q = x.reshape(n, n)
    Q = 0.5 * (Q + Q.T)

    A = sys.A_at(t)
    F = get_F(t)
    bpbsr = get_bpbsr(t)

    Fl = F.T @ l0
    Y = _safe_sqrtm(Q)

    l1 = bpbsr @ Fl
    l2 = Y @ Fl

    if np.linalg.norm(l1) < ABS_TOL or np.linalg.norm(l2) < ABS_TOL:
        S = np.eye(n)
    else:
        S = ell_valign(l1, l2)

    Z = Y @ S @ bpbsr

    dQ = A @ Q + Q @ A.T + Z + Z.T
    dQ = 0.5 * (dQ + dQ.T)
    return dQ.ravel()


def _iedist_ode_rhs(
    t: float,
    x: np.ndarray,
    l0: np.ndarray,
    sys: LinSys,
    n: int,
    get_F,
    get_bpbsr,
    get_gqg,
) -> np.ndarray:
    """RHS of internal shape matrix ODE (with disturbance).

    dQ/dt = A Q + Q A' + Z + Z' - (p1/p2) Q - (p2/p1) GQG

    Z = sqrtm(Q) · S · BPBsr        (Mink sum with control, alignment)
    S = ell_valign(BPBsr F' l0, sqrtm(Q) F' l0)

    p1 = sqrt(l0' F GQG F' l0)       (Mink diff with disturbance, parametric)
    p2 = sqrt(l0' F Q F' l0)

    The disturbance shrinks the IA via the parametric Minkowski
    difference formula (dual of how the EA handles control).
    """
    Q = x.reshape(n, n)
    Q = 0.5 * (Q + Q.T)

    A = sys.A_at(t)
    F = get_F(t)
    bpbsr = get_bpbsr(t)
    gqg = get_gqg(t)

    Fl = F.T @ l0

    # ── Mink sum with control (alignment) ────────────────────────────
    Y = _safe_sqrtm(Q)

    l1 = bpbsr @ Fl
    l2 = Y @ Fl

    if np.linalg.norm(l1) < ABS_TOL or np.linalg.norm(l2) < ABS_TOL:
        S = np.eye(n)
    else:
        S = ell_valign(l1, l2)

    Z = Y @ S @ bpbsr

    # ── Mink diff with disturbance (parametric) ──────────────────────
    p1_sq = max(float(Fl @ gqg @ Fl), 0.0)
    p2_sq = max(float(Fl @ Q @ Fl), 0.0)
    p1 = np.sqrt(p1_sq)
    p2 = np.sqrt(p2_sq)
    p1 = max(p1, ABS_TOL)
    p2 = max(p2, ABS_TOL)

    dQ = (A @ Q + Q @ A.T
          + Z + Z.T
          - (p1 / p2) * Q
          - (p2 / p1) * gqg)
    dQ = 0.5 * (dQ + dQ.T)
    return dQ.ravel()


# =====================================================================
# Discrete-time reach sets
# =====================================================================

def _discrete_precompute(
    sys: LinSys,
    N_steps: int,
) -> tuple[np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
    """Precompute BPB, BPBsr, GQG, GQGsr for the discrete-time case.

    For LTI discrete systems the bound matrices are constant, so we
    evaluate once at t=0 and replicate.
    """
    n = sys._n
    BPB = np.zeros((N_steps, n, n))
    BPBsr = np.zeros((N_steps, n, n))

    has_dist = sys.has_disturbance()
    GQG: Optional[np.ndarray] = np.zeros((N_steps, n, n)) if has_dist else None
    GQGsr: Optional[np.ndarray] = np.zeros((N_steps, n, n)) if has_dist else None

    for k in range(N_steps):
        t = float(k)
        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 _sanitize_shape(Q: np.ndarray, n: int) -> np.ndarray:
    """Symmetrize and clamp negative eigenvalues."""
    Q = 0.5 * (Q + Q.T)
    eigvals = np.linalg.eigvalsh(Q)
    if np.min(eigvals) < -ABS_TOL:
        Q = Q - np.min(eigvals) * np.eye(n)
    return Q


def reach_ea_d(
    sys: LinSys,
    X0: Ellipsoid,
    L0: np.ndarray,
    N_steps: int = 10,
) -> ReachSet:
    """Compute discrete-time external ellipsoidal approximation.

    For the system  x[k+1] = A x[k] + B u[k]  (+ G v[k]),
    iterates the EA shape recurrence directly (no ODE solving).

    **EA recurrence (no disturbance):**

    .. math::
        Q_{k+1} = (1 + 1/p)\\,A Q_k A^\\top + (1 + p)\\,BPB

    where  p = \\sqrt{l_k^\\top BPB\\, l_k} / \\sqrt{l_k^\\top A Q_k A^\\top l_k},
    l_{k+1} = A^{-\\top} l_k.

    **EA recurrence (with disturbance):**

    Like above for the Mink sum with control, then subtract the
    alignment-based Mink difference with GQG.

    Parameters
    ----------
    sys : LinSys
        Discrete-time linear system (``sys.discrete == True``).
    X0 : Ellipsoid
        Initial set.
    L0 : ndarray (n, m)
        Direction matrix.
    N_steps : int
        Number of discrete time steps (output has N_steps+1 entries,
        from k=0 to k=N_steps).

    Returns
    -------
    ReachSet
    """
    if sys.is_empty():
        raise ValueError("reach_ea_d: system must be non-empty.")
    if X0.is_empty():
        raise ValueError("reach_ea_d: initial set must be non-empty.")

    q0, Q0 = X0.parameters()
    assert q0 is not None and Q0 is not None
    n = q0.shape[0]

    L0 = np.asarray(L0, dtype=float)
    if L0.ndim == 1:
        L0 = L0.reshape(-1, 1)
    m = L0.shape[1]

    N = N_steps + 1  # k = 0, 1, ..., N_steps
    t_eval = np.arange(N, dtype=float)

    BPB, BPBsr, GQG, GQGsr = _discrete_precompute(sys, N)
    has_dist = GQG is not None

    A = sys.A_at(0.0)
    A_inv_T = np.linalg.inv(A).T

    # ── Centers: q[k+1] = A q[k] + B p[k] + G d[k] ─────────────────
    centers = np.zeros((N, n))
    centers[0] = q0
    for k in range(N_steps):
        q_k = centers[k]
        q_next = A @ q_k
        B_k = sys.B_at(float(k))
        if B_k is not None:
            U_k = sys.U_at(float(k))
            if U_k is not None and not U_k.is_empty():
                q_next += B_k @ U_k.center
        G_k = sys.G_at(float(k))
        if G_k is not None and sys.has_disturbance():
            V_k = sys.V_at(float(k))
            if V_k is not None and not V_k.is_empty():
                q_next += G_k @ V_k.center
        centers[k + 1] = q_next

    # ── EA shape iteration ───────────────────────────────────────────
    ea_shapes = np.zeros((N, m, n, n))
    ea_shapes[0, :, :, :] = Q0[np.newaxis, :, :]

    for j in range(m):
        l = L0[:, j].copy()
        Q = Q0.copy()

        for k in range(N_steps):
            bpb = BPB[k]
            AQA = A @ Q @ A.T
            AQA = 0.5 * (AQA + AQA.T)

            # Parametric Mink sum: Q+ = (1+1/p)*AQA + (1+p)*BPB
            lAQAl = max(float(l @ AQA @ l), 0.0)
            lBPBl = max(float(l @ bpb @ l), 0.0)
            s_AQA = np.sqrt(lAQAl)
            s_BPB = np.sqrt(lBPBl)
            s_AQA = max(s_AQA, ABS_TOL)
            s_BPB = max(s_BPB, ABS_TOL)

            p = s_BPB / s_AQA
            Q_next = (1.0 + 1.0 / p) * AQA + (1.0 + p) * bpb

            if has_dist:
                gqg = GQG[k]
                gqgsr = GQGsr[k]
                # Alignment Mink diff: subtract Z + Z'
                Y = _safe_sqrtm(Q_next)
                l1 = Y @ l
                l2 = gqgsr @ l
                if np.linalg.norm(l1) < ABS_TOL or np.linalg.norm(l2) < ABS_TOL:
                    S = np.eye(n)
                else:
                    S = ell_valign(l1, l2)
                Z = Y @ S @ gqgsr
                Q_next = Q_next - Z - Z.T

            Q_next = _sanitize_shape(Q_next, n)
            ea_shapes[k + 1, j] = Q_next
            Q = Q_next
            l = A_inv_T @ l

    return ReachSet(
        t_eval=t_eval,
        centers=centers,
        ea_shapes=ea_shapes,
        n_dirs=m,
        n=n,
        sys=sys,
        X0=X0,
        L0=L0,
    )


def reach_ia_d(
    sys: LinSys,
    X0: Ellipsoid,
    L0: np.ndarray,
    N_steps: int = 10,
) -> ReachSet:
    """Compute discrete-time internal ellipsoidal approximation.

    **IA recurrence (no disturbance):**

    .. math::
        Q_{k+1} = (\\sqrt{A Q_k A^\\top} + S\\,\\sqrt{BPB})^\\top
                   (\\sqrt{A Q_k A^\\top} + S\\,\\sqrt{BPB})

    where S = ell_valign(BPBsr l_k, sqrtm(AQA) l_k).

    **IA recurrence (with disturbance):**

    Like above for the Mink sum with control, then subtract the
    parametric Mink difference with GQG.

    Parameters
    ----------
    sys : LinSys
        Discrete-time linear system.
    X0 : Ellipsoid
        Initial set.
    L0 : ndarray (n, m)
        Direction matrix.
    N_steps : int
        Number of discrete time steps.

    Returns
    -------
    ReachSet
    """
    if sys.is_empty():
        raise ValueError("reach_ia_d: system must be non-empty.")
    if X0.is_empty():
        raise ValueError("reach_ia_d: initial set must be non-empty.")

    q0, Q0 = X0.parameters()
    assert q0 is not None and Q0 is not None
    n = q0.shape[0]

    L0 = np.asarray(L0, dtype=float)
    if L0.ndim == 1:
        L0 = L0.reshape(-1, 1)
    m = L0.shape[1]

    N = N_steps + 1
    t_eval = np.arange(N, dtype=float)

    BPB, BPBsr, GQG, GQGsr = _discrete_precompute(sys, N)
    has_dist = GQG is not None

    A = sys.A_at(0.0)
    A_inv_T = np.linalg.inv(A).T

    # ── Centers (same as EA) ─────────────────────────────────────────
    centers = np.zeros((N, n))
    centers[0] = q0
    for k in range(N_steps):
        q_k = centers[k]
        q_next = A @ q_k
        B_k = sys.B_at(float(k))
        if B_k is not None:
            U_k = sys.U_at(float(k))
            if U_k is not None and not U_k.is_empty():
                q_next += B_k @ U_k.center
        G_k = sys.G_at(float(k))
        if G_k is not None and sys.has_disturbance():
            V_k = sys.V_at(float(k))
            if V_k is not None and not V_k.is_empty():
                q_next += G_k @ V_k.center
        centers[k + 1] = q_next

    # ── IA shape iteration ───────────────────────────────────────────
    ia_shapes = np.zeros((N, m, n, n))
    ia_shapes[0, :, :, :] = Q0[np.newaxis, :, :]

    for j in range(m):
        l = L0[:, j].copy()
        Q = Q0.copy()

        for k in range(N_steps):
            bpbsr = BPBsr[k]
            AQA = A @ Q @ A.T
            AQA = 0.5 * (AQA + AQA.T)

            # Alignment Mink sum: Q+ = (Y + S * BPBsr)^T (Y + S * BPBsr)
            Y = _safe_sqrtm(AQA)
            l1 = bpbsr @ l
            l2 = Y @ l
            if np.linalg.norm(l1) < ABS_TOL or np.linalg.norm(l2) < ABS_TOL:
                S = np.eye(n)
            else:
                S = ell_valign(l1, l2)

            M = Y + S @ bpbsr
            Q_next = M.T @ M

            if has_dist:
                gqg = GQG[k]
                # Parametric Mink diff: subtract (p1/p2)*Q + (p2/p1)*GQG
                lQl = max(float(l @ Q_next @ l), 0.0)
                lGl = max(float(l @ gqg @ l), 0.0)
                s_Q = np.sqrt(lQl)
                s_G = np.sqrt(lGl)
                s_Q = max(s_Q, ABS_TOL)
                s_G = max(s_G, ABS_TOL)
                Q_next = Q_next - (s_G / s_Q) * Q_next - (s_Q / s_G) * gqg
                # Simplifies to: Q_next * (1 - s_G/s_Q) - (s_Q/s_G)*gqg
                # which can go negative if disturbance > control.

            Q_next = _sanitize_shape(Q_next, n)
            ia_shapes[k + 1, j] = Q_next
            Q = Q_next
            l = A_inv_T @ l

    return ReachSet(
        t_eval=t_eval,
        centers=centers,
        ia_shapes=ia_shapes,
        n_dirs=m,
        n=n,
        sys=sys,
        X0=X0,
        L0=L0,
    )
