from __future__ import annotations

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

import numpy as np
from scipy.special import sph_harm_y

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "results" / "swf_ism_foundation" / "result.json"
SPEC = ROOT / "research" / "swf_ism" / "SWF_ISM_FOUNDATION_SPEC.md"
ZS_RESULT = ROOT / "results" / "zs_foundation" / "result.json"

ABS_TOL = 2.0e-12
HARMONIC_TOL = 2.0e-11
PROJECTOR_TOL = 2.0e-12
DERIV_TOL = 2.0e-7
ENTROPY_TOL = 2.0e-12
L_MAX = 3
N_X = 40
N_PHI = 96


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


def sphere_quadrature() -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Tensor quadrature for dOmega = dphi d(cos(theta))."""
    x, wx = np.polynomial.legendre.leggauss(N_X)
    phi = 2.0 * math.pi * np.arange(N_PHI, dtype=float) / N_PHI
    wphi = np.full(N_PHI, 2.0 * math.pi / N_PHI, dtype=float)
    theta = np.arccos(x)
    th, ph = np.meshgrid(theta, phi, indexing="ij")
    weights = wx[:, None] * wphi[None, :]
    return x, th, ph, weights


def integrate(field: np.ndarray, weights: np.ndarray) -> complex:
    return np.sum(field * weights)


def entropy(rho: np.ndarray, weights: np.ndarray) -> float:
    if np.any(rho <= 0.0):
        return float("-inf")
    return float(-np.real(integrate(rho * np.log(rho), weights)))


def kl_divergence(rho: np.ndarray, rho0: np.ndarray, weights: np.ndarray) -> float:
    return float(np.real(integrate(rho * np.log(rho / rho0), weights)))


def harmonic_coefficients(rho: np.ndarray, th: np.ndarray, ph: np.ndarray, weights: np.ndarray, lmax: int) -> dict[str, complex]:
    coeffs: dict[str, complex] = {}
    for ell in range(lmax + 1):
        for m in range(-ell, ell + 1):
            y = sph_harm_y(ell, m, th, ph)
            coeffs[f"{ell},{m}"] = complex(integrate(rho * np.conjugate(y), weights))
    return coeffs


def harmonic_orthonormality(th: np.ndarray, ph: np.ndarray, weights: np.ndarray, lmax: int) -> float:
    funcs: list[np.ndarray] = []
    for ell in range(lmax + 1):
        for m in range(-ell, ell + 1):
            funcs.append(sph_harm_y(ell, m, th, ph))
    max_res = 0.0
    for i, yi in enumerate(funcs):
        for j, yj in enumerate(funcs):
            got = integrate(np.conjugate(yi) * yj, weights)
            target = 1.0 if i == j else 0.0
            max_res = max(max_res, abs(got - target))
    return float(max_res)


def exponential_family(kappa: float, x: np.ndarray, weights: np.ndarray) -> tuple[np.ndarray, np.ndarray, float, float, float]:
    p2_x = 0.5 * (3.0 * x * x - 1.0)
    p2 = np.repeat(p2_x[:, None], N_PHI, axis=1)
    unnorm = np.exp(kappa * p2)
    z = float(np.real(integrate(unnorm, weights)))
    rho = unnorm / z
    mean = float(np.real(integrate(rho * p2, weights)))
    var = float(np.real(integrate(rho * (p2 - mean) ** 2, weights)))
    return rho, p2, z, mean, var


def log_partition(kappa: float, x: np.ndarray, weights: np.ndarray) -> float:
    _, _, z, _, _ = exponential_family(kappa, x, weights)
    return math.log(z)


def maxent_derivative_probe(kappa: float, x: np.ndarray, weights: np.ndarray) -> dict[str, float | bool]:
    _, _, _, mean, var = exponential_family(kappa, x, weights)
    h = 1.0e-4
    lm = log_partition(kappa - h, x, weights)
    l0 = log_partition(kappa, x, weights)
    lp = log_partition(kappa + h, x, weights)
    d1 = (lp - lm) / (2.0 * h)
    d2 = (lp - 2.0 * l0 + lm) / (h * h)
    return {
        "kappa": kappa,
        "mean_P2": mean,
        "var_P2": var,
        "d_logZ_numeric": d1,
        "d2_logZ_numeric": d2,
        "d1_abs_error": abs(d1 - mean),
        "d2_abs_error": abs(d2 - var),
        "variance_nonnegative": bool(var >= -ABS_TOL),
        "derivative_identities_pass": bool(abs(d1 - mean) <= DERIV_TOL and abs(d2 - var) <= DERIV_TOL),
    }


def axisymmetry_projection_probe(x: np.ndarray, th: np.ndarray, ph: np.ndarray, weights: np.ndarray) -> dict[str, Any]:
    sin_th = np.sin(th)
    p2 = 0.5 * (3.0 * np.cos(th) ** 2 - 1.0)
    raw = np.exp(0.45 * sin_th * np.cos(ph) + 0.30 * p2)
    raw = raw / float(np.real(integrate(raw, weights)))

    # SO(2)_z projection: average over azimuth at fixed polar angle.
    projected_theta = np.mean(raw, axis=1, keepdims=True)
    projected = np.repeat(projected_theta, N_PHI, axis=1)
    projected = projected / float(np.real(integrate(projected, weights)))

    second_theta = np.mean(projected, axis=1, keepdims=True)
    projected_twice = np.repeat(second_theta, N_PHI, axis=1)
    projected_twice = projected_twice / float(np.real(integrate(projected_twice, weights)))

    raw_coeffs = harmonic_coefficients(raw, th, ph, weights, L_MAX)
    proj_coeffs = harmonic_coefficients(projected, th, ph, weights, L_MAX)

    raw_nonaxis = max((abs(v) for k, v in raw_coeffs.items() if int(k.split(",")[1]) != 0), default=0.0)
    proj_nonaxis = max((abs(v) for k, v in proj_coeffs.items() if int(k.split(",")[1]) != 0), default=0.0)

    return {
        "raw_normalization_error": abs(float(np.real(integrate(raw, weights))) - 1.0),
        "projected_normalization_error": abs(float(np.real(integrate(projected, weights))) - 1.0),
        "projected_min_density": float(np.min(projected)),
        "idempotence_Linf": float(np.max(np.abs(projected_twice - projected))),
        "raw_max_abs_m_nonzero_coefficient": float(raw_nonaxis),
        "projected_max_abs_m_nonzero_coefficient": float(proj_nonaxis),
        "projection_suppression_ratio": float(proj_nonaxis / raw_nonaxis) if raw_nonaxis > 0 else 0.0,
        "pass": bool(
            abs(float(np.real(integrate(projected, weights))) - 1.0) <= ABS_TOL
            and np.min(projected) >= 0.0
            and np.max(np.abs(projected_twice - projected)) <= PROJECTOR_TOL
            and proj_nonaxis <= HARMONIC_TOL
        ),
    }


def spectral_probe() -> dict[str, Any]:
    # Synthetic positive spectral measure used only as a consistency test of the stated
    # Kallen-Lehmann-type nonnegative-weight representation; not a physical spectrum.
    masses2 = np.array([0.4, 1.1, 2.7, 5.0], dtype=float)
    weights = np.array([0.15, 0.35, 0.30, 0.20], dtype=float)
    q2 = np.array([0.0, 0.25, 1.0, 4.0, 12.0], dtype=float)
    d_e = np.array([np.sum(weights / (qq + masses2)) for qq in q2], dtype=float)
    deriv = np.array([-np.sum(weights / (qq + masses2) ** 2) for qq in q2], dtype=float)
    return {
        "weights_sum": float(np.sum(weights)),
        "min_weight": float(np.min(weights)),
        "min_euclidean_propagator": float(np.min(d_e)),
        "max_derivative": float(np.max(deriv)),
        "positive_measure": bool(np.min(weights) >= 0.0 and abs(np.sum(weights) - 1.0) <= ABS_TOL),
        "positive_euclidean_response": bool(np.min(d_e) > 0.0),
        "monotone_euclidean_response": bool(np.max(deriv) < 0.0 and np.all(np.diff(d_e) < 0.0)),
    }


def main() -> None:
    zs_pass = False
    zs_status = "MISSING"
    if ZS_RESULT.exists():
        zs_obj = json.loads(ZS_RESULT.read_text(encoding="utf-8"))
        zs_status = str(zs_obj.get("status", "UNKNOWN"))
        zs_pass = bool(zs_obj.get("verdict") == "PASS" and zs_status == "PASS_ZS_FOUNDATION_TO_SWF_ISM_HANDOFF")

    x, th, ph, weights = sphere_quadrature()
    area = float(np.real(integrate(np.ones_like(th), weights)))

    rho0 = np.full_like(th, 1.0 / (4.0 * math.pi), dtype=float)
    norm0 = float(np.real(integrate(rho0, weights)))
    h0 = entropy(rho0, weights)
    hmax = math.log(4.0 * math.pi)
    coeff0 = harmonic_coefficients(rho0, th, ph, weights, L_MAX)
    a00_target = 1.0 / math.sqrt(4.0 * math.pi)
    a00_error = abs(coeff0["0,0"] - a00_target)
    higher_max = max((abs(v) for k, v in coeff0.items() if k != "0,0"), default=0.0)
    ortho_res = harmonic_orthonormality(th, ph, weights, L_MAX)

    # Reconstruct rho0 from computed coefficients through l=L_MAX.
    reconstruction = np.zeros_like(th, dtype=complex)
    for key, coeff in coeff0.items():
        ell, m = (int(v) for v in key.split(","))
        reconstruction += coeff * sph_harm_y(ell, m, th, ph)
    recon_error = float(np.max(np.abs(reconstruction.real - rho0)))
    recon_imag = float(np.max(np.abs(reconstruction.imag)))

    exp_probes: list[dict[str, Any]] = []
    for kappa in (-1.0, -0.35, 0.35, 1.0):
        rho, _, z, _, _ = exponential_family(kappa, x, weights)
        h = entropy(rho, weights)
        dkl = kl_divergence(rho, rho0, weights)
        deriv_probe = maxent_derivative_probe(kappa, x, weights)
        exp_probes.append({
            "kappa": kappa,
            "normalization_error": abs(float(np.real(integrate(rho, weights))) - 1.0),
            "min_density": float(np.min(rho)),
            "partition_Z": z,
            "entropy": h,
            "entropy_gap_from_uniform": hmax - h,
            "KL_to_uniform": dkl,
            "entropy_KL_identity_error": abs((hmax - h) - dkl),
            "derivatives": deriv_probe,
        })

    projector = axisymmetry_projection_probe(x, th, ph, weights)
    spectral = spectral_probe()

    checks = {
        "zs_handoff_passes": zs_pass,
        "sphere_measure_is_4pi": abs(area - 4.0 * math.pi) <= ABS_TOL,
        "uniform_state_nonnegative": bool(np.min(rho0) >= 0.0),
        "uniform_state_normalized": abs(norm0 - 1.0) <= ABS_TOL,
        "uniform_a00_matches": a00_error <= HARMONIC_TOL,
        "uniform_higher_harmonics_vanish": higher_max <= HARMONIC_TOL,
        "low_order_harmonics_orthonormal": ortho_res <= HARMONIC_TOL,
        "uniform_low_order_reconstruction": recon_error <= HARMONIC_TOL and recon_imag <= HARMONIC_TOL,
        "uniform_entropy_is_log_4pi": abs(h0 - hmax) <= ENTROPY_TOL,
        "nonuniform_exponential_states_positive_normalized": all(p["normalization_error"] <= ABS_TOL and p["min_density"] > 0.0 for p in exp_probes),
        "uniform_maxent_dominates_nonuniform_probes": all(p["entropy_gap_from_uniform"] > 0.0 for p in exp_probes),
        "relative_entropy_nonnegative": all(p["KL_to_uniform"] >= -ENTROPY_TOL for p in exp_probes),
        "entropy_relative_entropy_identity": all(p["entropy_KL_identity_error"] <= 5.0e-12 for p in exp_probes),
        "exponential_family_logZ_identities": all(p["derivatives"]["derivative_identities_pass"] for p in exp_probes),
        "axisymmetry_projection_is_valid_projector": bool(projector["pass"]),
        "positive_spectral_weight_probe": bool(spectral["positive_measure"] and spectral["positive_euclidean_response"] and spectral["monotone_euclidean_response"]),
    }

    claim_guards = {
        "swf_ism_is_information_state_layer_not_eft_positivity_gate": True,
        "answf_eft_filter_not_claimed_by_this_gate": True,
        "sanswf_gauge_symmetry_enforcement_not_claimed_by_this_gate": True,
        "synthetic_spectral_probe_not_physical_spectrum_prediction": True,
        "finite_harmonic_quadrature_is_numerical_consistency_probe": True,
        "no_measured_physics_target_used": True,
        "no_downstream_historical_rows_renumbered": True,
    }

    all_pass = all(checks.values()) and all(claim_guards.values())
    verdict = "PASS" if all_pass else "FAIL"
    status = "PASS_SWF_ISM_INFORMATION_STATE_FOUNDATION_TO_ANSWF_HANDOFF" if all_pass else "FAIL_SWF_ISM_INFORMATION_STATE_FOUNDATION"

    result = {
        "row": 0.1,
        "stage": "SWF-ISM",
        "gate": "swf_ism_information_state_foundation",
        "signature": "SWF_ISM_FOUNDATION_PREFLIGHT_v1",
        "context": "LIVE",
        "tier": "Critical",
        "blocking": True,
        "timestamp_utc": utc_now_z(),
        "verdict": verdict,
        "status": status,
        "upstream_dependency": {
            "required": "PASS_ZS_FOUNDATION_TO_SWF_ISM_HANDOFF",
            "observed": zs_status,
            "pass": zs_pass,
        },
        "checks": checks,
        "claim_guards": claim_guards,
        "metrics": {
            "quadrature": {
                "N_x": N_X,
                "N_phi": N_PHI,
                "sphere_area": area,
                "sphere_area_abs_error": abs(area - 4.0 * math.pi),
            },
            "uniform_information_state": {
                "normalization": norm0,
                "min_density": float(np.min(rho0)),
                "a00": {"real": coeff0["0,0"].real, "imag": coeff0["0,0"].imag},
                "a00_target": a00_target,
                "a00_abs_error": a00_error,
                "max_abs_higher_harmonic": float(higher_max),
                "harmonic_orthonormality_max_residual": ortho_res,
                "low_order_reconstruction_Linf": recon_error,
                "low_order_reconstruction_max_imag": recon_imag,
                "entropy": h0,
                "entropy_target_log_4pi": hmax,
                "entropy_abs_error": abs(h0 - hmax),
            },
            "exponential_family_probes": exp_probes,
            "axisymmetry_projection_probe": projector,
            "spectral_representation_probe": spectral,
        },
        "handoff": {
            "from": "SWF-ISM normalized nonnegative information state with harmonic/MaxEnt structure",
            "to": "ANSWF EFT admissibility / positivity layer",
            "exports": [
                "normalized nonnegative rho(Omega)",
                "harmonic coefficients a_lm",
                "maximum-entropy reference state",
                "Gaussian/spectral consistency structure inherited from ZS",
            ],
            "downstream_rows_renumbered": False,
        },
        "open_boundaries": [
            "This preflight validates the stated information-state mathematics; it does not derive a unique physical continuum field kernel or spectral density.",
            "The synthetic spectral measure is a representation-consistency probe only, not a particle-spectrum prediction.",
            "Mapping SWF-ISM moments into scattering amplitudes and enforcing EFT positivity is deferred to ANSWF.",
            "Gauge covariance, bounded-from-below cones, and partial-wave unitarity enforcement are deferred to SANSWF.",
        ],
        "specification": str(SPEC.relative_to(ROOT)),
        "notes": "Critical foundational preflight between ZS and ANSWF. Execution label 0.1 preserves historical row numbering.",
    }

    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.1, "stage": "SWF-ISM", "verdict": verdict, "status": status}, indent=2))
    if verdict != "PASS":
        raise SystemExit(1)


if __name__ == "__main__":
    main()
