"""Tests for closed-form BCC grid formulas (B_n, C_n, D_n).

Validates the formulas against numerical computation and verifies
Langlands duality (B_n = C_n in BCC), null eigenvalue properties,
and root counts.
"""

import numpy as np
import pytest

import sys
sys.path.insert(0, "src")

from watkins_nn.cascade import (
    _build_roots,
    _build_roots_extended,
    bn_bcc_formula,
    cn_bcc_formula,
    dn_bcc_formula,
    langlands_dual_check,
)


# ---------------------------------------------------------------------------
# Helper: compute BCC grid numerically from roots
# ---------------------------------------------------------------------------

def _numerical_grid(name: str) -> np.ndarray:
    """Compute the 3x3 BCC sign-pattern grid numerically."""
    if name.startswith("B") or name.startswith("C"):
        roots = _build_roots_extended(name)
    else:
        roots = _build_roots(name)
    g = 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)
        g[sp0, sp1] += 1
    return g


# ---------------------------------------------------------------------------
# B_n formula matches numerical computation (n = 3..12)
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("n", range(3, 13))
def test_bn_formula_matches_numerical(n):
    expected = _numerical_grid(f"B{n}")
    result = bn_bcc_formula(n)
    np.testing.assert_array_equal(result, expected)


# ---------------------------------------------------------------------------
# C_n formula matches numerical computation (n = 3..12)
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("n", range(3, 13))
def test_cn_formula_matches_numerical(n):
    expected = _numerical_grid(f"C{n}")
    result = cn_bcc_formula(n)
    np.testing.assert_array_equal(result, expected)


# ---------------------------------------------------------------------------
# D_n formula matches numerical computation (n = 4..12)
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("n", range(4, 13))
def test_dn_formula_matches_numerical(n):
    expected = _numerical_grid(f"D{n}")
    result = dn_bcc_formula(n)
    np.testing.assert_array_equal(result, expected)


# ---------------------------------------------------------------------------
# Langlands duality: B_n and C_n BCC grids are identical (n = 3..20)
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("n", range(3, 21))
def test_langlands_duality(n):
    assert langlands_dual_check(n)
    np.testing.assert_array_equal(bn_bcc_formula(n), cn_bcc_formula(n))


# ---------------------------------------------------------------------------
# Null eigenvalue tests: (1, 0, -1) is eigenvector with eigenvalue 0
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("n", range(3, 11))
def test_bn_null_eigenvalue(n):
    g = bn_bcc_formula(n)
    v = np.array([1, 0, -1])
    np.testing.assert_array_equal(g @ v, np.zeros(3, dtype=int))


@pytest.mark.parametrize("n", range(3, 11))
def test_cn_null_eigenvalue(n):
    g = cn_bcc_formula(n)
    v = np.array([1, 0, -1])
    np.testing.assert_array_equal(g @ v, np.zeros(3, dtype=int))


@pytest.mark.parametrize("n", range(4, 11))
def test_dn_null_eigenvalue(n):
    g = dn_bcc_formula(n)
    v = np.array([1, 0, -1])
    np.testing.assert_array_equal(g @ v, np.zeros(3, dtype=int))


# ---------------------------------------------------------------------------
# Root count tests
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("n", range(3, 11))
def test_bn_root_count(n):
    """B_n has 2n^2 roots."""
    roots = _build_roots_extended(f"B{n}")
    assert len(roots) == 2 * n * n


@pytest.mark.parametrize("n", range(3, 11))
def test_cn_root_count(n):
    """C_n has 2n^2 roots."""
    roots = _build_roots_extended(f"C{n}")
    assert len(roots) == 2 * n * n


@pytest.mark.parametrize("n", range(4, 11))
def test_dn_root_count(n):
    """D_n has 2n(n-1) roots."""
    roots = _build_roots(f"D{n}")
    assert len(roots) == 2 * n * (n - 1)


# ---------------------------------------------------------------------------
# Grid sum equals root count
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("n", range(3, 11))
def test_bn_grid_sum(n):
    g = bn_bcc_formula(n)
    assert g.sum() == 2 * n * n


@pytest.mark.parametrize("n", range(3, 11))
def test_cn_grid_sum(n):
    g = cn_bcc_formula(n)
    assert g.sum() == 2 * n * n


@pytest.mark.parametrize("n", range(4, 11))
def test_dn_grid_sum(n):
    g = dn_bcc_formula(n)
    assert g.sum() == 2 * n * (n - 1)


# ---------------------------------------------------------------------------
# Bisymmetry: all grids are bisymmetric (centrosymmetric)
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("n", range(3, 11))
def test_bn_bisymmetric(n):
    g = bn_bcc_formula(n)
    np.testing.assert_array_equal(g, g[::-1, ::-1])
    np.testing.assert_array_equal(g, g.T)


@pytest.mark.parametrize("n", range(4, 11))
def test_dn_bisymmetric(n):
    g = dn_bcc_formula(n)
    np.testing.assert_array_equal(g, g[::-1, ::-1])
    np.testing.assert_array_equal(g, g.T)


# ---------------------------------------------------------------------------
# Edge cases / validation
# ---------------------------------------------------------------------------

def test_bn_rejects_small_n():
    with pytest.raises(ValueError):
        bn_bcc_formula(2)


def test_cn_rejects_small_n():
    with pytest.raises(ValueError):
        cn_bcc_formula(2)


def test_dn_rejects_small_n():
    with pytest.raises(ValueError):
        dn_bcc_formula(3)
