#!/usr/bin/env python3
"""Reproduce the two-instanton modified-Mathieu NS capstone.

Source convention:

    H_s = -hbar^2 d^2/dQ^2 + 2 Lambda^2 cosh(Q),
    H_s psi = E_s psi,  psi in L^2(R).

The script solves the model-specific Nekrasov–Shatashvili condition

    dF_NS/da = pi*hbar*(n + 1/2)

at perturbative, one-instanton, and two-instanton order, then applies the
quantum Matone relation.  The default Lambda=hbar=1 results are compared
with the independently refined DCHE targets from the chapter.  The displayed
truncation shifts are diagnostics, not rigorous remainder bounds.

Use --recompute-direct to regenerate both the parity-grid and DCHE shooting
spectra with the neighboring modified-mathieu-dche-spectrum.py program.
This mode never reads DCHE_TARGETS. The exact map is Q=2x,
kappa=4*(Lambda/hbar)^2, E_s=hbar^2*A/4. It supports 1<=kappa<=9 and
2<=levels<=8, the direct solver's calibrated range. --json saves all
normalizations, settings, direct refinements, NS orders and truncation gaps.
For example:

    python3 modified-mathieu-ns-capstone.py --recompute-direct --json ns.json

Both .py files must be in the same directory. NumPy and SciPy are required.
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import math
import platform
import sys
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np

try:
    import scipy
    from scipy.optimize import brentq
    from scipy.special import loggamma
except ImportError as exc:  # pragma: no cover - user-facing dependency guard
    raise SystemExit(
        "modified-mathieu-ns-capstone.py requires SciPy and NumPy"
    ) from exc


DCHE_TARGETS = np.array(
    [
        3.059174596896,
        5.285125967380,
        7.714579573227,
        10.327666944456,
    ],
    dtype=float,
)


@dataclass(frozen=True)
class InstantonTerms:
    first: float
    second: float
    derivative_first: float
    derivative_second: float


def require(condition: bool, message: str) -> None:
    if not condition:
        raise ValueError(message)


def gamma_term(a_value: float, hbar: float, scale: float) -> float:
    """Return the real source-convention gamma(a,hbar,Lambda)."""

    logarithm = 0.5 * a_value * math.log(hbar * hbar / (scale * scale))
    phase = hbar * float(loggamma(1.0 + 1j * a_value / hbar).imag)
    return logarithm - 0.25 * math.pi * hbar + phase


def instanton_terms(
    a_value: float, hbar: float, scale: float
) -> InstantonTerms:
    a_squared = a_value * a_value
    hbar_squared = hbar * hbar
    first_denominator = a_squared + hbar_squared
    first = -2.0 * scale**4 / first_denominator
    derivative_first = 4.0 * a_value * scale**4 / first_denominator**2

    numerator = 7.0 * hbar_squared - 5.0 * a_squared
    denominator = (
        first_denominator**3 * (a_squared + 4.0 * hbar_squared)
    )
    second = scale**8 * numerator / denominator
    derivative_denominator = denominator * (
        6.0 * a_value / first_denominator
        + 2.0 * a_value / (a_squared + 4.0 * hbar_squared)
    )
    derivative_second = scale**8 * (
        -10.0 * a_value * denominator
        - numerator * derivative_denominator
    ) / denominator**2
    return InstantonTerms(
        first, second, derivative_first, derivative_second
    )


def free_energy_derivative(
    a_value: float, hbar: float, scale: float, instanton_order: int
) -> float:
    value = 2.0 * gamma_term(a_value, hbar, scale)
    terms = instanton_terms(a_value, hbar, scale)
    if instanton_order >= 1:
        value += terms.derivative_first
    if instanton_order >= 2:
        value += terms.derivative_second
    return value


def solve_flat_coordinate(
    index: int, hbar: float, scale: float, instanton_order: int
) -> float:
    target = math.pi * hbar * (index + 0.5)

    def equation(a_value: float) -> float:
        return (
            free_energy_derivative(
                a_value, hbar, scale, instanton_order
            )
            - target
        )

    lower = 1.0e-10 * hbar
    require(equation(lower) < 0.0, "unexpected lower NS bracket sign")
    upper = max(hbar, scale)
    for _ in range(100):
        if equation(upper) > 0.0:
            return float(
                brentq(
                    equation,
                    lower,
                    upper,
                    xtol=5.0e-14,
                    rtol=2.0e-14,
                    maxiter=160,
                )
            )
        upper *= 1.35
    raise ValueError(f"could not bracket NS root {index}")


def matone_energy(
    a_value: float, hbar: float, scale: float, instanton_order: int
) -> float:
    """Apply E_s=2u_s=a^2/4-(Lambda/4)d_Lambda F_inst."""

    terms = instanton_terms(a_value, hbar, scale)
    energy = 0.25 * a_value * a_value
    if instanton_order >= 1:
        energy -= terms.first
    if instanton_order >= 2:
        energy -= 2.0 * terms.second
    return energy


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--hbar", type=float, default=1.0)
    parser.add_argument("--Lambda", dest="scale", type=float, default=1.0)
    parser.add_argument("--levels", type=int, default=4)
    parser.add_argument("--recompute-direct", action="store_true",
                        help="regenerate grid and DCHE references; no stored targets")
    parser.add_argument("--direct-mode", choices=("default", "high"), default="high")
    parser.add_argument("--direct-tolerance", type=float,
                        help="optional absolute grid/DCHE agreement threshold in native A units")
    parser.add_argument("--json", type=Path,
                        help="full reproduction record for --recompute-direct")
    args = parser.parse_args()
    if not math.isfinite(args.hbar) or args.hbar <= 0:
        parser.error("--hbar must be positive and finite")
    if not math.isfinite(args.scale) or args.scale <= 0:
        parser.error("--Lambda must be positive and finite")
    if not 1 <= args.levels <= 12:
        parser.error("--levels must lie between 1 and 12")
    if args.recompute_direct:
        ratio = args.scale / args.hbar
        if not math.isfinite(ratio) or not 0.5 <= ratio <= 1.5:
            parser.error("--recompute-direct requires 1 <= kappa=4*(Lambda/hbar)^2 <= 9")
        if not 2 <= args.levels <= 8:
            parser.error("--recompute-direct supports 2 <= levels <= 8")
        if not math.isfinite(args.hbar*args.hbar) or args.hbar*args.hbar/4 == 0:
            parser.error("hbar^2/4 must be representable as a positive binary64 value")
        if args.direct_tolerance is not None and (
            not math.isfinite(args.direct_tolerance) or args.direct_tolerance <= 0
        ):
            parser.error("--direct-tolerance must be positive and finite")
    elif args.json is not None or args.direct_tolerance is not None:
        parser.error("--json and --direct-tolerance require --recompute-direct")
    return args


def load_direct_solver():
    """Load the actual public direct solver, retaining dataclass module identity."""
    path = Path(__file__).with_name("modified-mathieu-dche-spectrum.py")
    if not path.exists():
        raise ValueError("Download modified-mathieu-dche-spectrum.py into the same directory.")
    name = "modified_mathieu_direct_for_ns"
    specification = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(specification)
    sys.modules[name] = module
    specification.loader.exec_module(module)
    return module, path


def recompute_direct_comparison(arguments: argparse.Namespace) -> int:
    """Compare dimensionless NS approximants with freshly regenerated spectra."""
    direct, direct_path = load_direct_solver()
    ratio = arguments.scale / arguments.hbar
    kappa = 4*ratio*ratio
    energy_scale = arguments.hbar*arguments.hbar / 4
    settings = direct.settings_for_mode(arguments.direct_mode)
    report = {
        "environment": {"Python": platform.python_version(), "NumPy": np.__version__, "SciPy": scipy.__version__},
        "inputs": {key: str(value) if isinstance(value, Path) else value for key, value in vars(arguments).items()},
        "source_operator": "-hbar_s^2 d_Q^2 + 2 Lambda^2 cosh(Q) on L2(R)",
        "direct_operator": "-d_x^2 + 2 kappa cosh(2x) on L2(R)",
        "normalization_map": {
            "coordinate": "Q=2*x", "kappa": kappa,
            "energy": "E_s=(hbar_s^2/4)*A", "energy_scale": energy_scale,
            "normalized_wavefunction": "psi_s(Q)=Phi(Q/2)/sqrt(2) when both L2 norms are one",
            "NS_dimensionless_coordinate": "a_s/hbar_s",
            "NS_dimensionless_energy": "E_s/hbar_s^2",
            "loggamma_branch": "analytic logGamma(1+i*a_s/hbar_s) used by scipy.special.loggamma",
            "DCHE": "alpha=1; gamma=2; delta=-4*kappa; q_D=2*kappa-A-1/4; negative zeta ray",
        },
        "reference_generation": {
            "stored_targets_used": False,
            "direct_program": direct_path.name,
            "direct_program_sha256": hashlib.sha256(direct_path.read_bytes()).hexdigest(),
            "NS_program_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
            "settings": asdict(settings),
            "independence_scope": "grid and DCHE solve different representations; DCHE bracketing is initialized by the computed grid eigenvalue",
        },
        "checks": [],
        "qualification": "finite instanton truncation; empirical floating-point refinements; no remainder bound, Borel resummation or interval certificate",
    }

    def check(name, value, tolerance):
        passed = math.isfinite(value) and value <= tolerance
        report["checks"].append({"name": name, "value": value, "tolerance": tolerance, "passed": passed})
        print(f"{'PASS' if passed else 'FAIL'}: {name} = {value:.3e} (limit {tolerance:.3e})", flush=True)

    print("Modified-Mathieu NS comparison with regenerated direct spectra")
    print(f"Lambda={arguments.scale:g}; hbar_s={arguments.hbar:g}; kappa={kappa:g}; levels={arguments.levels}")
    print("Q=2x; E_s=(hbar_s^2/4) A; no stored targets are used.")
    print("Environment:", report["environment"])
    print("Computing parity grids and DCHE shooting with profile", arguments.direct_mode, flush=True)
    try:
        grid = direct.grid_spectrum(kappa, settings.domain, settings.grid_coarse_cells,
                                    (arguments.levels+1)//2)
        spectral_rows = direct.build_rows(kappa, grid, settings, arguments.levels)
        off_shell_energy = (spectral_rows[0].dche+spectral_rows[1].dche)/2
        abel_spread, invariants = direct.abel_invariant_spread(
            kappa, off_shell_energy, settings.series_fine_order, settings.s_fine_min,
            settings.s_fine_max, settings.zeta_fine_step_factor,
        )
        thresholds = direct.ACCEPTANCE_THRESHOLDS[arguments.direct_mode]
        report["direct_refinements"] = {
            "fine_grid_even_A": grid.fine_even.tolist(), "fine_grid_odd_A": grid.fine_odd.tolist(),
            "Richardson_even_A": grid.even.tolist(), "Richardson_odd_A": grid.odd.tolist(),
            "maximum_Richardson_shift_A": grid.maximum_shift,
            "sturm_bisection_stagnation_width_A": grid.sturm_width,
            "maximum_Richardson_shift_E_s": energy_scale*grid.maximum_shift,
            "within_run_Abel_spread": abel_spread,
            "Abel_off_shell_A": off_shell_energy, "Abel_samples": invariants,
            "Abel_scope": "within-run propagation check; not an independent endpoint-normalization certificate",
            "native_A_thresholds": thresholds,
        }
        # Preserve all direct solver checks, including spectral order, parity,
        # potential minimum and its own documented refinement thresholds.
        try:
            direct.run_acceptance_checks(arguments.direct_mode, kappa, spectral_rows, grid, abel_spread)
            report["checks"].append({"name": "direct solver calibrated checks", "passed": True})
        except RuntimeError as error:
            report["checks"].append({"name": "direct solver calibrated checks", "passed": False, "message": str(error)})
        print("\n n parity        grid E_s        DCHE E_s          NS order 0          NS order 1          NS order 2", flush=True)
        rows = []
        ns_residuals = []
        for row in spectral_rows:
            if not all(math.isfinite(energy_scale*value) and energy_scale*value > 0
                       for value in (row.grid, row.dche)):
                raise ValueError("source energy is outside the representable binary64 range")
            ns = []
            for order in range(3):
                # Exact homogeneity removes large/small common dimensional
                # factors from root finding and preserves the logGamma branch.
                a_reduced = solve_flat_coordinate(row.index, 1.0, ratio, order)
                e_reduced = matone_energy(a_reduced, 1.0, ratio, order)
                e_source = arguments.hbar*arguments.hbar*e_reduced
                if not math.isfinite(e_source):
                    raise ValueError("NS source energy is outside the representable binary64 range")
                residual = abs(free_energy_derivative(a_reduced, 1.0, ratio, order)
                               - math.pi*(row.index+0.5))
                ns_residuals.append(residual)
                ns.append({"instanton_order": order, "a_s_over_hbar_s": a_reduced,
                    "a_s": arguments.hbar*a_reduced, "E_s": e_source,
                    "normalized_NS_root_residual": residual,
                    "gap_to_grid_E_s": abs(e_source-energy_scale*row.grid),
                    "gap_to_DCHE_E_s": abs(e_source-energy_scale*row.dche)})
            rows.append({"index": row.index, "parity": row.parity,
                "grid_A": row.grid, "DCHE_A": row.dche,
                "grid_E_s": energy_scale*row.grid, "DCHE_E_s": energy_scale*row.dche,
                "grid_DCHE_gap_A": row.cross_residual,
                "grid_DCHE_gap_E_s": energy_scale*row.cross_residual,
                "DCHE_joint_refinement_shift_A": row.dche_refinement_shift,
                "DCHE_joint_refinement_shift_E_s": energy_scale*row.dche_refinement_shift,
                "NS_approximants": ns,
                "NS_gaps_decrease_with_order": all(ns[n+1]["gap_to_DCHE_E_s"] < ns[n]["gap_to_DCHE_E_s"] for n in (0, 1))})
            print(f"{row.index:2d} {row.parity:>4s} {energy_scale*row.grid:17.11f} {energy_scale*row.dche:17.11f} "
                  + " ".join(f"{item['E_s']:19.11f}" for item in ns), flush=True)
        report["rows"] = rows
        check("normalized NS root residual", max(ns_residuals), 2e-11)
        direct_tolerance = arguments.direct_tolerance or thresholds["grid_dche_gap"]
        check("regenerated grid versus DCHE gap in A", max(row.cross_residual for row in spectral_rows), direct_tolerance)
        check("DCHE joint refinement in A", max(row.dche_refinement_shift for row in spectral_rows), thresholds["dche_refinement_shift"])
        check("grid Richardson shift in A", grid.maximum_shift, thresholds["grid_refinement_shift"])
        print("Finite-instanton gaps are truncation diagnostics, not verification thresholds.")
    except (RuntimeError, ValueError, OverflowError, ZeroDivisionError) as error:
        report["execution_error"] = str(error)
        print("Numerical execution failed:", error, file=sys.stderr)
    passed = "execution_error" not in report and all(item["passed"] for item in report["checks"])
    report["verification_passed"] = passed
    print("Recomputed comparison:", "PASS" if passed else "FAIL")
    if arguments.json is not None:
        arguments.json.parent.mkdir(parents=True, exist_ok=True)
        arguments.json.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n", encoding="utf-8")
        print("Full record:", arguments.json)
    return 0 if passed else 1


def main() -> int:
    arguments = parse_arguments()
    if arguments.recompute_direct:
        return recompute_direct_comparison(arguments)
    require(math.isfinite(arguments.hbar) and arguments.hbar > 0.0,
            "--hbar must be positive and finite")
    require(math.isfinite(arguments.scale) and arguments.scale > 0.0,
            "--Lambda must be positive and finite")
    require(1 <= arguments.levels <= 12, "--levels must lie between 1 and 12")

    rows: list[list[float]] = []
    coordinates: list[list[float]] = []
    for index in range(arguments.levels):
        level_energies = []
        level_coordinates = []
        for order in range(3):
            a_value = solve_flat_coordinate(
                index, arguments.hbar, arguments.scale, order
            )
            level_coordinates.append(a_value)
            level_energies.append(
                matone_energy(
                    a_value, arguments.hbar, arguments.scale, order
                )
            )
        rows.append(level_energies)
        coordinates.append(level_coordinates)

    print("Modified-Mathieu NS capstone")
    print(f"Python {platform.python_version()}; NumPy {np.__version__}; "
          f"SciPy {scipy.__version__}")
    print(
        f"Lambda={arguments.scale:g}; hbar_s={arguments.hbar:g}; "
        f"levels={arguments.levels}"
    )
    print("quantization: d_a F_NS = pi*hbar_s*(n+1/2)")
    print("energy map: E_s=a_s^2/4-(Lambda/4)d_Lambda F_NS^inst")
    print()

    calibrated = (
        abs(arguments.scale - 1.0) <= 1.0e-14
        and abs(arguments.hbar - 1.0) <= 1.0e-14
        and arguments.levels <= len(DCHE_TARGETS)
    )
    if calibrated:
        print(
            " n       perturbative       one instanton      two instantons  "
            "       DCHE target        final gap"
        )
    else:
        print(" n       perturbative       one instanton      two instantons")

    final_gaps: list[float] = []
    for index, energies in enumerate(rows):
        prefix = (
            f"{index:2d}  {energies[0]:18.12f}  {energies[1]:18.12f}  "
            f"{energies[2]:18.12f}"
        )
        if calibrated:
            gap = abs(energies[2] - float(DCHE_TARGETS[index]))
            final_gaps.append(gap)
            print(
                f"{prefix}  {DCHE_TARGETS[index]:18.12f}  {gap:12.3e}"
            )
        else:
            print(prefix)

    print()
    maximum_ns_residual = 0.0
    for index, level_coordinates in enumerate(coordinates):
        target = math.pi * arguments.hbar * (index + 0.5)
        for order, a_value in enumerate(level_coordinates):
            maximum_ns_residual = max(
                maximum_ns_residual,
                abs(
                    free_energy_derivative(
                        a_value, arguments.hbar, arguments.scale, order
                    )
                    - target
                ),
            )
    print(f"maximum NS root residual: {maximum_ns_residual:.3e}")
    if calibrated:
        first_gaps = [abs(row[0] - DCHE_TARGETS[i]) for i, row in enumerate(rows)]
        second_gaps = [abs(row[1] - DCHE_TARGETS[i]) for i, row in enumerate(rows)]
        require(
            all(second_gaps[i] < first_gaps[i] for i in range(len(rows))),
            "one-instanton energies did not improve every calibrated level",
        )
        require(
            all(final_gaps[i] < second_gaps[i] for i in range(len(rows))),
            "two-instanton energies did not improve every calibrated level",
        )
        require(max(final_gaps) <= 6.0e-5,
                "two-instanton/DCHE regression gate failed")
        print(f"maximum two-instanton/DCHE gap: {max(final_gaps):.3e}")
    else:
        print(
            "Calibration warning: this parameter choice has no embedded "
            "DCHE target; the finite instanton truncation has a "
            "parameter-dependent domain of usefulness."
        )
    require(maximum_ns_residual <= 2.0e-11, "NS root residual gate failed")
    print(
        "Scope: finite instanton truncation in Lambda^4; this does not "
        "Borel resum the hbar expansion or bound the omitted instantons."
    )
    print("Regression gates: PASS")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except ValueError as error:
        print(f"error: {error}", file=sys.stderr)
        sys.exit(2)
