"""Tests for exceptional BCC grid formulas.

Verifies the symbolic closed-form BCC grids for E₆, E₇, E₈, F₄, G₂
against numerical computation from root coordinates.

Key results:
    E_n: a = 1 + 2^(n-3), b = 2f, c = 4·C(f,2) + correction
    F₄:  [[5,5,5],[5,8,5],[5,5,5]]  (bisymmetric, nearly uniform)
    G₂:  [[1,1,3],[1,0,1],[3,1,1]]  (NOT bisymmetric — unique!)
"""

import math
import numpy as np
import pytest

from watkins_nn.cascade import (
    bcc_grid,
    bcc_eigenvalues,
    en_bcc_formula,
    f4_bcc_formula,
    g2_bcc_formula,
    exceptional_bcc_grid,
    democratic_deficit,
    EXCEPTIONAL_BCC_GRIDS,
    EXCEPTIONAL_BCC_EIGENVALUES,
    EXCEPTIONAL_INVARIANTS,
    G2_BCC_GRID,
    _g2_roots,
    _build_roots_extended,
)


# ── Formula correctness: match numerical computation ────────────────


class TestFormulasMatchNumerical:
    """Symbolic formulas must reproduce the numerical BCC grids exactly."""

    @pytest.mark.parametrize("n", [6, 7, 8])
    def test_en_formula_matches_numerical(self, n):
        formula = en_bcc_formula(n)
        numerical = bcc_grid(f"E{n}", (0, 1))
        np.testing.assert_array_equal(formula, numerical)

    def test_f4_formula_matches_numerical(self):
        formula = f4_bcc_formula()
        numerical = bcc_grid("F4", (0, 1))
        np.testing.assert_array_equal(formula, numerical)

    def test_g2_formula_matches_numerical(self):
        formula = g2_bcc_formula()
        # Compute G₂ grid numerically
        roots = _g2_roots()
        grid = np.zeros((3, 3), dtype=int)
        for r in roots:
            sp0 = 0 if r[0] > 1e-10 else (2 if r[0] < -1e-10 else 1)
            sp1 = 0 if r[1] > 1e-10 else (2 if r[1] < -1e-10 else 1)
            grid[sp0, sp1] += 1
        np.testing.assert_array_equal(formula, grid)

    @pytest.mark.parametrize("name", ["E6", "E7", "E8", "F4", "G2"])
    def test_exceptional_bcc_grid_dispatcher(self, name):
        grid = exceptional_bcc_grid(name)
        if name == "G2":
            expected = g2_bcc_formula()
        elif name == "F4":
            expected = f4_bcc_formula()
        else:
            expected = en_bcc_formula(int(name[1]))
        np.testing.assert_array_equal(grid, expected)


# ── Root count verification ─────────────────────────────────────────


class TestRootCounts:
    """Grid entries must sum to |Phi|."""

    @pytest.mark.parametrize("name,expected_roots", [
        ("E6", 72), ("E7", 126), ("E8", 240), ("F4", 48), ("G2", 12),
    ])
    def test_grid_sum_equals_root_count(self, name, expected_roots):
        grid = exceptional_bcc_grid(name)
        assert grid.sum() == expected_roots
        assert grid.sum() == EXCEPTIONAL_INVARIANTS[name]["roots"]


# ── Bisymmetry and structure ────────────────────────────────────────


class TestBisymmetry:
    """E₆, E₇, E₈, F₄ are bisymmetric; G₂ is not."""

    @pytest.mark.parametrize("name", ["E6", "E7", "E8", "F4"])
    def test_bisymmetric(self, name):
        g = exceptional_bcc_grid(name)
        # Corners equal
        assert g[0, 0] == g[0, 2] == g[2, 0] == g[2, 2]
        # Edges equal
        assert g[0, 1] == g[1, 0] == g[1, 2] == g[2, 1]
        # Center is unique
        assert g[1, 1] == EXCEPTIONAL_BCC_GRIDS[name][2]

    def test_g2_not_bisymmetric(self):
        g = g2_bcc_formula()
        assert g[0, 0] != g[0, 2], "G₂ must NOT be bisymmetric"
        assert g[0, 0] == 1
        assert g[0, 2] == 3


# ── Spinor decomposition of corner entry a ──────────────────────────


class TestSpinorCornerFormula:
    """a = 1 + 2^(n-3) for E_n, a = 1 + 2^(rank-2) for F₄."""

    @pytest.mark.parametrize("n,expected_a", [(6, 9), (7, 17), (8, 33)])
    def test_en_corner_is_1_plus_spinor(self, n, expected_a):
        a = EXCEPTIONAL_BCC_GRIDS[f"E{n}"][0]
        assert a == expected_a
        spinor_count = 2 ** (n - 3)
        assert a == 1 + spinor_count

    def test_f4_corner_decomposition(self):
        a = EXCEPTIONAL_BCC_GRIDS["F4"][0]
        # 1 long root (e₀+e₁) + 4 spinor roots ((½,½,±½,±½))
        assert a == 1 + 4
        assert a == 1 + 2 ** (4 - 2)  # rank=4, no parity constraint


# ── Edge entry b ────────────────────────────────────────────────────


class TestEdgeFormula:
    """b = 2f for E-types (f = free interior coords), b = 2f+1 for F₄."""

    @pytest.mark.parametrize("name,f", [("E6", 3), ("E7", 4), ("E8", 6)])
    def test_en_edge_is_2f(self, name, f):
        b = EXCEPTIONAL_BCC_GRIDS[name][1]
        assert b == 2 * f

    def test_f4_edge(self):
        b = EXCEPTIONAL_BCC_GRIDS["F4"][1]
        # 2 free interior coords (x₂, x₃) → 4 long roots + 1 short root
        assert b == 2 * 2 + 1


# ── Eigenvalue verification ────────────────────────────────────────


class TestEigenvalues:
    @pytest.mark.parametrize("name", ["E6", "E7", "E8", "F4"])
    def test_eigenvalues_match(self, name):
        evals = bcc_eigenvalues(name, (0, 1))
        expected = EXCEPTIONAL_BCC_EIGENVALUES[name]
        for got, exp in zip(evals, expected):
            assert pytest.approx(got, abs=1e-8) == exp

    def test_g2_eigenvalues(self):
        g = g2_bcc_formula()
        evals = sorted(np.linalg.eigvals(g.astype(float)).real, reverse=True)
        expected = EXCEPTIONAL_BCC_EIGENVALUES["G2"]
        for got, exp in zip(evals, expected):
            assert pytest.approx(got, abs=1e-8) == exp

    def test_e6_e7_integer_eigenvalues(self):
        """E₆ and E₇ have all-integer eigenvalues (democratic)."""
        for name in ["E6", "E7"]:
            evals = EXCEPTIONAL_BCC_EIGENVALUES[name]
            for e in evals:
                assert e == int(e)

    def test_e8_irrational_eigenvalues(self):
        """E₈ eigenvalues involve √33."""
        e1, e2, e3 = EXCEPTIONAL_BCC_EIGENVALUES["E8"]
        assert pytest.approx(e1 + e2) == 126
        assert pytest.approx(e1 - e2) == 6 * math.sqrt(33)
        assert e3 == 0.0

    def test_f4_eigenvalues_involve_sqrt51(self):
        e1, e2, e3 = EXCEPTIONAL_BCC_EIGENVALUES["F4"]
        assert pytest.approx(e1 + e2) == 18
        assert pytest.approx(e1 - e2) == 2 * math.sqrt(51)
        assert e3 == 0.0


# ── Democratic eigenvector ──────────────────────────────────────────


class TestDemocraticDeficit:
    """(1,1,1) is an eigenvector iff democratic deficit = 0."""

    def test_e6_democratic(self):
        assert democratic_deficit("E6") == 0

    def test_e7_democratic(self):
        assert democratic_deficit("E7") == 0

    def test_e8_deficit_is_6(self):
        assert democratic_deficit("E8") == 6

    def test_f4_deficit_is_3(self):
        assert democratic_deficit("F4") == 3

    @pytest.mark.parametrize("name", ["E6", "E7"])
    def test_democratic_means_111_eigenvector(self, name):
        g = exceptional_bcc_grid(name)
        v = np.array([1, 1, 1])
        Mv = g @ v
        # Mv should be proportional to v
        assert Mv[0] == Mv[1] == Mv[2]

    @pytest.mark.parametrize("name", ["E8", "F4"])
    def test_nondemocratic_breaks_111(self, name):
        g = exceptional_bcc_grid(name)
        v = np.array([1, 1, 1])
        Mv = g @ v
        assert Mv[0] != Mv[1], f"{name} should break (1,1,1) eigenvector"


# ── Null eigenvalue ─────────────────────────────────────────────────


class TestNullEigenvalue:
    """(1,0,-1) eigenvalue = 0 for bisymmetric, = -2 for G₂."""

    @pytest.mark.parametrize("name", ["E6", "E7", "E8", "F4"])
    def test_bisymmetric_null(self, name):
        g = exceptional_bcc_grid(name)
        v = np.array([1, 0, -1])
        assert np.allclose(g @ v, 0)

    def test_g2_antinull(self):
        g = g2_bcc_formula()
        v = np.array([1, 0, -1])
        Mv = g @ v
        assert Mv[0] == -2
        assert Mv[1] == 0
        assert Mv[2] == 2


# ── Trace and branching identities ──────────────────────────────────


class TestTraceIdentities:
    def test_e8_trace_equals_e7_roots(self):
        g = exceptional_bcc_grid("E8")
        assert np.trace(g) == 126  # = |E₇|

    def test_e7_trace_equals_60(self):
        g = exceptional_bcc_grid("E7")
        # trace = 2a + c = 34 + 26 = 60
        assert np.trace(g) == 60

    def test_e6_trace_equals_30(self):
        g = exceptional_bcc_grid("E6")
        # trace = 2a + c = 18 + 12 = 30 = |E₆|/rank · (rank-2) + ...
        assert np.trace(g) == 30

    @pytest.mark.parametrize("name", ["E6", "E7", "E8", "F4"])
    def test_trace_equals_eigenvalue_sum(self, name):
        g = exceptional_bcc_grid(name)
        evals = EXCEPTIONAL_BCC_EIGENVALUES[name]
        assert pytest.approx(np.trace(g)) == sum(evals)


# ── G₂ uniqueness ──────────────────────────────────────────────────


class TestG2Uniqueness:
    def test_g2_center_is_zero(self):
        """G₂ is the only simple root system with BCC center = 0."""
        g = g2_bcc_formula()
        assert g[1, 1] == 0

    def test_g2_no_zero_eigenvalue(self):
        """G₂ is the only system with no zero BCC eigenvalue."""
        evals = EXCEPTIONAL_BCC_EIGENVALUES["G2"]
        for e in evals:
            assert abs(e) > 0.1

    def test_g2_eigenvalues_involve_sqrt6(self):
        """G₂ char poly = (λ+2)(λ²−4λ−2), roots involve √6."""
        e1, e2, e3 = EXCEPTIONAL_BCC_EIGENVALUES["G2"]
        assert pytest.approx(e1 + e2) == 4  # from λ²−4λ−2
        assert pytest.approx(e1 - e2) == 2 * math.sqrt(6)
        assert pytest.approx(e3) == -2

    def test_g2_root_count(self):
        roots = _g2_roots()
        assert len(roots) == 12

    def test_g2_short_long_decomposition(self):
        """6 short roots + 6 long roots = 12."""
        roots = _g2_roots()
        short = [r for r in roots if sum(x**2 for x in r) < 3]
        long_ = [r for r in roots if sum(x**2 for x in r) >= 3]
        assert len(short) == 6
        assert len(long_) == 6


# ── Cross-system patterns ──────────────────────────────────────────


class TestCrossSystemPatterns:
    def test_en_a_formula_power_of_2_plus_1(self):
        """a(E_n) = 1 + 2^(n-3) for n = 6, 7, 8."""
        for n in [6, 7, 8]:
            a = EXCEPTIONAL_BCC_GRIDS[f"E{n}"][0]
            assert a == 1 + 2 ** (n - 3)

    def test_en_b_values_are_even(self):
        """All E-type edge entries are even (no short roots)."""
        for name in ["E6", "E7", "E8"]:
            b = EXCEPTIONAL_BCC_GRIDS[name][1]
            assert b % 2 == 0

    def test_f4_grid_nearly_uniform(self):
        """F₄ has a = b = 5 (corners = edges), the 'most uniform' grid."""
        a, b, c = EXCEPTIONAL_BCC_GRIDS["F4"]
        assert a == b == 5

    def test_en_invalid_rank(self):
        with pytest.raises(ValueError):
            en_bcc_formula(5)
        with pytest.raises(ValueError):
            en_bcc_formula(9)
