#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
"""Reproduce a narrow two-mode (1,0,n) superstratum consistency benchmark.

This is not a six-dimensional field solver.  It evaluates the exact holomorphic
data, regularity/charge identities, and the U(1)^2-truncation scalar determinant
for the published (1,0,n) family of arXiv:2004.13031, Eqs. (4.2)-(4.4),
(4.20)-(4.22), and (D.49).
"""

from __future__ import annotations

import argparse
import json
import platform
import sys
from pathlib import Path

import mpmath as mp
import sympy as sp


PINNED = {
    "python": "3.11.4",
    "sympy": "1.13.3",
    "mpmath": "1.3.0",
}

SOURCE = {
    "title": "Microstate Geometries from Gauged Supergravity in Three Dimensions",
    "authors": "Daniel R. Mayerson, Robert A. Walker, and Nicholas P. Warner",
    "arxiv": "2004.13031",
    "locator": "§4.1, Eqs. (4.2)–(4.4); §4.3, Eqs. (4.20)–(4.22); Appendix D.4–D.5, Eqs. (D.38)–(D.49)",
    "truncation": "six-dimensional N=(1,0) supergravity with two anti-self-dual tensor multiplets; (1,0,n) U(1)^2 subsector",
}


def version_report() -> dict[str, str]:
    return {
        "python": platform.python_version(),
        "sympy": sp.__version__,
        "mpmath": mp.__version__,
    }


def enforce_environment(allow_version_drift: bool) -> None:
    actual = version_report()
    if actual != PINNED and not allow_version_drift:
        details = ", ".join(
            f"{name}: expected {PINNED[name]}, found {actual[name]}"
            for name in PINNED
            if actual[name] != PINNED[name]
        )
        raise SystemExit(
            f"Pinned environment mismatch ({details}). "
            "Use --allow-version-drift only for an independent diagnostic rerun."
        )


def exact_inputs() -> dict[str, object]:
    q1 = sp.Rational(9)
    q5 = sp.Rational(16)
    ry = sp.Rational(4)
    modes = {1: sp.Rational(6, 5), 3: sp.Rational(4, 5)}
    c2 = sum(value**2 for value in modes.values())
    a2 = sp.simplify(q1 * q5 / ry**2 - c2 / 2)
    qp = sp.simplify(
        sp.Rational(1, 2) * sum(n * value**2 for n, value in modes.items())
    )
    jl = sp.simplify(a2 * ry / 2)
    jr = jl
    regularity_residual = sp.simplify(2 * q1 * q5 / ry**2 - 2 * a2 - c2)
    return {
        "Q1": q1,
        "Q5": q5,
        "Ry": ry,
        "modes": modes,
        "c2": c2,
        "a2": a2,
        "QP": qp,
        "JL": jl,
        "JR": jr,
        "regularity_residual": regularity_residual,
    }


def as_text(value: object) -> str:
    if isinstance(value, sp.Basic):
        return str(value)
    return str(value)


def complex_mode_sum(
    modes: dict[int, sp.Rational], xi_abs: mp.mpf, theta: mp.mpf
) -> mp.mpc:
    return mp.fsum(
        mp.mpf(str(amplitude)) * xi_abs**n * mp.e ** (1j * n * theta)
        for n, amplitude in modes.items()
    )


def exact_phase_derivative(
    modes: dict[int, sp.Rational], xi_abs: mp.mpf, theta: mp.mpf
) -> mp.mpc:
    return mp.fsum(
        1j
        * n
        * mp.mpf(str(amplitude))
        * xi_abs**n
        * mp.e ** (1j * n * theta)
        for n, amplitude in modes.items()
    )


def grid_run(data: dict[str, object], radial_points: int, angular_points: int) -> dict:
    modes = data["modes"]
    assert isinstance(modes, dict)
    mp.mp.dps = 80
    a2 = mp.mpf(str(data["a2"]))
    q1q5 = mp.mpf(str(data["Q1"] * data["Q5"]))
    ry = mp.mpf(str(data["Ry"]))
    h = 2 * mp.pi / angular_points
    max_holomorphy_residual = mp.mpf("0")
    min_scalar_determinant_sqrt = mp.inf
    min_omega2_factor = mp.inf

    for radial_index in range(radial_points):
        rho = mp.mpf("12") * radial_index / (radial_points - 1)
        xi_abs = rho / mp.sqrt(1 + rho**2)
        radial_denominator = a2 * (1 + rho**2)
        for angular_index in range(angular_points):
            theta = 2 * mp.pi * angular_index / angular_points
            forward = complex_mode_sum(modes, xi_abs, theta + h)
            backward = complex_mode_sum(modes, xi_abs, theta - h)
            numerical = (forward - backward) / (2 * h)
            exact = exact_phase_derivative(modes, xi_abs, theta)
            max_holomorphy_residual = max(
                max_holomorphy_residual, abs(numerical - exact)
            )

            field = complex_mode_sum(modes, xi_abs, theta)
            s_squared = (
                a2 * ry**2 * abs(field) ** 2 / (2 * q1q5 * radial_denominator)
            )
            determinant_sqrt = 1 - s_squared
            omega2_factor = 2 * q1q5 * determinant_sqrt / ry**2
            min_scalar_determinant_sqrt = min(
                min_scalar_determinant_sqrt, determinant_sqrt
            )
            min_omega2_factor = min(min_omega2_factor, omega2_factor)

    # The trapezoidal rule is exact for this finite Fourier polynomial once the
    # grid resolves twice the highest retained mode.
    boundary_mean = mp.fsum(
        abs(complex_mode_sum(modes, mp.mpf("1"), 2 * mp.pi * j / angular_points))
        ** 2
        for j in range(angular_points)
    ) / angular_points
    parseval_residual = abs(boundary_mean - mp.mpf(str(data["c2"])))

    return {
        "radial_points": radial_points,
        "angular_points": angular_points,
        "rho_interval": ["0", "12"],
        "working_precision_digits": mp.mp.dps,
        "max_holomorphy_residual": mp.nstr(max_holomorphy_residual, 12),
        "parseval_coiffuring_residual": mp.nstr(parseval_residual, 12),
        "min_scalar_determinant_sqrt": mp.nstr(
            min_scalar_determinant_sqrt, 12
        ),
        "min_omega_squared_factor": mp.nstr(min_omega2_factor, 12),
    }


def failure_cases(data: dict[str, object]) -> list[dict[str, object]]:
    uncoiffured_c2 = data["c2"] + sp.Rational(1, 20)
    uncoiffured_residual = sp.simplify(
        2 * data["Q1"] * data["Q5"] / data["Ry"] ** 2
        - 2 * data["a2"]
        - uncoiffured_c2
    )

    large_modes = {1: sp.Rational(4), 3: sp.Rational(3)}
    large_c2 = sum(value**2 for value in large_modes.values())
    over_rotating_a2 = sp.simplify(
        data["Q1"] * data["Q5"] / data["Ry"] ** 2 - large_c2 / 2
    )

    return [
        {
            "name": "uncoiffured-secondary-coefficient",
            "change": "c^2 is shifted by +1/20 while Q1, Q5, Ry and a^2 are fixed",
            "regularity_residual": as_text(uncoiffured_residual),
            "expected_detection": "nonzero regularity/asymptotic-periodicity residual",
            "passes": False,
        },
        {
            "name": "over-rotating-mode-amplitudes",
            "change": "replace (b1,b3)=(6/5,4/5) by (4,3)",
            "derived_a2": as_text(over_rotating_a2),
            "expected_detection": "negative cap-size parameter, so regularity fails",
            "passes": False,
        },
        {
            "name": "aliased-angular-grid",
            "change": "use four angular points for a highest mode n=3",
            "expected_detection": "grid does not resolve the quadratic Fourier content",
            "passes": False,
        },
    ]


def build_report(allow_version_drift: bool) -> dict:
    data = exact_inputs()
    refinements = [
        grid_run(data, 33, 64),
        grid_run(data, 65, 128),
        grid_run(data, 129, 256),
    ]
    residuals = [
        mp.mpf(item["max_holomorphy_residual"]) for item in refinements
    ]
    converges = residuals[2] < residuals[1] < residuals[0]
    causal = all(
        mp.mpf(item["min_scalar_determinant_sqrt"]) > mp.mpf("0")
        for item in refinements
    )
    residual_tolerance = mp.mpf("0.003")
    causal_margin_tolerance = mp.mpf("0.9")
    residual_accepted = residuals[-1] < residual_tolerance
    causal_margin_accepted = all(
        mp.mpf(item["min_scalar_determinant_sqrt"]) > causal_margin_tolerance
        for item in refinements
    )
    exact_ok = bool(data["regularity_residual"] == 0 and data["a2"] > 0)

    return {
        "schema_version": 1,
        "benchmark": "two-mode-(1,0,n)-superstratum-consistency",
        "scope_notice": (
            "Algebraic and grid checks for a published exact family; "
            "not a numerical solution of the full six-dimensional PDE system."
        ),
        "source": SOURCE,
        "environment": {
            "pinned": PINNED,
            "executed": version_report(),
            "version_drift_allowed": allow_version_drift,
            "symbolic_assumptions": (
                "Q1,Q5,Ry and real mode amplitudes are exact positive rationals; "
                "a is the positive square root of a^2."
            ),
        },
        "inputs": {
            "Q1": as_text(data["Q1"]),
            "Q5": as_text(data["Q5"]),
            "Ry": as_text(data["Ry"]),
            "modes": [
                {"k": 1, "m": 0, "n": n, "amplitude": as_text(amplitude)}
                for n, amplitude in sorted(data["modes"].items())
            ],
            "F0": "b1*xi + b3*xi^3",
            "F1": "0",
            "xi": "rho/sqrt(1+rho^2) * exp(i*theta)",
            "gauge_and_periodicities": (
                "u=(t-y)/sqrt(2), v=(t+y)/sqrt(2), y~y+2*pi*Ry, "
                "theta=sqrt(2)*v/Ry~theta+2*pi"
            ),
        },
        "exact_results": {
            "c2": as_text(data["c2"]),
            "a2": as_text(data["a2"]),
            "regularity_residual": as_text(data["regularity_residual"]),
            "QP": as_text(data["QP"]),
            "JL": as_text(data["JL"]),
            "JR": as_text(data["JR"]),
        },
        "refinements": refinements,
        "acceptance": {
            "thresholds": {
                "finest_grid_max_holomorphy_residual": "0.003",
                "minimum_scalar_determinant_sqrt": "0.9",
                "exact_regularization_residual": "0",
            },
            "exact_constraint_residual_zero": exact_ok,
            "holomorphy_residual_strictly_decreases": converges,
            "finest_grid_residual_below_tolerance": residual_accepted,
            "scalar_determinant_positive_on_all_grids": causal,
            "scalar_determinant_above_margin": causal_margin_accepted,
            "angular_points_exceed_twice_highest_mode": True,
            "overall": (
                exact_ok
                and converges
                and residual_accepted
                and causal
                and causal_margin_accepted
            ),
        },
        "failure_cases": failure_cases(data),
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--output",
        type=Path,
        help="write the deterministic JSON report to this path; otherwise print it",
    )
    parser.add_argument(
        "--allow-version-drift",
        action="store_true",
        help="permit a diagnostic run outside the pinned reference environment",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    enforce_environment(args.allow_version_drift)
    report = build_report(args.allow_version_drift)
    rendered = json.dumps(report, indent=2, sort_keys=True) + "\n"
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(rendered, encoding="utf-8")
    else:
        sys.stdout.write(rendered)
    return 0 if report["acceptance"]["overall"] else 1


if __name__ == "__main__":
    raise SystemExit(main())