from __future__ import annotations

import datetime
import itertools
import json
import math
from pathlib import Path
from typing import Any

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "results" / "sanswf_symmetry" / "result.json"
SPEC = ROOT / "research" / "sanswf" / "SANSWF_SYMMETRY_STABILITY_SPEC.md"
ANSWF_RESULT = ROOT / "results" / "answf_eft" / "result.json"

ABS_TOL = 2.0e-11
UNITARITY_LIMIT = 8.0 * math.pi
LAMBDA_H = 0.13
LAMBDA_A = 0.20
LAMBDA_P = 0.10
N_H = 4
N_ADJ = 24


def utc_now_z() -> str:
    return datetime.datetime.now(datetime.UTC).isoformat(timespec="seconds").replace("+00:00", "Z")


def suN_generators(n: int) -> np.ndarray:
    """Hermitian fundamental generators with Tr(T^a T^b)=1/2 delta^{ab}."""
    gs: list[np.ndarray] = []
    for i in range(n):
        for j in range(i + 1, n):
            s = np.zeros((n, n), dtype=complex)
            s[i, j] = s[j, i] = 0.5
            gs.append(s)
            a = np.zeros((n, n), dtype=complex)
            a[i, j] = -0.5j
            a[j, i] = 0.5j
            gs.append(a)
    for k in range(1, n):
        d = np.zeros(n, dtype=float)
        d[:k] = 1.0
        d[k] = -float(k)
        gs.append((np.diag(d) / math.sqrt(2.0 * k * (k + 1.0))).astype(complex))
    return np.asarray(gs)


def suN_structure_constants(gens: np.ndarray) -> np.ndarray:
    n = len(gens)
    f = np.empty((n, n, n), dtype=float)
    for a in range(n):
        for b in range(n):
            comm = gens[a] @ gens[b] - gens[b] @ gens[a]
            for c in range(n):
                f[a, b, c] = float((-2.0j * np.trace(comm @ gens[c])).real)
    return f


def adjoint_quartic_derivative_tensor(gens: np.ndarray) -> np.ndarray:
    """Fourth derivative of Tr(Sigma^4) w.r.t. real adjoint components at Sigma=0."""
    prod2 = np.einsum("aij,bjk->abik", gens, gens)
    tr4 = np.einsum("abij,cdji->abcd", prod2, prod2).real
    d4 = np.zeros_like(tr4)
    for perm in itertools.permutations(range(4)):
        d4 += np.transpose(tr4, perm)
    return d4


def full_contact_matrix(d4_adj: np.ndarray) -> tuple[np.ndarray, list[tuple[int, int]]]:
    """
    Contact-amplitude matrix from fourth derivatives of
      V = lambda_H/4 (h^2)^2 + lambda_A Tr(Sigma^4) + lambda_p/4 h^2 Sigma^2.

    The two-particle basis contains unordered component pairs (i<=j), with a
    factor sqrt(2) for identical-particle channels. This convention is stated
    explicitly so the a0=M/(16*pi) normalization is reproducible.
    """
    n = N_H + N_ADJ
    pairs = [(i, j) for i in range(n) for j in range(i, n)]
    mat = np.zeros((len(pairs), len(pairs)), dtype=float)

    for p, (i, j) in enumerate(pairs):
        norm_i = math.sqrt(1.0 + float(i == j))
        for q, (k, ell) in enumerate(pairs):
            norm_f = math.sqrt(1.0 + float(k == ell))
            inds = [i, j, k, ell]
            n_h = sum(idx < N_H for idx in inds)
            raw = 0.0
            if n_h == 4:
                a, b, c, d = inds
                raw = 2.0 * LAMBDA_H * (
                    float(a == b and c == d)
                    + float(a == c and b == d)
                    + float(a == d and b == c)
                )
            elif n_h == 2:
                hs = [idx for idx in inds if idx < N_H]
                ss = [idx - N_H for idx in inds if idx >= N_H]
                raw = LAMBDA_P * float(hs[0] == hs[1] and ss[0] == ss[1])
            elif n_h == 0:
                a, b, c, d = [idx - N_H for idx in inds]
                raw = LAMBDA_A * d4_adj[a, b, c, d]
            mat[p, q] = raw / (norm_i * norm_f)
    return mat, pairs


def singlet_matrix(d4_adj: np.ndarray) -> dict[str, Any]:
    # Exact projection in the same normalized identical-particle convention.
    a_coeff = 6.0
    c_coeff = 0.5 * math.sqrt(N_H * N_ADJ)  # 2*sqrt(6)

    total_adj = 0.0
    for a in range(N_ADJ):
        for b in range(N_ADJ):
            total_adj += d4_adj[a, a, b, b]
    b_coeff = total_adj / (2.0 * N_ADJ)

    m = np.array(
        [
            [a_coeff * LAMBDA_H, c_coeff * LAMBDA_P],
            [c_coeff * LAMBDA_P, b_coeff * LAMBDA_A],
        ],
        dtype=float,
    )
    eig = np.linalg.eigvalsh(m)
    return {
        "coefficient_A_over_lambdaH": a_coeff,
        "coefficient_B_over_lambdaA": float(b_coeff),
        "coefficient_C_over_lambdaP": c_coeff,
        "matrix": m.tolist(),
        "eigenvalues": eig.tolist(),
        "max_abs_eigenvalue": float(np.max(np.abs(eig))),
        "max_abs_a0": float(np.max(np.abs(eig)) / (16.0 * math.pi)),
    }


def adjoint_direction_ratio(x: np.ndarray) -> float:
    x = np.asarray(x, dtype=float)
    x = x - np.mean(x)
    q2 = float(np.sum(x * x))
    if q2 <= 0.0:
        return float("nan")
    return float(np.sum(x ** 4) / (q2 * q2))


def two_level_ratio(k: int, n: int = 5) -> float:
    a = 1.0
    b = -k / (n - k)
    q2 = k * a * a + (n - k) * b * b
    q4 = k * a ** 4 + (n - k) * b ** 4
    return q4 / (q2 * q2)


def bfb_probe() -> dict[str, Any]:
    ratios = [two_level_ratio(k) for k in range(1, 5)]
    r_min = 7.0 / 30.0
    r_max = 13.0 / 20.0

    # Three-distinct-root stationary patterns permitted by the Lagrange cubic.
    # Multiplicities (3,1,1) -> {0, +a, -a}: r=1/2.
    # Multiplicities (2,2,1) -> {+a,+a,-a,-a,0}: r=1/4.
    three_root_stationary = [0.5, 0.25]

    rng = np.random.default_rng(230817)
    sampled = []
    for _ in range(25000):
        v = rng.normal(size=5)
        sampled.append(adjoint_direction_ratio(v))
    sampled_arr = np.asarray(sampled)

    # Exact BFB reduction in x=H^dag H>=0 and y=Tr Sigma^2>=0:
    # V/y^2 = lambda_H t^2 + lambda_p t + lambda_A r, t=x/y>=0.
    bfb_margin = LAMBDA_P + 2.0 * math.sqrt(LAMBDA_H * LAMBDA_A * r_min)
    if LAMBDA_P >= 0.0:
        min_normalized_v = LAMBDA_A * r_min
        minimizer_t = 0.0
    else:
        minimizer_t = -LAMBDA_P / (2.0 * LAMBDA_H)
        min_normalized_v = LAMBDA_A * r_min - (LAMBDA_P ** 2) / (4.0 * LAMBDA_H)

    return {
        "r_min_exact": r_min,
        "r_max_exact": r_max,
        "two_level_stationary_ratios": ratios,
        "three_root_stationary_ratios": three_root_stationary,
        "random_stress_min": float(np.min(sampled_arr)),
        "random_stress_max": float(np.max(sampled_arr)),
        "all_random_samples_inside_exact_interval": bool(np.all(sampled_arr >= r_min - ABS_TOL) and np.all(sampled_arr <= r_max + ABS_TOL)),
        "lambda_H": LAMBDA_H,
        "lambda_A": LAMBDA_A,
        "lambda_p": LAMBDA_P,
        "corrected_mixed_BFB_margin": bfb_margin,
        "min_normalized_quartic_V_over_y2": min_normalized_v,
        "minimizer_t_x_over_y": minimizer_t,
        "pass": bool(
            LAMBDA_H > 0.0
            and LAMBDA_A > 0.0
            and bfb_margin > 0.0
            and min_normalized_v > 0.0
            and np.all(sampled_arr >= r_min - ABS_TOL)
            and np.all(sampled_arr <= r_max + ABS_TOL)
        ),
    }


def gauge_mass_probe(gens: np.ndarray, f: np.ndarray) -> dict[str, Any]:
    gram = np.einsum("aij,bji->ab", gens, gens).real
    gram_residual = float(np.max(np.abs(gram - 0.5 * np.eye(len(gens)))))

    closure_residual = 0.0
    for a in range(len(gens)):
        for b in range(len(gens)):
            lhs = gens[a] @ gens[b] - gens[b] @ gens[a]
            rhs = 1.0j * np.einsum("c,cij->ij", f[a, b], gens)
            closure_residual = max(closure_residual, float(np.max(np.abs(lhs - rhs))))

    # Canonical SU(5)->SU(3)xSU(2)xU(1) adjoint direction.
    sigma0 = np.diag([2.0, 2.0, 2.0, -3.0, -3.0]).astype(complex) / math.sqrt(60.0)
    v = np.array([2.0 * np.trace(sigma0 @ t).real for t in gens], dtype=float)
    recon = np.einsum("a,aij->ij", v, gens)
    vev_reconstruction = float(np.max(np.abs(recon - sigma0)))

    # Index-safe form: R_ad=f^{acd} v_c, M^2=g^2 R R^T. Set g=1 for structural test.
    rmat = np.einsum("acd,c->ad", f, v)
    mass2 = rmat @ rmat.T
    eig = np.linalg.eigvalsh(mass2)
    zero_count = int(np.sum(np.abs(eig) <= 1.0e-10))
    positive = eig[eig > 1.0e-10]

    return {
        "generator_gram_residual": gram_residual,
        "lie_closure_residual": closure_residual,
        "vev_reconstruction_residual": vev_reconstruction,
        "mass2_symmetry_residual": float(np.max(np.abs(mass2 - mass2.T))),
        "mass2_min_eigenvalue": float(np.min(eig)),
        "unbroken_zero_modes": zero_count,
        "broken_positive_modes": int(len(positive)),
        "broken_mass2_min": float(np.min(positive)),
        "broken_mass2_max": float(np.max(positive)),
        "pass": bool(
            gram_residual <= ABS_TOL
            and closure_residual <= ABS_TOL
            and vev_reconstruction <= ABS_TOL
            and np.min(eig) >= -ABS_TOL
            and zero_count == 12
            and len(positive) == 12
        ),
    }


def main() -> None:
    answf_status = "MISSING"
    answf_pass = False
    if ANSWF_RESULT.exists():
        obj = json.loads(ANSWF_RESULT.read_text(encoding="utf-8"))
        answf_status = str(obj.get("status", "UNKNOWN"))
        answf_pass = bool(obj.get("verdict") == "PASS" and answf_status == "PASS_ANSWF_EFT_POSITIVITY_TO_SANSWF_HANDOFF")

    gens = suN_generators(5)
    f = suN_structure_constants(gens)
    gauge = gauge_mass_probe(gens, f)
    bfb = bfb_probe()
    d4_adj = adjoint_quartic_derivative_tensor(gens)
    singlet = singlet_matrix(d4_adj)
    full_mat, pairs = full_contact_matrix(d4_adj)
    matrix_symmetry = float(np.max(np.abs(full_mat - full_mat.T)))
    full_eig = np.linalg.eigvalsh(full_mat)
    max_abs_eig = float(np.max(np.abs(full_eig)))
    max_abs_a0 = max_abs_eig / (16.0 * math.pi)

    # The largest eigenvalue must agree with the singlet block for this benchmark.
    singlet_full_agreement = abs(max_abs_eig - singlet["max_abs_eigenvalue"])

    checks = {
        "answf_handoff_passes": answf_pass,
        "su5_generators_and_closure_consistent": bool(gauge["generator_gram_residual"] <= ABS_TOL and gauge["lie_closure_residual"] <= ABS_TOL),
        "su5_breaking_mass_matrix_psd": bool(gauge["pass"]),
        "su5_breaking_has_12_unbroken_and_12_broken_generators": bool(gauge["unbroken_zero_modes"] == 12 and gauge["broken_positive_modes"] == 12),
        "corrected_su5_adjoint_BFB_passes": bool(bfb["pass"]),
        "full_contact_matrix_symmetric": bool(matrix_symmetry <= ABS_TOL),
        "full_406_channel_partial_wave_unitarity": bool(max_abs_eig < UNITARITY_LIMIT and max_abs_a0 < 0.5),
        "singlet_block_contains_global_maximum_at_benchmark": bool(singlet_full_agreement <= 2.0e-10),
    }

    claim_guards = {
        "old_a0_0p12_value_not_used": True,
        "old_lambdaeig_1p2_value_not_used": True,
        "unitarity_convention_explicitly_defined_in_gate": True,
        "bfb_uses_TrSigma4_direction_ratio_not_assumed_equal_to_TrSigma2_squared": True,
        "gauge_mass_matrix_uses_index_safe_adjoint_map_product": True,
        "su5_benchmark_vev_is_structural_normalization_not_gut_scale_prediction": True,
        "anomaly_cancellation_and_unified_threshold_matching_deferred_to_utfanswf_embedding": True,
        "historical_rows_not_renumbered": True,
    }

    verdict = "PASS" if all(checks.values()) and all(claim_guards.values()) else "FAIL"
    status = "PASS_SANSWF_SYMMETRY_STABILITY_TO_UTFANSWF_HANDOFF" if verdict == "PASS" else "FAIL_SANSWF_SYMMETRY_STABILITY"

    result = {
        "row": 0.3,
        "stage": "SANSWF",
        "gate": "sanswf_symmetry_stability",
        "signature": "SANSWF_SYMMETRY_STABILITY_PREFLIGHT_v1",
        "context": "LIVE",
        "tier": "Critical",
        "blocking": True,
        "timestamp_utc": utc_now_z(),
        "verdict": verdict,
        "status": status,
        "upstream_dependency": {
            "required": "PASS_ANSWF_EFT_POSITIVITY_TO_SANSWF_HANDOFF",
            "observed": answf_status,
            "pass": answf_pass,
        },
        "benchmark": {
            "lambda_H": LAMBDA_H,
            "lambda_A": LAMBDA_A,
            "lambda_p": LAMBDA_P,
            "generator_normalization": "Tr(Ta Tb)=1/2 delta_ab",
            "two_particle_basis": "unordered real-component pairs with sqrt(2) identical-channel normalization",
            "partial_wave_convention": "a0 = M_const/(16*pi)",
        },
        "checks": checks,
        "claim_guards": claim_guards,
        "metrics": {
            "gauge_SU5_breaking": gauge,
            "bfb_SU5_adjoint": bfb,
            "singlet_contact_matrix": singlet,
            "full_contact_matrix": {
                "n_real_components": N_H + N_ADJ,
                "n_symmetric_two_particle_channels": len(pairs),
                "matrix_symmetry_residual": matrix_symmetry,
                "min_eigenvalue": float(np.min(full_eig)),
                "max_eigenvalue": float(np.max(full_eig)),
                "max_abs_eigenvalue": max_abs_eig,
                "unitarity_limit_8pi": UNITARITY_LIMIT,
                "max_abs_a0": max_abs_a0,
                "a0_limit": 0.5,
                "singlet_global_max_abs_eigenvalue_difference": singlet_full_agreement,
            },
        },
        "corrections_to_legacy_writeup": [
            "For V=lambda_A Tr(Sigma^4), the generic BFB cone must account for the SU(5) traceless-adjoint ratio r=Tr(Sigma^4)/(Tr(Sigma^2))^2 in [7/30,13/20]; setting r=1 is not the exact general condition.",
            "With the explicit canonical real-component/two-particle normalization used by this gate, the benchmark full contact matrix has max |lambda_eig| = %.12g and max |a0| = %.12g; legacy values a0~0.12 and lambda_eig~1.2 are not used." % (max_abs_eig, max_abs_a0),
            "The gauge-boson mass matrix is implemented in the index-safe form M^2=g^2 R R^T with R_ad=f^{acd} v_c, avoiding ambiguous repeated indices in the compact manuscript expression.",
        ],
        "handoff": {
            "from": "SANSWF gauge-consistent, BFB-stable, perturbatively unitary benchmark slice",
            "to": "UTFANSWF unified embedding / anomaly / threshold / seesaw layer",
            "downstream_rows_renumbered": False,
        },
        "open_boundaries": [
            "This gate closes the stated SU(5) Higgs+adjoint tree-level scalar contact matrix at the benchmark; loop-level unitarity is not claimed.",
            "The exact BFB interval is for a traceless Hermitian 5x5 adjoint direction in the stated minimal potential; additional quartic invariants would require a new cone analysis.",
            "Anomaly cancellation, heavy-threshold unification, seesaw consistency, and proton-sector consequences remain downstream UTFANSWF embedding tasks.",
        ],
        "specification": str(SPEC.relative_to(ROOT)),
        "notes": "Critical SANSWF preflight with explicit SU(5) group contractions, corrected BFB geometry, and full 406-channel contact-matrix unitarity check.",
    }

    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps(result, indent=2, sort_keys=True), encoding="utf-8")
    print(json.dumps({"row": 0.3, "stage": "SANSWF", "verdict": verdict, "status": status, "max_abs_eigenvalue": max_abs_eig, "max_abs_a0": max_abs_a0}, indent=2))
    if verdict != "PASS":
        raise SystemExit(1)


if __name__ == "__main__":
    main()
