"""Ellipsoid plotting in 1D, 2D, and 3D.

Provides functions to visualize individual or multiple ellipsoids
using Matplotlib.

1D
    Plots a horizontal interval [q − √Q, q + √Q].

2D
    Generates boundary points  x(θ) = q + Q^{1/2} [cos θ, sin θ]^T
    for θ ∈ [0, 2π).  Optionally fills the interior.

3D
    Generates a mesh on the unit sphere, transforms via q + Q^{1/2} s,
    and renders as a surface.

Reference
---------
MATLAB Ellipsoidal Toolbox v3, ``@ellipsoid/plot.m``.
"""

from __future__ import annotations

from typing import Optional, Sequence, Union

import numpy as np
import scipy.linalg

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

__all__ = [
    "plot_ellipsoid",
    "plot_ellipsoids",
]

# Default visual parameters
_N_PTS_2D = 200
_N_PTS_3D = 30  # grid resolution per axis for sphere mesh


# =====================================================================
# Boundary-point generators
# =====================================================================

def _boundary_1d(E: Ellipsoid) -> tuple[float, float]:
    """Return the (lo, hi) interval for a 1D ellipsoid."""
    q = float(E.center[0])
    Q = float(E.shape[0, 0])
    r = np.sqrt(max(Q, 0.0))
    return q - r, q + r


def _boundary_2d(E: Ellipsoid, n_pts: int = _N_PTS_2D) -> np.ndarray:
    """Return boundary points of a 2D ellipsoid.

    Returns
    -------
    pts : ndarray (n_pts, 2)
        Boundary coordinates.
    """
    q = E.center
    Q = E.shape.copy()
    sd, rk = E.dimension()
    if rk < sd:
        Q = regularize(Q)

    S = np.real(scipy.linalg.sqrtm(Q))
    theta = np.linspace(0, 2 * np.pi, n_pts)
    circle = np.column_stack([np.cos(theta), np.sin(theta)])  # (n_pts, 2)
    pts = circle @ S.T + q[np.newaxis, :]
    return pts


def _boundary_3d(
    E: Ellipsoid, n_pts: int = _N_PTS_3D,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Return surface mesh arrays (X, Y, Z) for a 3D ellipsoid.

    Returns
    -------
    X, Y, Z : ndarray (n_pts, n_pts)
        Mesh-grid coordinate arrays suitable for ``plot_surface``.
    """
    q = E.center
    Q = E.shape.copy()
    sd, rk = E.dimension()
    if rk < sd:
        Q = regularize(Q)

    S = np.real(scipy.linalg.sqrtm(Q))

    u = np.linspace(0, 2 * np.pi, n_pts)
    v = np.linspace(0, np.pi, n_pts)

    # Unit sphere parametrization
    sx = np.outer(np.cos(u), np.sin(v))
    sy = np.outer(np.sin(u), np.sin(v))
    sz = np.outer(np.ones_like(u), np.cos(v))

    # Transform each surface point: q + S @ [sx, sy, sz]
    X = np.zeros_like(sx)
    Y = np.zeros_like(sy)
    Z = np.zeros_like(sz)

    for i in range(n_pts):
        for j in range(n_pts):
            pt = S @ np.array([sx[i, j], sy[i, j], sz[i, j]]) + q
            X[i, j] = pt[0]
            Y[i, j] = pt[1]
            Z[i, j] = pt[2]

    return X, Y, Z


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

def plot_ellipsoid(
    E: Ellipsoid,
    ax=None,
    color: Optional[str] = None,
    fill: bool = True,
    alpha: float = 0.3,
    linewidth: float = 1.5,
    linestyle: str = "-",
    label: Optional[str] = None,
    n_pts: Optional[int] = None,
    wireframe: bool = False,
    show: bool = False,
):
    """Plot a single ellipsoid.

    Automatically detects 1D / 2D / 3D from the ellipsoid dimension.

    Parameters
    ----------
    E : Ellipsoid
        The ellipsoid to plot.
    ax : matplotlib Axes or None
        Target axes.  If *None*, a new figure is created.
    color : str or None
        Color string (e.g. ``'blue'``, ``'#ff0000'``).  If *None*,
        Matplotlib default color cycle is used.
    fill : bool
        If True, fill the interior (2D: filled polygon, 3D: opaque surface).
    alpha : float
        Transparency (0 = transparent, 1 = opaque).
    linewidth : float
        Boundary line width (2D).
    linestyle : str
        Boundary line style (2D).
    label : str or None
        Legend label.
    n_pts : int or None
        Number of boundary/mesh points (overrides defaults).
    wireframe : bool
        If True, render 3D ellipsoid as wireframe instead of surface.
    show : bool
        If True, call ``plt.show()`` after plotting.

    Returns
    -------
    ax : matplotlib Axes
        The axes used for plotting.
    """
    import matplotlib.pyplot as plt

    if E.is_empty():
        raise ValueError("Cannot plot an empty ellipsoid.")

    sd, _ = E.dimension()

    if sd == 1:
        return _plot_1d(E, ax, color, alpha, linewidth, label, show)
    elif sd == 2:
        return _plot_2d(E, ax, color, fill, alpha, linewidth, linestyle,
                        label, n_pts, show)
    elif sd == 3:
        return _plot_3d(E, ax, color, fill, alpha, label, n_pts,
                        wireframe, show)
    else:
        raise ValueError(
            f"Cannot plot a {sd}D ellipsoid directly. "
            f"Use projection() to reduce to 1D/2D/3D first."
        )


def plot_ellipsoids(
    ellipsoids: Sequence[Ellipsoid],
    ax=None,
    colors: Optional[Sequence[str]] = None,
    fill: bool = True,
    alpha: float = 0.3,
    linewidth: float = 1.5,
    labels: Optional[Sequence[str]] = None,
    n_pts: Optional[int] = None,
    wireframe: bool = False,
    title: Optional[str] = None,
    show: bool = False,
):
    """Plot multiple ellipsoids on the same axes.

    All ellipsoids must have the same dimension.

    Parameters
    ----------
    ellipsoids : sequence of Ellipsoid
    ax : matplotlib Axes or None
    colors : sequence of str or None
        One color per ellipsoid.  Cycles the default palette if None.
    fill : bool
    alpha : float
    linewidth : float
    labels : sequence of str or None
    n_pts : int or None
    wireframe : bool
        3D only: wireframe vs surface.
    title : str or None
    show : bool

    Returns
    -------
    ax : matplotlib Axes
    """
    import matplotlib.pyplot as plt

    if not ellipsoids:
        raise ValueError("No ellipsoids to plot.")

    dims = set()
    for E in ellipsoids:
        if E.is_empty():
            continue
        sd, _ = E.dimension()
        dims.add(sd)
    if len(dims) != 1:
        raise ValueError(
            f"All ellipsoids must have the same dimension; got {dims}."
        )
    sd = dims.pop()

    # Resolve default color cycle
    prop_cycle = plt.rcParams["axes.prop_cycle"]
    default_colors = [p["color"] for p in prop_cycle]
    n = len(ellipsoids)

    if colors is None:
        colors = [default_colors[i % len(default_colors)] for i in range(n)]
    if labels is None:
        labels = [None] * n

    # Create axes if needed
    if ax is None:
        if sd == 3:
            import mpl_toolkits.mplot3d  # noqa: F401
            fig = plt.figure()
            ax = fig.add_subplot(111, projection="3d")
        else:
            fig, ax = plt.subplots()

    for i, E in enumerate(ellipsoids):
        if E.is_empty():
            continue
        plot_ellipsoid(
            E,
            ax=ax,
            color=colors[i],
            fill=fill,
            alpha=alpha,
            linewidth=linewidth,
            label=labels[i],
            n_pts=n_pts,
            wireframe=wireframe,
        )

    if title:
        ax.set_title(title)

    if any(lb is not None for lb in labels):
        ax.legend()

    if sd == 2:
        ax.set_aspect("equal")
    elif sd == 3:
        _set_3d_equal_aspect(ax)

    if show:
        plt.show()

    return ax


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

def _plot_1d(E, ax, color, alpha, linewidth, label, show):
    import matplotlib.pyplot as plt

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

    lo, hi = _boundary_1d(E)
    kwargs = {}
    if color:
        kwargs["color"] = color
    if label:
        kwargs["label"] = label

    ax.plot([lo, hi], [0, 0], linewidth=linewidth * 2, solid_capstyle="round",
            **kwargs)
    ax.plot([lo, lo], [-0.05, 0.05], linewidth=linewidth, **kwargs)
    ax.plot([hi, hi], [-0.05, 0.05], linewidth=linewidth,
            color=kwargs.get("color"))

    ax.set_ylim(-0.3, 0.3)
    ax.set_xlabel("x")
    ax.axhline(0, color="gray", linewidth=0.5, zorder=0)

    if show:
        plt.show()
    return ax


def _plot_2d(E, ax, color, fill, alpha, linewidth, linestyle, label, n_pts, show):
    import matplotlib.pyplot as plt

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

    pts_count = n_pts or _N_PTS_2D
    pts = _boundary_2d(E, pts_count)

    kwargs = {}
    if color:
        kwargs["color"] = color
    if label:
        kwargs["label"] = label

    if fill:
        from matplotlib.patches import Polygon
        poly = Polygon(pts, closed=True, alpha=alpha,
                       facecolor=color or "C0",
                       edgecolor=color or "C0",
                       linewidth=linewidth, linestyle=linestyle)
        if label:
            poly.set_label(label)
        ax.add_patch(poly)
        ax.autoscale_view()
    else:
        ax.plot(pts[:, 0], pts[:, 1], linewidth=linewidth,
                linestyle=linestyle, **kwargs)

    ax.set_xlabel("x₁")
    ax.set_ylabel("x₂")

    if show:
        plt.show()
    return ax


def _plot_3d(E, ax, color, fill, alpha, label, n_pts, wireframe, show):
    import matplotlib.pyplot as plt
    import mpl_toolkits.mplot3d  # noqa: F401 — registers the '3d' projection

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

    pts_count = n_pts or _N_PTS_3D
    X, Y, Z = _boundary_3d(E, pts_count)

    if wireframe:
        ax.plot_wireframe(X, Y, Z, color=color or "C0", alpha=alpha,
                          linewidth=0.5, label=label)
    else:
        ax.plot_surface(X, Y, Z, color=color or "C0", alpha=alpha,
                        shade=True, label=label)

    ax.set_xlabel("x₁")
    ax.set_ylabel("x₂")
    ax.set_zlabel("x₃")

    if show:
        plt.show()
    return ax


def _set_3d_equal_aspect(ax):
    """Set equal aspect ratio for a 3D axes."""
    limits = np.array([
        ax.get_xlim3d(),
        ax.get_ylim3d(),
        ax.get_zlim3d(),
    ])
    center = limits.mean(axis=1)
    radius = 0.5 * (limits[:, 1] - limits[:, 0]).max()
    ax.set_xlim3d([center[0] - radius, center[0] + radius])
    ax.set_ylim3d([center[1] - radius, center[1] + radius])
    ax.set_zlim3d([center[2] - radius, center[2] + radius])
