"""Arbitrary-precision symplectic mode audit for AP1-R2c.

Only the real 2x2 midpoint/Yoshida transfer is evaluated in Decimal
arithmetic.  Background interpolation and the canonical frequency formula
remain the already audited Float64 inputs; extra precision is used solely to
prevent loss of the decaying solution and its Wronskian in the tachyonic band.
"""
from __future__ import annotations

from decimal import Decimal, localcontext
import json
from pathlib import Path

import numpy as np

from ap1_r2c_mode_inheritance import inherited_trajectory
from chi_background_closure import (
    ChiParameters,
    PV_J,
    _interp_mid,
    _omega2_at,
    curvature_uv_tail_n,
    frozen_renormalization_constants,
    order0_uv_tail_n,
    physical_shells,
    portal_terms,
)


def _D(value: float) -> Decimal:
    return Decimal.from_float(float(value))


def propagate_endpoint_decimal(
    trajectory,
    k: np.ndarray,
    weights: np.ndarray | None = None,
    digits: int = 60,
    record_history: bool = False,
) -> dict:
    if digits < 32:
        raise ValueError("at least 32 decimal digits required")
    p = ChiParameters()
    tr = trajectory.validated()
    k = np.asarray(k, dtype=float)
    if weights is not None:
        weights = np.asarray(weights, dtype=float)
        if weights.shape != k.shape:
            raise ValueError("weights must match k")
    mass2_0 = float(portal_terms(tr.sigma[0], tr.theta[0], p)[0])
    muR2 = (p.muR_over_Lambda * p.Lambda) ** 2
    omega2_0 = k[None, :] ** 2 * np.exp(-2.0 * tr.N[0]) + mass2_0 + PV_J[:, None] * muR2
    if np.any(omega2_0 <= 0.0):
        raise ValueError("positive initial frequencies required")

    with localcontext() as ctx:
        ctx.prec = digits
        two = Decimal(2)
        ur, ui, vr, vi, target = [], [], [], [], []
        for sector in range(4):
            ur_s, ui_s, vr_s, vi_s, target_s = [], [], [], [], []
            for omega2 in omega2_0[sector]:
                omega = _D(float(omega2)).sqrt()
                u0 = (Decimal(1) / (two * omega)).sqrt()
                ur_s.append(u0); ui_s.append(Decimal(0))
                vr_s.append(Decimal(0)); vi_s.append(-omega * u0)
                target_s.append(-two * (u0 * (-omega * u0)))
            ur.append(ur_s); ui.append(ui_s); vr.append(vr_s); vi.append(vi_s); target.append(target_s)

        def resolved_at(node: int) -> dict[str, float]:
            if weights is None:
                raise ValueError("weights required for source history")
            a = _D(float(np.exp(tr.N[node])))
            a3 = a**3
            H = _D(float(tr.H[node]))
            mass2 = _D(float(portal_terms(tr.sigma[node], tr.theta[node], p)[0]))
            muR2_d = _D(muR2)
            rho_each, pressure_each, chi_each = [], [], []
            for sector in range(4):
                rho_s = Decimal(0); pressure_s = Decimal(0); chi_s = Decimal(0)
                regulator = Decimal(int(PV_J[sector])) * muR2_d
                for mode in range(len(k)):
                    f2 = (ur[sector][mode] ** 2 + ui[sector][mode] ** 2) / a3
                    shift = Decimal("1.5") * H
                    fdot_re = (vr[sector][mode] - shift * ur[sector][mode]) / a3.sqrt()
                    fdot_im = (vi[sector][mode] - shift * ui[sector][mode]) / a3.sqrt()
                    fdot2 = fdot_re**2 + fdot_im**2
                    physical_k2 = _D(float(k[mode] ** 2 * np.exp(-2.0 * tr.N[node])))
                    total_mass2 = mass2 + regulator
                    weight = _D(float(weights[mode]))
                    rho_s += weight * Decimal("0.5") * (fdot2 + (physical_k2 + total_mass2) * f2)
                    pressure_s += weight * (Decimal("0.5") * fdot2 - (physical_k2 / Decimal(6) + Decimal("0.5") * total_mass2) * f2)
                    chi_s += weight * f2
                rho_each.append(rho_s); pressure_each.append(pressure_s); chi_each.append(chi_s)
            pv = (Decimal(1), Decimal(-3), Decimal(3), Decimal(-1))
            return {
                "rho": float(sum(c * value for c, value in zip(pv, rho_each))),
                "pressure": float(sum(c * value for c, value in zip(pv, pressure_each))),
                "chi2": float(sum(c * value for c, value in zip(pv, chi_each))),
                "gross_rho": float(sum(abs(c * value) for c, value in zip(pv, rho_each))),
            }

        history = {key: [] for key in ("rho", "pressure", "chi2", "gross_rho")}
        if record_history:
            initial_source = resolved_at(0)
            for key in history:
                history[key].append(initial_source[key])

        cbrt2 = 2.0 ** (1.0 / 3.0)
        gammas = (1.0 / (2.0 - cbrt2), -cbrt2 / (2.0 - cbrt2), 1.0 / (2.0 - cbrt2))
        max_wronskian = Decimal(0)
        for index in range(len(tr.N) - 1):
            dN = float(tr.N[index + 1] - tr.N[index])
            cursor = float(tr.N[index])
            for gamma in gammas:
                substep = gamma * dN
                midpoint = cursor + 0.5 * substep
                alpha = _D(substep) / (two * _D(_interp_mid(tr.N, tr.H, midpoint)))
                for sector in range(4):
                    omega2_values = _omega2_at(midpoint, tr, k, PV_J[sector] * muR2, p)
                    for mode, omega2_value in enumerate(omega2_values):
                        omega2 = _D(float(omega2_value))
                        denominator = Decimal(1) + alpha * alpha * omega2
                        factor = Decimal(1) - alpha * alpha * omega2
                        old_ur, old_ui = ur[sector][mode], ui[sector][mode]
                        old_vr, old_vi = vr[sector][mode], vi[sector][mode]
                        ur[sector][mode] = (factor * old_ur + two * alpha * old_vr) / denominator
                        ui[sector][mode] = (factor * old_ui + two * alpha * old_vi) / denominator
                        vr[sector][mode] = (-two * alpha * omega2 * old_ur + factor * old_vr) / denominator
                        vi[sector][mode] = (-two * alpha * omega2 * old_ui + factor * old_vi) / denominator
                cursor += substep
            for sector in range(4):
                for mode in range(len(k)):
                    observed = -two * (ur[sector][mode] * vi[sector][mode] - ui[sector][mode] * vr[sector][mode])
                    error = abs(observed / target[sector][mode] - Decimal(1))
                    max_wronskian = max(max_wronskian, error)
            if record_history:
                source = resolved_at(index + 1)
                for key in history:
                    history[key].append(source[key])

        endpoint = []
        for sector in range(4):
            endpoint.append([float((ur[sector][mode] ** 2 + ui[sector][mode] ** 2).sqrt()) for mode in range(len(k))])
        result = {
            "digits": digits,
            "wronskian_relative_error": float(max_wronskian),
            "physical_endpoint_abs_u": endpoint[0],
            "all_sector_endpoint_abs_u": endpoint,
        }
        if weights is not None:
            a = _D(float(np.exp(tr.N[-1])))
            a3 = a**3
            H = _D(float(tr.H[-1]))
            mass2 = _D(float(portal_terms(tr.sigma[-1], tr.theta[-1], p)[0]))
            muR2_d = _D(muR2)
            rho_each, pressure_each, chi_each = [], [], []
            for sector in range(4):
                rho_s = Decimal(0); pressure_s = Decimal(0); chi_s = Decimal(0)
                regulator = Decimal(int(PV_J[sector])) * muR2_d
                for mode in range(len(k)):
                    f2 = (ur[sector][mode] ** 2 + ui[sector][mode] ** 2) / a3
                    shift = Decimal("1.5") * H
                    fdot_re = (vr[sector][mode] - shift * ur[sector][mode]) / a3.sqrt()
                    fdot_im = (vi[sector][mode] - shift * ui[sector][mode]) / a3.sqrt()
                    fdot2 = fdot_re**2 + fdot_im**2
                    physical_k2 = _D(float(k[mode] ** 2 * np.exp(-2.0 * tr.N[-1])))
                    total_mass2 = mass2 + regulator
                    weight = _D(float(weights[mode]))
                    rho_s += weight * Decimal("0.5") * (fdot2 + (physical_k2 + total_mass2) * f2)
                    pressure_s += weight * (Decimal("0.5") * fdot2 - (physical_k2 / Decimal(6) + Decimal("0.5") * total_mass2) * f2)
                    chi_s += weight * f2
                rho_each.append(rho_s); pressure_each.append(pressure_s); chi_each.append(chi_s)
            pv = (Decimal(1), Decimal(-3), Decimal(3), Decimal(-1))
            result["resolved_pv_endpoint"] = {
                "rho": float(sum(c * value for c, value in zip(pv, rho_each))),
                "pressure": float(sum(c * value for c, value in zip(pv, pressure_each))),
                "chi2": float(sum(c * value for c, value in zip(pv, chi_each))),
                "gross_rho": float(sum(abs(c * value) for c, value in zip(pv, rho_each))),
            }
        if record_history:
            result["resolved_pv_history"] = history
        return result


def completed_history_decimal(
    trajectory,
    k: np.ndarray,
    weights: np.ndarray,
    digits: int = 60,
) -> dict:
    """Return the inherited renormalized source history on one trajectory."""

    p = ChiParameters()
    tr = trajectory.validated()
    result = propagate_endpoint_decimal(
        tr, k, weights, digits, record_history=True
    )
    resolved = result["resolved_pv_history"]
    tail0 = order0_uv_tail_n(tr, p)
    tail24 = curvature_uv_tail_n(tr, p, nodes_per_octave=6)
    constants = frozen_renormalization_constants(p)
    mass2 = np.asarray(portal_terms(tr.sigma, tr.theta, p)[0], dtype=float)
    C = constants["C_chi2"]
    A = constants["A_g"]
    B = constants["B_G_rhs"]
    result["completed_history"] = {
        "rho": np.asarray(resolved["rho"]) + tail0["rho"] + tail24["rho"] - 0.5 * C * mass2 + A + 3.0 * B * tr.H**2,
        "pressure": np.asarray(resolved["pressure"]) + tail0["pressure"] + tail24["pressure"] + 0.5 * C * mass2 - A - B * (2.0 * tr.Hdot + 3.0 * tr.H**2),
        "chi2": np.asarray(resolved["chi2"]) + tail0["chi2"] + tail24["chi2"] - C,
    }
    return result


def build_high_precision_report(
    state_path: Path,
    levels: tuple[int, ...] = (1025, 2049, 4097),
    digits: int = 60,
) -> dict:
    p = ChiParameters()
    k, weights = physical_shells(p.Lambda, nodes=8)
    runs = []
    endpoints = []
    history_records = []
    for nodes in levels:
        tr = inherited_trajectory(state_path, nodes)
        result = completed_history_decimal(tr, k, weights, digits)
        endpoint = np.asarray(result["physical_endpoint_abs_u"], dtype=float)
        endpoints.append(endpoint)
        completed_history = result["completed_history"]
        completed = {key: float(value[-1]) for key, value in completed_history.items()}
        history_records.append((tr.N.copy(), completed_history))
        runs.append({
            "nodes": nodes,
            "digits": digits,
            "wronskian_relative_error": result["wronskian_relative_error"],
            "physical_endpoint_abs_u_min": float(np.min(endpoint)),
            "physical_endpoint_abs_u_max": float(np.max(endpoint)),
            "resolved_pv_endpoint": result["resolved_pv_endpoint"],
            "completed_renormalized_endpoint": completed,
        })
    changes = [
        float(np.max(np.abs(left - right) / np.maximum(np.abs(right), 1.0e-300)))
        for left, right in zip(endpoints[:-1], endpoints[1:])
    ]
    source_changes = {}
    for key in ("rho", "pressure", "chi2"):
        values = [run["resolved_pv_endpoint"][key] for run in runs]
        source_changes[key] = [
            abs(left - right) / max(abs(right), 1.0e-300)
            for left, right in zip(values[:-1], values[1:])
        ]
    completed_changes = {}
    for key in ("rho", "pressure", "chi2"):
        values = [run["completed_renormalized_endpoint"][key] for run in runs]
        completed_changes[key] = [
            abs(left - right) / max(abs(right), 1.0e-300)
            for left, right in zip(values[:-1], values[1:])
        ]
    history_changes = {key: [] for key in ("rho", "pressure", "chi2")}
    for coarse, fine in zip(history_records[:-1], history_records[1:]):
        coarse_N, coarse_history = coarse
        fine_N, fine_history = fine
        for key in history_changes:
            coarse_values = np.asarray(coarse_history[key])
            fine_values = np.interp(coarse_N, fine_N, np.asarray(fine_history[key]))
            scale = max(float(np.max(np.abs(fine_values))), 1.0e-300)
            history_changes[key].append(float(np.max(np.abs(coarse_values - fine_values)) / scale))
    gates = {
        "wronskian_below_1e_11": bool(max(run["wronskian_relative_error"] for run in runs) < 1.0e-11),
        "endpoint_refines_monotonically": bool(len(changes) < 2 or changes[-1] < changes[-2]),
        "finite": bool(all(
            np.isfinite(run["wronskian_relative_error"])
            and np.isfinite(run["physical_endpoint_abs_u_min"])
            and np.isfinite(run["physical_endpoint_abs_u_max"])
            and all(np.isfinite(value) for value in run["resolved_pv_endpoint"].values())
            and all(np.isfinite(value) for value in run["completed_renormalized_endpoint"].values())
            for run in runs
        )),
        "resolved_source_refines_monotonically": bool(all(
            len(values) < 2 or values[-1] < values[-2]
            for values in source_changes.values()
        )),
        "completed_endpoint_refines_monotonically": bool(all(
            len(values) < 2 or values[-1] < values[-2]
            for values in completed_changes.values()
        )),
        "completed_history_refines_monotonically": bool(all(
            len(values) < 2 or values[-1] < values[-2]
            for values in history_changes.values()
        )),
    }
    return {
        "schema": "apeiron-ap1-r2c-high-precision-modes-v1.0",
        "classification": "HIGH_PRECISION_INHERITANCE_AND_RENORMALIZED_HISTORY_PASS_CANDIDATE_WARD_PENDING" if all(gates.values()) else "HIGH_PRECISION_INHERITANCE_OR_SOURCE_HISTORY_NOT_PASS",
        "method": "Decimal real symplectic transfer and resolved PV endpoint; canonical UV tails and frozen finite counterterms; audited Float64 background coefficients; no endpoint vacuum reset",
        "old_solver_or_physical_map_called": False,
        "runs": runs,
        "successive_endpoint_max_relative_changes": changes,
        "successive_resolved_source_relative_changes": source_changes,
        "successive_completed_endpoint_relative_changes": completed_changes,
        "successive_completed_history_max_relative_changes": history_changes,
        "gates": gates,
        "claim_boundary": "inherited renormalized source-history gate only; Ward identity on a new self-consistent candidate and observables remain unreleased",
    }


def main() -> None:
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("state", type=Path)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--digits", type=int, default=60)
    args = parser.parse_args()
    rendered = json.dumps(build_high_precision_report(args.state, digits=args.digits), indent=2) + "\n"
    if args.output:
        args.output.write_text(rendered, encoding="utf-8")
    else:
        print(rendered, end="")


if __name__ == "__main__":
    main()