"""Phase 4B tests — Reach tube plotting.

Tests plot_ea_tube, plot_ia_tube, and plot_reach using the Agg backend.
3D tests (2D systems) are skipped if mpl_toolkits.mplot3d is unavailable.
"""

import matplotlib
matplotlib.use("Agg")

import numpy as np
import pytest
import matplotlib.pyplot as plt

from ellipsoidal_toolbox.core.ellipsoid import Ellipsoid
from ellipsoidal_toolbox.systems.linsys import LinSys
from ellipsoidal_toolbox.systems.reach import (
    reach_ea, reach_ia, reach_ea_d, reach_ia_d, ReachSet,
)
from ellipsoidal_toolbox.visualization.plot_reach import (
    plot_ea_tube, plot_ia_tube, plot_reach,
)

try:
    import mpl_toolkits.mplot3d  # noqa: F401
    _HAS_3D = True
except ImportError:
    _HAS_3D = False

_skip_3d = pytest.mark.skipif(not _HAS_3D, reason="mpl_toolkits.mplot3d unavailable")


# ── Fixtures ─────────────────────────────────────────────────────────

def _make_1d_ea() -> ReachSet:
    A = np.array([[0.0]])
    B = np.array([[1.0]])
    U = Ellipsoid([0.0], np.array([[1.0]]))
    sys = LinSys(A, B, U)
    X0 = Ellipsoid([0.0], np.array([[0.1]]))
    L0 = np.array([[1.0]])
    return reach_ea(sys, X0, L0, (0, 2), N_steps=20)


def _make_1d_ia() -> ReachSet:
    A = np.array([[0.0]])
    B = np.array([[1.0]])
    U = Ellipsoid([0.0], np.array([[1.0]]))
    sys = LinSys(A, B, U)
    X0 = Ellipsoid([0.0], np.array([[0.1]]))
    L0 = np.array([[1.0]])
    return reach_ia(sys, X0, L0, (0, 2), N_steps=20)


def _make_2d_ea() -> ReachSet:
    A = -np.eye(2)
    B = np.eye(2)
    U = Ellipsoid(np.zeros(2), np.eye(2))
    sys = LinSys(A, B, U)
    X0 = Ellipsoid(np.zeros(2), np.eye(2))
    L0 = np.array([[1, 0], [0, 1]], dtype=float)
    return reach_ea(sys, X0, L0, (0, 2), N_steps=15)


def _make_2d_ia() -> ReachSet:
    A = -np.eye(2)
    B = np.eye(2)
    U = Ellipsoid(np.zeros(2), np.eye(2))
    sys = LinSys(A, B, U)
    X0 = Ellipsoid(np.zeros(2), np.eye(2))
    L0 = np.array([[1, 0], [0, 1]], dtype=float)
    return reach_ia(sys, X0, L0, (0, 2), N_steps=15)


def _make_discrete_ea() -> ReachSet:
    A = 0.9 * np.eye(2)
    B = np.eye(2)
    U = Ellipsoid(np.zeros(2), np.eye(2))
    sys = LinSys(A, B, U, discrete=True)
    X0 = Ellipsoid(np.zeros(2), np.eye(2))
    L0 = np.eye(2)
    return reach_ea_d(sys, X0, L0, N_steps=8)


# =====================================================================
# 1D tube plots
# =====================================================================

class TestTube1D:
    def test_ea_tube_1d(self):
        rs = _make_1d_ea()
        ax = plot_ea_tube(rs)
        assert ax is not None
        plt.close("all")

    def test_ia_tube_1d(self):
        rs = _make_1d_ia()
        ax = plot_ia_tube(rs)
        assert ax is not None
        plt.close("all")

    def test_ea_tube_1d_no_center(self):
        rs = _make_1d_ea()
        ax = plot_ea_tube(rs, center=False, color="red")
        assert ax is not None
        plt.close("all")

    def test_ea_tube_1d_with_title(self):
        rs = _make_1d_ea()
        ax = plot_ea_tube(rs, title="1D Reach Tube")
        assert ax.get_title() == "1D Reach Tube"
        plt.close("all")

    def test_ea_tube_1d_on_existing_axes(self):
        fig, ax = plt.subplots()
        rs = _make_1d_ea()
        ax2 = plot_ea_tube(rs, ax=ax)
        assert ax2 is ax
        plt.close("all")


# =====================================================================
# 2D tube plots (require 3D projection)
# =====================================================================

@_skip_3d
class TestTube2D:
    def test_ea_tube_2d(self):
        rs = _make_2d_ea()
        ax = plot_ea_tube(rs, n_pts=30)
        assert ax is not None
        plt.close("all")

    def test_ia_tube_2d(self):
        rs = _make_2d_ia()
        ax = plot_ia_tube(rs, n_pts=30)
        assert ax is not None
        plt.close("all")

    def test_ea_tube_2d_with_options(self):
        rs = _make_2d_ea()
        ax = plot_ea_tube(rs, color="blue", alpha=0.2, center=True,
                          center_color="red", title="EA Tube", n_pts=20)
        assert ax is not None
        plt.close("all")

    def test_ea_tube_discrete(self):
        rs = _make_discrete_ea()
        ax = plot_ea_tube(rs, n_pts=20)
        assert ax is not None
        plt.close("all")

    def test_ea_tube_on_existing_axes(self):
        import mpl_toolkits.mplot3d  # noqa: F401
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
        rs = _make_2d_ea()
        ax2 = plot_ea_tube(rs, ax=ax, n_pts=20)
        assert ax2 is ax
        plt.close("all")


# =====================================================================
# plot_reach (combined EA + IA)
# =====================================================================

class TestPlotReach:
    def test_1d_ea_only(self):
        rs = _make_1d_ea()
        ax = plot_reach(rs)
        assert ax is not None
        plt.close("all")

    def test_1d_ia_only(self):
        rs = _make_1d_ia()
        ax = plot_reach(rs)
        assert ax is not None
        plt.close("all")

    @_skip_3d
    def test_2d_ea_only(self):
        rs = _make_2d_ea()
        ax = plot_reach(rs, n_pts=20)
        assert ax is not None
        plt.close("all")

    @_skip_3d
    def test_2d_ia_only(self):
        rs = _make_2d_ia()
        ax = plot_reach(rs, n_pts=20)
        assert ax is not None
        plt.close("all")

    def test_no_data_raises(self):
        rs = ReachSet(t_eval=np.array([0.0]), centers=np.array([[0.0]]),
                       n=1, n_dirs=0)
        with pytest.raises(ValueError, match="no EA or IA"):
            plot_reach(rs)


# =====================================================================
# Error cases
# =====================================================================

class TestErrors:
    def test_ea_tube_no_data(self):
        rs = _make_1d_ia()  # only IA
        with pytest.raises(ValueError, match="no EA"):
            plot_ea_tube(rs)

    def test_ia_tube_no_data(self):
        rs = _make_1d_ea()  # only EA
        with pytest.raises(ValueError, match="no IA"):
            plot_ia_tube(rs)

    def test_high_dim_raises(self):
        A = -np.eye(3)
        B = np.eye(3)
        U = Ellipsoid(np.zeros(3), np.eye(3))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(3), np.eye(3))
        L0 = np.eye(3)
        rs = reach_ea(sys, X0, L0, (0, 1), N_steps=5)
        with pytest.raises(ValueError, match="projection"):
            plot_ea_tube(rs)

    @_skip_3d
    def test_high_dim_projected_works(self):
        A = -np.eye(3)
        B = np.eye(3)
        U = Ellipsoid(np.zeros(3), np.eye(3))
        sys = LinSys(A, B, U)
        X0 = Ellipsoid(np.zeros(3), np.eye(3))
        L0 = np.eye(3)
        rs = reach_ea(sys, X0, L0, (0, 1), N_steps=5)
        basis = np.array([[1, 0], [0, 1], [0, 0]], dtype=float)
        rs_proj = rs.projection(basis)
        ax = plot_ea_tube(rs_proj, n_pts=15)
        assert ax is not None
        plt.close("all")
