"""Phase 4A tests — Ellipsoid plotting.

Uses the non-interactive Agg backend so no display is needed.
Tests verify that plotting functions run without error, return
the correct axes types, and produce expected graphical elements.
"""

import matplotlib
matplotlib.use("Agg")

import numpy as np
import pytest
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon

from ellipsoidal_toolbox.core.ellipsoid import Ellipsoid
from ellipsoidal_toolbox.visualization.plot_ellipsoid import (
    plot_ellipsoid,
    plot_ellipsoids,
    _boundary_1d,
    _boundary_2d,
    _boundary_3d,
)

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")


# =====================================================================
# Boundary generators
# =====================================================================

class TestBoundary1D:
    def test_unit_interval(self):
        E = Ellipsoid([0.0], np.array([[1.0]]))
        lo, hi = _boundary_1d(E)
        assert lo == pytest.approx(-1.0)
        assert hi == pytest.approx(1.0)

    def test_shifted(self):
        E = Ellipsoid([3.0], np.array([[4.0]]))
        lo, hi = _boundary_1d(E)
        assert lo == pytest.approx(1.0)
        assert hi == pytest.approx(5.0)


class TestBoundary2D:
    def test_unit_circle_shape(self):
        E = Ellipsoid(np.zeros(2), np.eye(2))
        pts = _boundary_2d(E, n_pts=100)
        assert pts.shape == (100, 2)

    def test_unit_circle_on_circle(self):
        E = Ellipsoid(np.zeros(2), np.eye(2))
        pts = _boundary_2d(E, n_pts=200)
        radii = np.linalg.norm(pts, axis=1)
        np.testing.assert_allclose(radii, 1.0, atol=1e-6)

    def test_shifted_center(self):
        E = Ellipsoid([5.0, 3.0], np.eye(2))
        pts = _boundary_2d(E, n_pts=100)
        center = pts.mean(axis=0)
        np.testing.assert_allclose(center, [5.0, 3.0], atol=0.1)

    def test_scaled_ellipse(self):
        Q = np.diag([4.0, 1.0])
        E = Ellipsoid(np.zeros(2), Q)
        pts = _boundary_2d(E, n_pts=200)
        assert np.max(np.abs(pts[:, 0])) > 1.8
        assert np.max(np.abs(pts[:, 1])) < 1.2


class TestBoundary3D:
    def test_unit_sphere_shape(self):
        E = Ellipsoid(np.zeros(3), np.eye(3))
        X, Y, Z = _boundary_3d(E, n_pts=10)
        assert X.shape == (10, 10)
        assert Y.shape == (10, 10)
        assert Z.shape == (10, 10)

    def test_unit_sphere_on_sphere(self):
        E = Ellipsoid(np.zeros(3), np.eye(3))
        X, Y, Z = _boundary_3d(E, n_pts=20)
        R = np.sqrt(X**2 + Y**2 + Z**2)
        np.testing.assert_allclose(R, 1.0, atol=1e-6)

    def test_shifted_center(self):
        E = Ellipsoid([1.0, 2.0, 3.0], np.eye(3))
        X, Y, Z = _boundary_3d(E, n_pts=10)
        assert X.mean() == pytest.approx(1.0, abs=0.3)
        assert Y.mean() == pytest.approx(2.0, abs=0.3)
        assert Z.mean() == pytest.approx(3.0, abs=0.3)


# =====================================================================
# plot_ellipsoid — 1D
# =====================================================================

class TestPlot1D:
    def test_returns_axes(self):
        E = Ellipsoid([0.0], np.array([[1.0]]))
        ax = plot_ellipsoid(E)
        assert ax is not None
        plt.close("all")

    def test_with_options(self):
        E = Ellipsoid([0.0], np.array([[1.0]]))
        ax = plot_ellipsoid(E, color="red", label="test", linewidth=3)
        assert ax is not None
        plt.close("all")

    def test_on_existing_axes(self):
        fig, ax = plt.subplots()
        E = Ellipsoid([0.0], np.array([[1.0]]))
        ax2 = plot_ellipsoid(E, ax=ax)
        assert ax2 is ax
        plt.close("all")


# =====================================================================
# plot_ellipsoid — 2D
# =====================================================================

class TestPlot2D:
    def test_returns_axes(self):
        E = Ellipsoid(np.zeros(2), np.eye(2))
        ax = plot_ellipsoid(E)
        assert ax is not None
        plt.close("all")

    def test_filled(self):
        E = Ellipsoid(np.zeros(2), np.eye(2))
        ax = plot_ellipsoid(E, fill=True, color="blue", alpha=0.5)
        patches = [p for p in ax.patches if isinstance(p, Polygon)]
        assert len(patches) >= 1
        plt.close("all")

    def test_unfilled(self):
        E = Ellipsoid(np.zeros(2), np.eye(2))
        ax = plot_ellipsoid(E, fill=False, color="green")
        assert len(ax.lines) >= 1
        plt.close("all")

    def test_with_label(self):
        E = Ellipsoid(np.zeros(2), np.eye(2))
        ax = plot_ellipsoid(E, label="E1", fill=True)
        plt.close("all")

    def test_custom_n_pts(self):
        E = Ellipsoid(np.zeros(2), np.eye(2))
        ax = plot_ellipsoid(E, n_pts=50)
        assert ax is not None
        plt.close("all")

    def test_on_existing_axes(self):
        fig, ax = plt.subplots()
        E = Ellipsoid(np.zeros(2), np.eye(2))
        ax2 = plot_ellipsoid(E, ax=ax)
        assert ax2 is ax
        plt.close("all")

    def test_degenerate_ellipsoid(self):
        Q = np.array([[1, 0], [0, 0]])
        E = Ellipsoid(np.zeros(2), Q)
        ax = plot_ellipsoid(E, fill=False)
        assert ax is not None
        plt.close("all")


# =====================================================================
# plot_ellipsoid — 3D
# =====================================================================

@_skip_3d
class TestPlot3D:
    def test_returns_axes(self):
        E = Ellipsoid(np.zeros(3), np.eye(3))
        ax = plot_ellipsoid(E)
        assert ax is not None
        plt.close("all")

    def test_wireframe(self):
        E = Ellipsoid(np.zeros(3), np.eye(3))
        ax = plot_ellipsoid(E, wireframe=True)
        assert ax is not None
        plt.close("all")

    def test_surface(self):
        E = Ellipsoid(np.zeros(3), np.eye(3))
        ax = plot_ellipsoid(E, fill=True, alpha=0.4, color="red")
        assert ax is not None
        plt.close("all")

    def test_on_existing_axes(self):
        import mpl_toolkits.mplot3d  # noqa: F401
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
        E = Ellipsoid(np.zeros(3), np.eye(3))
        ax2 = plot_ellipsoid(E, ax=ax, n_pts=8)
        assert ax2 is ax
        plt.close("all")


# =====================================================================
# plot_ellipsoid — error cases
# =====================================================================

class TestPlotErrors:
    def test_empty_ellipsoid_raises(self):
        with pytest.raises(ValueError, match="empty"):
            plot_ellipsoid(Ellipsoid())

    def test_4d_raises(self):
        E = Ellipsoid(np.zeros(4), np.eye(4))
        with pytest.raises(ValueError, match="4D"):
            plot_ellipsoid(E)


# =====================================================================
# plot_ellipsoids — multiple
# =====================================================================

class TestPlotMultiple:
    def test_two_2d_ellipsoids(self):
        E1 = Ellipsoid(np.zeros(2), np.eye(2))
        E2 = Ellipsoid([2.0, 0.0], 0.5 * np.eye(2))
        ax = plot_ellipsoids([E1, E2])
        assert ax is not None
        plt.close("all")

    def test_with_colors_and_labels(self):
        E1 = Ellipsoid(np.zeros(2), np.eye(2))
        E2 = Ellipsoid([2.0, 0.0], 0.5 * np.eye(2))
        ax = plot_ellipsoids([E1, E2],
                             colors=["red", "blue"],
                             labels=["E1", "E2"],
                             title="Two Ellipsoids")
        plt.close("all")

    @_skip_3d
    def test_three_3d_ellipsoids(self):
        ells = [
            Ellipsoid(np.zeros(3), np.eye(3)),
            Ellipsoid([2, 0, 0], 0.5 * np.eye(3)),
            Ellipsoid([0, 2, 0], np.diag([1, 2, 1])),
        ]
        ax = plot_ellipsoids(ells, n_pts=8)
        assert ax is not None
        plt.close("all")

    def test_mixed_dimensions_raises(self):
        E2 = Ellipsoid(np.zeros(2), np.eye(2))
        E3 = Ellipsoid(np.zeros(3), np.eye(3))
        with pytest.raises(ValueError, match="same dimension"):
            plot_ellipsoids([E2, E3])

    def test_empty_list_raises(self):
        with pytest.raises(ValueError, match="No ellipsoids"):
            plot_ellipsoids([])

    def test_on_existing_axes(self):
        fig, ax = plt.subplots()
        ells = [
            Ellipsoid(np.zeros(2), np.eye(2)),
            Ellipsoid([1, 0], 0.5 * np.eye(2)),
        ]
        ax2 = plot_ellipsoids(ells, ax=ax)
        assert ax2 is ax
        plt.close("all")

    def test_1d_multiple(self):
        ells = [
            Ellipsoid([0.0], np.array([[1.0]])),
            Ellipsoid([3.0], np.array([[0.5]])),
        ]
        ax = plot_ellipsoids(ells, colors=["red", "blue"])
        assert ax is not None
        plt.close("all")
