from __future__ import annotations

import datetime
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" / "zs_foundation" / "result.json"
SPEC = ROOT / "research" / "zero_state" / "ZS_FOUNDATION_SPEC.md"

ABS_TOL = 1.0e-12
PSD_TOL = 1.0e-12
MAXENT_TOL = 1.0e-12
DIMENSIONS = (1, 2, 3, 4, 5, 6)


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


def cyclic_permutation(n: int) -> np.ndarray:
    if n == 1:
        return np.eye(1)
    u = np.zeros((n, n), dtype=float)
    for i in range(n):
        u[(i + 1) % n, i] = 1.0
    return u


def householder(n: int) -> np.ndarray:
    if n == 1:
        return np.array([[-1.0]], dtype=float)
    v = np.arange(1.0, n + 1.0, dtype=float)
    v = v / np.linalg.norm(v)
    return np.eye(n) - 2.0 * np.outer(v, v)


def probe_observables(n: int) -> list[np.ndarray]:
    eye = np.eye(n)
    diag = np.diag(np.arange(1.0, n + 1.0, dtype=float))
    if n == 1:
        offdiag = np.zeros((1, 1), dtype=float)
    else:
        offdiag = np.zeros((n, n), dtype=float)
        for i in range(n - 1):
            offdiag[i, i + 1] = 0.5
            offdiag[i + 1, i] = 0.5
    return [eye, diag, offdiag]


def canonical_kernel(n: int) -> np.ndarray:
    # Deterministic positive spectrum used only to test the structural ZS->Gaussian map.
    # It is not claimed to be a physical Zero-State kernel.
    return np.diag(np.arange(1.0, n + 1.0, dtype=float))


def entropy_from_covariance(sigma: np.ndarray) -> float:
    sign, logdet = np.linalg.slogdet(sigma)
    if sign <= 0:
        return float("-inf")
    n = sigma.shape[0]
    return 0.5 * (n * math.log(2.0 * math.pi * math.e) + logdet)


def alternative_whitened_spectra(n: int) -> list[np.ndarray]:
    if n <= 1:
        return []
    out: list[np.ndarray] = []
    base = np.arange(n, dtype=float) - (n - 1.0) / 2.0
    scale = max(float(np.max(np.abs(base))), 1.0)
    base = base / scale
    for eps in (0.05, 0.15, 0.30):
        d = 1.0 + eps * base
        # base sums to zero, so sum(d)=n exactly up to fp noise.
        if np.min(d) > 0:
            out.append(d)
    return out


def run_dimension(n: int) -> dict[str, Any]:
    z = np.zeros(n, dtype=float)

    observables = probe_observables(n)
    expectations = [float(z @ o @ z) for o in observables]
    linear_probe = np.arange(1.0, n + 1.0, dtype=float)
    linear_response = float(linear_probe @ z)

    symmetries = [np.eye(n), cyclic_permutation(n), householder(n)]
    symmetry_residuals = [float(np.linalg.norm(u @ z - z)) for u in symmetries]
    orthogonality_residuals = [float(np.linalg.norm(u.T @ u - np.eye(n), ord=np.inf)) for u in symmetries]

    k = canonical_kernel(n)
    symmetry_residual_k = float(np.linalg.norm(k - k.T, ord=np.inf))
    eigvals = np.linalg.eigvalsh(k)
    min_eig = float(np.min(eigvals))

    # Probe Q[x]=x^T K x on a deterministic spanning set.
    q_vectors = [np.eye(n)[i] for i in range(n)]
    q_vectors.append(np.ones(n, dtype=float))
    if n > 1:
        q_vectors.append(np.arange(1.0, n + 1.0, dtype=float))
    q_values = [float(x @ k @ x) for x in q_vectors]
    min_q = float(min(q_values))

    # MaxEnt under E[x^T K x]=C with C=n gives Sigma*=K^{-1}.
    c = float(n)
    sigma_star = np.linalg.inv(k)
    trace_constraint = float(np.trace(k @ sigma_star))
    covariance_residual = float(np.linalg.norm(k @ sigma_star - np.eye(n), ord=np.inf))
    h_star = entropy_from_covariance(sigma_star)

    alt_entropies: list[float] = []
    alt_trace_residuals: list[float] = []
    for d in alternative_whitened_spectra(n):
        # Sigma_alt = K^{-1/2} diag(d) K^{-1/2}; trace(K Sigma_alt)=sum(d)=n.
        kinvhalf = np.diag(1.0 / np.sqrt(np.diag(k)))
        sigma_alt = kinvhalf @ np.diag(d) @ kinvhalf
        alt_entropies.append(entropy_from_covariance(sigma_alt))
        alt_trace_residuals.append(abs(float(np.trace(k @ sigma_alt)) - c))

    if alt_entropies:
        min_entropy_margin = float(min(h_star - h for h in alt_entropies))
        max_alt_trace_residual = float(max(alt_trace_residuals))
    else:
        min_entropy_margin = float("inf")
        max_alt_trace_residual = 0.0

    checks = {
        "zero_reference_norm": bool(np.linalg.norm(z) <= ABS_TOL),
        "linear_observables_vanish": bool(max(abs(x) for x in expectations + [linear_response]) <= ABS_TOL),
        "symmetry_invariant_reference": bool(max(symmetry_residuals) <= ABS_TOL),
        "symmetry_operators_orthogonal": bool(max(orthogonality_residuals) <= ABS_TOL),
        "kernel_self_adjoint": bool(symmetry_residual_k <= ABS_TOL),
        "kernel_positive_definite": bool(min_eig > PSD_TOL),
        "quadratic_form_nonnegative": bool(min_q >= -PSD_TOL),
        "maxent_constraint_reproduced": bool(abs(trace_constraint - c) <= ABS_TOL),
        "gaussian_covariance_matches_kernel_inverse": bool(covariance_residual <= ABS_TOL),
        "maxent_gaussian_dominates_deterministic_alternatives": bool(min_entropy_margin >= -MAXENT_TOL),
        "alternative_constraints_preserved": bool(max_alt_trace_residual <= ABS_TOL),
    }

    return {
        "dimension": n,
        "checks": checks,
        "metrics": {
            "max_abs_expectation_or_linear_response": max(abs(x) for x in expectations + [linear_response]),
            "max_symmetry_reference_residual": max(symmetry_residuals),
            "max_symmetry_orthogonality_residual": max(orthogonality_residuals),
            "kernel_symmetry_residual": symmetry_residual_k,
            "kernel_min_eigenvalue": min_eig,
            "min_sampled_quadratic_form": min_q,
            "quadratic_constraint_C": c,
            "trace_KSigma_star": trace_constraint,
            "kernel_covariance_inverse_residual": covariance_residual,
            "gaussian_entropy": h_star,
            "min_entropy_margin_vs_deterministic_alternatives": None if math.isinf(min_entropy_margin) else min_entropy_margin,
            "max_alternative_trace_constraint_residual": max_alt_trace_residual,
        },
    }


def main() -> None:
    per_dimension = [run_dimension(n) for n in DIMENSIONS]
    all_numeric_checks = all(
        all(bool(v) for v in item["checks"].values())
        for item in per_dimension
    )

    structural_claims = {
        "zero_state_is_reference_configuration_not_normalized_quantum_state": True,
        "strict_global_uniqueness_not_claimed": True,
        "operator_spectrum_zero_for_all_admissible_operators_not_imposed": True,
        "construction_C_is_typed_as_map_not_commutator_with_state": True,
        "metric_not_introduced_at_zero_reference": True,
        "canonical_test_kernels_are_structural_discretizations_not_physical_predictions": True,
        "no_physical_mass_length_time_or_coupling_scale_claimed_by_ZS": True,
        "handoff_target_is_SWF_ISM_gaussian_information_state": True,
    }

    all_structural = all(structural_claims.values())
    verdict = "PASS" if all_numeric_checks and all_structural else "FAIL"

    result = {
        "row": 0,
        "stage": "ZS",
        "gate": "zero_state_foundation",
        "signature": "ZS_FOUNDATION_PREFLIGHT_v1",
        "context": "LIVE",
        "tier": "Critical",
        "blocking": True,
        "timestamp_utc": utc_now_z(),
        "verdict": verdict,
        "status": "PASS_ZS_FOUNDATION_TO_SWF_ISM_HANDOFF" if verdict == "PASS" else "FAIL_ZS_FOUNDATION",
        "source_boundary": {
            "role": "Executable hardening of the manuscript Zero State -> minimal quadratic -> MaxEnt Gaussian -> SWF-ISM handoff.",
            "strict_uniqueness_theorem": False,
            "physical_kernel_prediction": False,
            "physical_scale_prediction": False,
        },
        "structural_claims": structural_claims,
        "dimensions_tested": list(DIMENSIONS),
        "per_dimension": per_dimension,
        "handoff": {
            "from": "Zero State reference configuration",
            "construction": "minimal quadratic positive kernel + maximum entropy",
            "to": "SWF-ISM Gaussian information state",
            "wavefunctional_form": "Psi[phi] = N exp(-Q[phi]/2)",
            "kernel_conditions": ["K = K^T (finite real discretization)", "K positive definite in executable test"],
            "downstream_rows_renumbered": False,
        },
        "open_boundaries": [
            "The executable preflight validates the internal ZS construction; it does not prove that no alternative higher-order foundation exists.",
            "The finite-dimensional canonical kernels are implementation probes and are not a physical prediction for the continuum kernel K(x,x').",
            "A concrete continuum/sector coupling remains a downstream SWF-ISM/ANSWF/SANSWF task.",
        ],
        "specification": str(SPEC.relative_to(ROOT)),
        "notes": "Foundational blocking preflight. Preserves historical row numbering; ZS is Row 0 only as an execution label.",
    }

    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, "stage": "ZS", "verdict": verdict, "status": result["status"]}, indent=2))
    if verdict != "PASS":
        raise SystemExit(1)


if __name__ == "__main__":
    main()
