"""Tests for Theorem X — The Exceptional Cascade."""

import math
import numpy as np
import pytest

from watkins_nn.cascade import (
    cascade_chain,
    bcc_grid,
    bcc_eigenvalues,
    bcc_perron_alpha,
    e8_roots,
    CASCADE_ORBITS,
    MINUSCULE_DIMS,
    E8_BCC_EIGENVALUE_1,
    E8_BCC_EIGENVALUE_2,
)


# ── Root system sizes ────────────────────────────────────────────────

class TestRootSystems:
    def test_e8_root_count(self):
        assert len(e8_roots()) == 240

    @pytest.mark.parametrize("name,expected", [
        ("E8", 240), ("E7", 126), ("E6", 72),
        ("D8", 112), ("D5", 40), ("D4", 24), ("D3", 12),
        ("F4", 48), ("A2", 6), ("A1", 2),
    ])
    def test_root_counts(self, name, expected):
        g = bcc_grid(name, (0, 1))
        assert g.sum() == expected


# ── Cascade layer sizes ──────────────────────────────────────────────

class TestCascadeLayers:
    @pytest.fixture(scope="class")
    def chain(self):
        return cascade_chain()

    def test_seven_layers(self, chain):
        assert len(chain) == 7

    @pytest.mark.parametrize("idx,expected_size", [
        (0, 114), (1, 54), (2, 32), (3, 16), (4, 12), (5, 6), (6, 4),
    ])
    def test_layer_sizes(self, chain, idx, expected_size):
        assert chain[idx]["layer_size"] == expected_size

    def test_layers_sum_to_240(self, chain):
        total = sum(c["layer_size"] for c in chain)
        # 240 = E8 roots, minus A1 (2 roots) remaining
        assert total == 240 - 2

    def test_minuscule_dimensions(self, chain):
        """Each layer = 2 × minuscule dim (±1 for weight-zero roots)."""
        for c in chain:
            step = c["step"]
            if step not in MINUSCULE_DIMS:
                continue
            half = c["layer_size"] // 2
            fund = MINUSCULE_DIMS[step]
            assert abs(half - fund) <= 1, f"{step}: {half} vs {fund}"


# ── BCC orbit structure ──────────────────────────────────────────────

class TestBCCOrbits:
    @pytest.fixture(scope="class")
    def chain(self):
        return cascade_chain()

    @pytest.mark.parametrize("step", list(CASCADE_ORBITS.keys()))
    def test_orbit_multiplicities(self, chain, step):
        layer = next(c for c in chain if c["step"] == step)
        co, ed, fa, ce = (layer["bcc_corner"], layer["bcc_edge"],
                          layer["bcc_face"], layer["bcc_center"])
        expected = CASCADE_ORBITS[step]
        assert (co, ed, fa, ce) == expected, f"{step}: got {(co,ed,fa,ce)}"

    @pytest.mark.parametrize("step", [
        "E8→E7", "E7→E6", "E6→D5", "D5→D4", "D4→D3",
    ])
    def test_orbit_sum(self, chain, step):
        """Orbit sum = layer size for layers with full 3D projection."""
        layer = next(c for c in chain if c["step"] == step)
        co, ed, fa, ce = (layer["bcc_corner"], layer["bcc_edge"],
                          layer["bcc_face"], layer["bcc_center"])
        assert 8 * co + 12 * ed + 6 * fa + ce == layer["layer_size"]


# ── Three structural laws ────────────────────────────────────────────

class TestStructuralLaws:
    @pytest.fixture(scope="class")
    def chain(self):
        return cascade_chain()

    def test_spinor_corner_law(self, chain):
        """Only E-type layers have nonzero corners; 16 = 8 + 4 + 4."""
        e_corners = [c["bcc_corner"] for c in chain if c["above"].startswith("E")]
        d_corners = [c["bcc_corner"] for c in chain if c["above"].startswith("D")]
        assert e_corners == [8, 4, 4]
        assert all(c == 0 for c in d_corners)
        assert sum(e_corners) == 16

    def test_edge_saturation_law(self, chain):
        """All 12 edge roots live in D₃→A₂."""
        for c in chain:
            if c["step"] == "D3→A2":
                assert c["bcc_edge"] == 1  # 12 cells × 1 = 12
            else:
                assert c["bcc_edge"] == 0, f"{c['step']} has edge={c['bcc_edge']}"

    def test_null_vector_cascade(self, chain):
        """(1,0,−1) is null for every layer except D₃→A₂."""
        for c in chain:
            if c["step"] == "D3→A2":
                assert not c["null_eigenvalue"], "D3→A2 should be anti-null"
            elif c["layer_size"] > 0:
                assert c["null_eigenvalue"], f"{c['step']} should be null"


# ── Branching identity ───────────────────────────────────────────────

class TestBranchingIdentity:
    def test_e8_trace_equals_e7_roots(self):
        g = bcc_grid("E8", (0, 1))
        assert np.trace(g) == 126  # |E7|

    def test_e8_offdiag_equals_complement(self):
        g = bcc_grid("E8", (0, 1))
        assert g.sum() - np.trace(g) == 114  # |E8\E7|


# ── Eigenvalue structure ─────────────────────────────────────────────

class TestBCCEigenvalues:
    def test_e6_integer_eigenvalues(self):
        evals = bcc_eigenvalues("E6", (0, 1))
        assert pytest.approx(evals[0], abs=1e-8) == 24
        assert pytest.approx(evals[1], abs=1e-8) == 6
        assert pytest.approx(evals[2], abs=1e-8) == 0

    def test_e7_integer_eigenvalues(self):
        evals = bcc_eigenvalues("E7", (0, 1))
        assert pytest.approx(evals[0], abs=1e-8) == 42
        assert pytest.approx(evals[1], abs=1e-8) == 18
        assert pytest.approx(evals[2], abs=1e-8) == 0

    def test_e8_irrational_eigenvalues(self):
        evals = bcc_eigenvalues("E8", (0, 1))
        assert pytest.approx(evals[0], rel=1e-8) == E8_BCC_EIGENVALUE_1
        assert pytest.approx(evals[1], rel=1e-8) == E8_BCC_EIGENVALUE_2
        assert pytest.approx(evals[2], abs=1e-8) == 0

    def test_sqrt33_origin(self):
        """√33 = √(1 + 2⁵) = √(1 + dim S⁺(SO(10)))."""
        _, disc = bcc_perron_alpha("E8", (0, 1))
        assert pytest.approx(disc) == 33
        assert 33 == 1 + 32  # 32 = 2⁵ = dim(S⁺ SO(10))


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

class TestDemocraticEigenvector:
    def test_e6_democratic(self):
        alpha, disc = bcc_perron_alpha("E6", (0, 1))
        assert pytest.approx(alpha, abs=1e-8) == 1.0

    def test_e7_democratic(self):
        alpha, disc = bcc_perron_alpha("E7", (0, 1))
        assert pytest.approx(alpha, abs=1e-8) == 1.0

    def test_e8_breaks_democracy(self):
        alpha, _ = bcc_perron_alpha("E8", (0, 1))
        expected = (1 + math.sqrt(33)) / 8
        assert pytest.approx(alpha, rel=1e-8) == expected
        assert alpha != pytest.approx(1.0, abs=0.01)

    def test_b_over_delta_ratio(self):
        """E₆,E₇ have b/δ=1; E₈ has b/δ=2."""
        for name, expected in [("E6", 1.0), ("E7", 1.0), ("E8", 2.0)]:
            g = bcc_grid(name, (0, 1))
            a, b, c = int(g[0, 0]), int(g[0, 1]), int(g[1, 1])
            delta = 2 * a - c
            assert b / delta == pytest.approx(expected)


# ── F₄ structure ─────────────────────────────────────────────────────

class TestF4:
    def test_f4_root_count(self):
        g = bcc_grid("F4", (0, 1))
        assert g.sum() == 48

    def test_f4_uniform_offdiag(self):
        g = bcc_grid("F4", (0, 1))
        assert g[0, 0] == g[0, 1] == g[0, 2] == 5
        assert g[1, 1] == 8

    def test_f4_null_vector(self):
        g = bcc_grid("F4", (0, 1))
        v = np.array([1, 0, -1])
        assert np.allclose(g @ v, 0)

    def test_f4_all_projections_identical(self):
        grids = set()
        for c in [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]:
            g = bcc_grid("F4", c)
            grids.add(tuple(g.flatten()))
        assert len(grids) == 1

    def test_26_center_is_dim_fund_f4(self):
        """The E₈\\E₇ center multiplicity = dim(fund F₄) = 26."""
        chain = cascade_chain()
        e8_e7 = chain[0]
        assert e8_e7["bcc_center"] == 26  # = dim(fund F₄)


# ── Projection universality ──────────────────────────────────────────

class TestProjectionUniversality:
    @pytest.mark.parametrize("name", ["E8", "E7"])
    def test_all_projections_identical(self, name):
        """Coordinate-pair projections give the same grid for E8/E7."""
        base = tuple(bcc_grid(name, (0, 1)).flatten())
        for coords in [(0, 2), (1, 3), (2, 5), (3, 4), (4, 5)]:
            g = bcc_grid(name, coords)
            assert tuple(g.flatten()) == base

    def test_e6_projections_within_subspace(self):
        """E6 projections identical for coords within {0,..,4}."""
        base = tuple(bcc_grid("E6", (0, 1)).flatten())
        for coords in [(0, 2), (1, 3), (2, 3), (2, 4)]:
            g = bcc_grid("E6", coords)
            assert tuple(g.flatten()) == base
