"""Reach tube visualization.

Provides functions to plot EA/IA reach tubes as stacked time-indexed
ellipsoid cross-sections.

For a 2D system (n=2), the reach tube is rendered as a 3D plot with
two spatial axes (x₁, x₂) and one time axis.  Each time step
contributes an ellipse boundary at its corresponding time height.

For a 1D system (n=1), the tube is a 2D band plot over time.

For higher-dimensional systems, project the ReachSet to 2D first
via ``ReachSet.projection(basis)``.

Reference
---------
MATLAB Ellipsoidal Toolbox v3, ``control/@reach/plot_ea.m``,
``control/@reach/plot_ia.m``.
"""

from __future__ import annotations

from typing import Optional, Sequence

import numpy as np
import scipy.linalg

from ..core.config import ABS_TOL
from ..core.ellipsoid import Ellipsoid
from ..core.helpers import regularize
from ..systems.reach import ReachSet
from .plot_ellipsoid import _boundary_2d, _N_PTS_2D

__all__ = [
    "plot_ea_tube",
    "plot_ia_tube",
    "plot_reach",
]

_TUBE_PTS = 80  # boundary points per ellipse cross-section


# =====================================================================
# Public API
# =====================================================================

def plot_ea_tube(
    rs: ReachSet,
    ax=None,
    color: Optional[str] = None,
    alpha: float = 0.15,
    center: bool = True,
    center_color: str = "black",
    n_pts: int = _TUBE_PTS,
    title: Optional[str] = None,
    show: bool = False,
):
    """Plot the external approximation tube.

    Parameters
    ----------
    rs : ReachSet
        Must have ``ea_shapes`` computed.
    ax : matplotlib Axes or None
        Target axes.  Created automatically if None.
    color : str or None
        Color for the tube surface / lines.
    alpha : float
        Transparency.
    center : bool
        If True, overlay the center trajectory.
    center_color : str
        Color for the center trajectory line.
    n_pts : int
        Boundary resolution per cross-section.
    title : str or None
    show : bool

    Returns
    -------
    ax : matplotlib Axes
    """
    if rs.ea_shapes is None:
        raise ValueError("plot_ea_tube: no EA data in ReachSet.")
    return _plot_tube(
        rs, "ea", ax, color or "C0", alpha, center, center_color,
        n_pts, title, show,
    )


def plot_ia_tube(
    rs: ReachSet,
    ax=None,
    color: Optional[str] = None,
    alpha: float = 0.15,
    center: bool = True,
    center_color: str = "black",
    n_pts: int = _TUBE_PTS,
    title: Optional[str] = None,
    show: bool = False,
):
    """Plot the internal approximation tube.

    Parameters
    ----------
    rs : ReachSet
        Must have ``ia_shapes`` computed.
    ax, color, alpha, center, center_color, n_pts, title, show
        Same as :func:`plot_ea_tube`.

    Returns
    -------
    ax : matplotlib Axes
    """
    if rs.ia_shapes is None:
        raise ValueError("plot_ia_tube: no IA data in ReachSet.")
    return _plot_tube(
        rs, "ia", ax, color or "C1", alpha, center, center_color,
        n_pts, title, show,
    )


def plot_reach(
    rs: ReachSet,
    ax=None,
    ea_color: str = "C0",
    ia_color: str = "C1",
    alpha: float = 0.15,
    center: bool = True,
    center_color: str = "black",
    n_pts: int = _TUBE_PTS,
    title: Optional[str] = None,
    show: bool = False,
):
    """Plot both EA and IA tubes on the same axes.

    Skips whichever approximation is not computed.

    Parameters
    ----------
    rs : ReachSet
    ax : matplotlib Axes or None
    ea_color, ia_color : str
    alpha : float
    center : bool
    center_color : str
    n_pts : int
    title : str or None
    show : bool

    Returns
    -------
    ax : matplotlib Axes
    """
    has_ea = rs.ea_shapes is not None
    has_ia = rs.ia_shapes is not None
    if not has_ea and not has_ia:
        raise ValueError("plot_reach: no EA or IA data in ReachSet.")

    if has_ea:
        ax = _plot_tube(
            rs, "ea", ax, ea_color, alpha, center, center_color,
            n_pts, None, False,
        )
        center = False  # don't draw center twice

    if has_ia:
        ax = _plot_tube(
            rs, "ia", ax, ia_color, alpha, False, center_color,
            n_pts, None, False,
        )

    import matplotlib.pyplot as plt
    if title:
        ax.set_title(title)
    if show:
        plt.show()
    return ax


# =====================================================================
# Internal implementation
# =====================================================================

def _plot_tube(rs, kind, ax, color, alpha, draw_center, center_color,
               n_pts, title, show):
    """Core tube plotter, dispatches by state dimension."""
    n = rs.n
    if n == 1:
        return _plot_tube_1d(rs, kind, ax, color, alpha, draw_center,
                             center_color, title, show)
    elif n == 2:
        return _plot_tube_2d(rs, kind, ax, color, alpha, draw_center,
                             center_color, n_pts, title, show)
    else:
        raise ValueError(
            f"plot_{kind}_tube: state dimension is {n}. "
            f"Use rs.projection(basis) to reduce to 1D or 2D first."
        )


def _plot_tube_1d(rs, kind, ax, color, alpha, draw_center, center_color,
                  title, show):
    """1D reach tube: band plot over time."""
    import matplotlib.pyplot as plt

    if ax is None:
        fig, ax = plt.subplots()

    t = rs.t_eval
    N = len(t)
    shapes = rs.ea_shapes if kind == "ea" else rs.ia_shapes

    upper = np.zeros(N)
    lower = np.zeros(N)
    for k in range(N):
        q = rs.centers[k, 0]
        # Take union of all directions (envelope)
        r_max = 0.0
        for j in range(rs.n_dirs):
            Q_val = float(shapes[k, j, 0, 0])
            r_max = max(r_max, np.sqrt(max(Q_val, 0.0)))
        upper[k] = q + r_max
        lower[k] = q - r_max

    ax.fill_between(t, lower, upper, alpha=alpha, color=color,
                    label=f"{kind.upper()} tube")
    ax.plot(t, upper, color=color, linewidth=0.8)
    ax.plot(t, lower, color=color, linewidth=0.8)

    if draw_center:
        ax.plot(t, rs.centers[:, 0], color=center_color, linewidth=1.5,
                linestyle="--", label="center")

    ax.set_xlabel("t")
    ax.set_ylabel("x")
    if title:
        ax.set_title(title)
    if show:
        plt.show()
    return ax


def _plot_tube_2d(rs, kind, ax, color, alpha, draw_center, center_color,
                  n_pts, title, show):
    """2D reach tube: 3D plot (x₁, x₂, t)."""
    import matplotlib.pyplot as plt
    try:
        import mpl_toolkits.mplot3d  # noqa: F401
    except ImportError:
        raise ImportError(
            "3D reach tube plotting requires mpl_toolkits.mplot3d. "
            "Install a matching matplotlib version."
        )

    if ax is None:
        fig = plt.figure(figsize=(10, 7))
        ax = fig.add_subplot(111, projection="3d")

    t = rs.t_eval
    N = len(t)
    shapes = rs.ea_shapes if kind == "ea" else rs.ia_shapes

    theta = np.linspace(0, 2 * np.pi, n_pts)
    circle = np.column_stack([np.cos(theta), np.sin(theta)])

    for k in range(N):
        q = rs.centers[k]
        for j in range(rs.n_dirs):
            Q = shapes[k, j]
            Q = 0.5 * (Q + Q.T)
            rk_mat = np.linalg.matrix_rank(Q, tol=ABS_TOL)
            if rk_mat < 2:
                Q = regularize(Q)
            S = np.real(scipy.linalg.sqrtm(Q))
            pts = circle @ S.T + q[np.newaxis, :]
            ax.plot(
                pts[:, 0], pts[:, 1],
                zs=t[k], zdir="z",
                color=color, alpha=max(alpha * 2, 0.3), linewidth=0.6,
            )

    # Fill between consecutive cross-sections for a surface effect
    if N > 1 and rs.n_dirs > 0:
        for j in range(rs.n_dirs):
            X_surf = np.zeros((N, n_pts))
            Y_surf = np.zeros((N, n_pts))
            T_surf = np.zeros((N, n_pts))
            for k in range(N):
                q = rs.centers[k]
                Q = shapes[k, j]
                Q = 0.5 * (Q + Q.T)
                rk_mat = np.linalg.matrix_rank(Q, tol=ABS_TOL)
                if rk_mat < 2:
                    Q = regularize(Q)
                S = np.real(scipy.linalg.sqrtm(Q))
                pts = circle @ S.T + q[np.newaxis, :]
                X_surf[k] = pts[:, 0]
                Y_surf[k] = pts[:, 1]
                T_surf[k] = t[k]

            ax.plot_surface(
                X_surf, Y_surf, T_surf,
                color=color, alpha=alpha, shade=False,
            )

    if draw_center:
        ax.plot(
            rs.centers[:, 0], rs.centers[:, 1], t,
            color=center_color, linewidth=2, linestyle="--",
            label="center",
        )

    ax.set_xlabel("x₁")
    ax.set_ylabel("x₂")
    ax.set_zlabel("t")
    if title:
        ax.set_title(title)
    if show:
        plt.show()
    return ax
