"""Tolerance freeze and 96-shell held-out for the exact AP1-M1 junction.

The eligible junction candidate is the exact stored HARD-PASS index zero.  No
interpolated root is eligible.  Time-grid and local curvature-tail budgets are
consumed by the 4/2/1 pilot; only the 96-shell momentum quadrature is held out.
Every acceptance cap is frozen before that rule is evaluated once.
"""
from __future__ import annotations

from dataclasses import asdict
from datetime import datetime, timezone
import json
from pathlib import Path
from typing import Any

import numpy as np

from ap1_m1_background_preflight import decode_hard_pass_state
from ap1_m1_junction_convergence import (
    JunctionPilotConfig,
    complete_frozen_trajectory,
    evaluate_prefix_level,
    initial_probe_delta,
)
from ap1_m1_junction_validation import file_sha256, load_json
from ap1_r2c_self_consistent_candidate import classical_terms
from chi_background_closure import ChiParameters, curvature_uv_tail_n


EXISTING_WRONSKIAN_CAP = 1.0e-11
EXISTING_FRIEDMANN_CAP = 1.0e-6
RAYCHAUDHURI_PILOT_CEILING = 4.0e-8

METRIC_KEYS = (
    "abs_delta_N_physical",
    "relative_delta_required_standard_density",
    "abs_delta_quantum_rho_over_H2",
    "abs_delta_quantum_pressure_over_H2",
    "abs_delta_quantum_chi2_over_scale",
    "abs_delta_friedmann_normalized",
    "abs_delta_raychaudhuri_normalized",
)

# Arithmetic budgets prevent a bitwise-zero pilot delta from becoming an
# impossible zero-tolerance held-out gate.  They are declared before the
# 96-shell value is evaluated.
REPRESENTATION_FLOORS = {
    "abs_delta_N_physical": 1.0e-12,
    "relative_delta_required_standard_density": 1.0e-12,
    "abs_delta_quantum_rho_over_H2": 1.0e-12,
    "abs_delta_quantum_pressure_over_H2": 1.0e-12,
    "abs_delta_quantum_chi2_over_scale": 2.5e-15,
    "abs_delta_friedmann_normalized": 1.0e-20,
    "abs_delta_raychaudhuri_normalized": 1.0e-12,
}


class JunctionHeldoutError(RuntimeError):
    """The exact-state held-out chain is incomplete or has changed."""


def _finite_metrics(values: dict[str, Any]) -> bool:
    return all(
        key in values and np.isfinite(float(values[key])) for key in METRIC_KEYS
    )


def heldout_caps(last_momentum_delta: dict[str, Any]) -> dict[str, float]:
    if not _finite_metrics(last_momentum_delta):
        raise JunctionHeldoutError("last momentum budget is incomplete")
    return {
        key: max(float(last_momentum_delta[key]), REPRESENTATION_FLOORS[key])
        for key in METRIC_KEYS
    }


def delta_gates(
    observed: dict[str, Any], caps: dict[str, Any]
) -> dict[str, bool]:
    if not _finite_metrics(observed):
        return {f"heldout_{key}_within_cap": False for key in METRIC_KEYS}
    return {
        f"heldout_{key}_within_cap": abs(float(observed[key]))
        <= float(caps[key])
        for key in METRIC_KEYS
    }


def build_tolerance_manifest(pilot_path: Path) -> dict[str, Any]:
    pilot = load_json(pilot_path)
    expected_config = json.loads(json.dumps(asdict(JunctionPilotConfig())))
    if pilot.get("schema") != "apeiron-ap1-m1-junction-convergence-pilot-v1.1":
        raise JunctionHeldoutError("exact-initial junction pilot v1.1 required")
    if pilot.get("config") != expected_config:
        raise JunctionHeldoutError("junction pilot configuration changed")
    expected_rule = (
        "use exact stored HARD-PASS index 0; determine the physical-N offset "
        "by Friedmann inversion; test Raychaudhuri independently against a "
        "pre-frozen pilot-derived tolerance"
    )
    if pilot.get("junction_rule", {}).get("rule") != expected_rule:
        raise JunctionHeldoutError("exact stored-state rule changed")

    momentum_deltas = pilot["momentum_pilot"][
        "successive_exact_initial_probe_deltas"
    ]
    curvature_deltas = pilot["curvature_tail_time_pilot"][
        "successive_exact_initial_probe_deltas"
    ]
    if len(momentum_deltas) != 3 or len(curvature_deltas) != 2:
        raise JunctionHeldoutError("pilot budget chain is incomplete")
    momentum_caps = heldout_caps(momentum_deltas[-1])
    curvature_budget = {
        key: float(curvature_deltas[-1][key]) for key in METRIC_KEYS
    }
    fine_probe = pilot["momentum_pilot"]["levels"][-1][
        "exact_initial_HARD_PASS_junction_probe"
    ]
    roundoff = pilot["roundoff_pilot"]["exact_initial_probe_delta"]
    gates = {
        "pilot_structural_gates_pass": bool(
            pilot.get("all_structural_pilot_gates_pass")
        ),
        "exact_stored_index_zero_rule_frozen": True,
        "curvature_tail_4_2_1_budget_complete": all(
            _finite_metrics(value) for value in curvature_deltas
        ),
        "momentum_24_32_48_64_budget_complete": all(
            _finite_metrics(value) for value in momentum_deltas
        ),
        "roundoff_60_80_bitwise_stable": all(
            float(roundoff[key]) == 0.0 for key in METRIC_KEYS
        ),
        "heldout_96_not_evaluated": (
            pilot.get("held_out", {}).get("momentum_nodes") == 96
            and not bool(pilot.get("held_out", {}).get("evaluated"))
        ),
        "candidate_not_checkpoint_or_seed": (
            not bool(pilot.get("checkpoint_released"))
            and not bool(pilot.get("seed_released"))
        ),
        "fine_probe_below_preregistered_ray_ceiling": abs(
            float(fine_probe["raychaudhuri_normalized"])
        )
        < RAYCHAUDHURI_PILOT_CEILING,
        "positive_density_and_frequency_at_exact_state": (
            float(fine_probe["required_standard_density_Mpl4"]) > 0.0
            and bool(fine_probe["positive_vacuum_frequency_margin"])
        ),
        "background_and_kernel_not_started": (
            not bool(pilot.get("new_AP1_M1_background_started"))
            and not bool(pilot.get("physical_response_kernel_started"))
        ),
    }
    passed = bool(all(gates.values()))
    return {
        "schema": "apeiron-ap1-m1-exact-initial-junction-tolerances-v1.0",
        "updated_utc": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "classification": (
            "M1_EXACT_INITIAL_JUNCTION_TOLERANCES_FROZEN_96_SHELL_HELDOUT_UNSEEN"
            if passed
            else "M1_EXACT_INITIAL_JUNCTION_TOLERANCE_FREEZE_NOT_PASS"
        ),
        "input_sha256": {pilot_path.name: file_sha256(pilot_path)},
        "junction_rule": pilot["junction_rule"],
        "exploratory_root_scan_role": pilot["exploratory_root_scan_role"],
        "cap_derivation": (
            "The 96-shell held-out must remain within the final 48->64 shell "
            "delta or an independently declared Float64/solver representation "
            "floor, whichever is larger. The Raychaudhuri ceiling 4e-8 was "
            "rounded upward from the visible pilot value 3.6256e-8 and frozen "
            "before the 96-shell evaluation; it is also far below the unchanged "
            "8e-5 background residual cap."
        ),
        "thresholds": {
            "heldout_momentum_nodes": 96,
            "evaluation_stride": 2,
            "prefix_intervals": 256,
            "decimal_digits": 80,
            "heldout_delta_caps": momentum_caps,
            "representation_floors": REPRESENTATION_FLOORS,
            "max_abs_raychaudhuri_normalized": RAYCHAUDHURI_PILOT_CEILING,
            "max_abs_friedmann_normalized": EXISTING_FRIEDMANN_CAP,
            "max_wronskian_relative_error": EXISTING_WRONSKIAN_CAP,
            "require_exact_stored_index": 0,
            "require_positive_standard_density": True,
            "require_positive_inherited_frequency_margin": True,
        },
        "curvature_tail_discretization_budget_from_2_to_1": curvature_budget,
        "gates": gates,
        "all_tolerance_freeze_gates_pass": passed,
        "tolerances_frozen_before_96_shell_heldout": passed,
        "heldout_evaluated": False,
        "checkpoint_released": False,
        "seed_released": False,
        "new_AP1_M1_background_started": False,
        "physical_response_kernel_started": False,
        "next_required": (
            "evaluate the 96-shell exact index-zero probe once; do not change "
            "the state rule, budgets, floors or residual ceilings afterwards"
        ),
        "claim_boundary": (
            "junction tolerance freeze only; no held-out result, accepted M1 "
            "seed, background, response kernel or observable"
        ),
    }


def build_heldout_report(
    apeiron_root: Path,
    pilot_path: Path,
    manifest_path: Path,
) -> dict[str, Any]:
    pilot = load_json(pilot_path)
    manifest = load_json(manifest_path)
    if not manifest.get("all_tolerance_freeze_gates_pass") or not manifest.get(
        "tolerances_frozen_before_96_shell_heldout"
    ):
        raise JunctionHeldoutError("96-shell tolerance manifest is open")
    if manifest["input_sha256"].get(pilot_path.name) != file_sha256(pilot_path):
        raise JunctionHeldoutError("pilot changed after the tolerance freeze")
    thresholds = manifest["thresholds"]
    if int(thresholds["heldout_momentum_nodes"]) != 96:
        raise JunctionHeldoutError("held-out momentum rule changed")

    state_path = apeiron_root / "HANDOFF/CURRENT_STATE_LATEST.npz"
    frozen = decode_hard_pass_state(state_path)
    full_trajectory = complete_frozen_trajectory(frozen)
    full_curvature_tail = curvature_uv_tail_n(
        full_trajectory, ChiParameters(), nodes_per_octave=6
    )
    heldout = evaluate_prefix_level(
        frozen,
        int(thresholds["prefix_intervals"]),
        int(thresholds["evaluation_stride"]),
        int(thresholds["heldout_momentum_nodes"]),
        int(thresholds["decimal_digits"]),
        full_curvature_tail,
    )
    baseline = pilot["momentum_pilot"]["levels"][-1]
    delta = initial_probe_delta(baseline, heldout)
    probe = heldout["exact_initial_HARD_PASS_junction_probe"]
    state = frozen[0]
    kinetic = classical_terms(state[:5])
    gates = {
        **delta_gates(delta, thresholds["heldout_delta_caps"]),
        "heldout_level_finite": bool(heldout["finite"]),
        "exact_stored_index_zero_used": (
            int(probe["stored_index"])
            == int(thresholds["require_exact_stored_index"])
            and float(probe["N_local"]) == float(state[5]) == 0.0
        ),
        "friedmann_below_frozen_cap": abs(float(probe["friedmann_normalized"]))
        <= float(thresholds["max_abs_friedmann_normalized"]),
        "raychaudhuri_below_frozen_cap": abs(
            float(probe["raychaudhuri_normalized"])
        )
        <= float(thresholds["max_abs_raychaudhuri_normalized"]),
        "positive_standard_density": float(
            probe["required_standard_density_Mpl4"]
        )
        > 0.0,
        "positive_inherited_frequency_margin": bool(
            probe["positive_vacuum_frequency_margin"]
        )
        and float(probe["vacuum_frequency_margin_Mpl2"]) > 0.0,
        "positive_H_PX_K": (
            float(state[4]) > 0.0
            and float(kinetic["PX"]) > 0.0
            and float(kinetic["K"]) > 0.0
        ),
        "wronskian_below_frozen_cap": float(heldout["wronskian_relative_error"])
        <= float(thresholds["max_wronskian_relative_error"]),
        "old_solver_or_physical_map_not_called": True,
        "no_interpolated_root_used": True,
    }
    passed = bool(all(gates.values()))
    seed = None
    if passed:
        seed = {
            "stored_HARD_PASS_index": 0,
            "N_local": float(state[5]),
            "N_physical": float(probe["N_physical"]),
            "N_offset_physical_minus_local": float(
                probe["N_offset_physical_minus_local"]
            ),
            "sigma": float(state[0]),
            "sigmadot": float(state[1]),
            "theta": float(state[2]),
            "thetadot": float(state[3]),
            "H_Mpl": float(state[4]),
            "Hdot_Mpl2": float(full_trajectory.Hdot[0]),
            "rho_q_Mpl4": float(probe["quantum_rho_Mpl4"]),
            "p_q_Mpl4": float(probe["quantum_pressure_Mpl4"]),
            "chi2_Mpl2": float(probe["quantum_chi2_Mpl2"]),
            "required_standard_density_Mpl4": float(
                probe["required_standard_density_Mpl4"]
            ),
            "standard_pressure_Mpl4": float(probe["standard_pressure_Mpl4"]),
            "friedmann_normalized": float(probe["friedmann_normalized"]),
            "raychaudhuri_normalized": float(probe["raychaudhuri_normalized"]),
            "vacuum_frequency_margin_Mpl2": float(
                probe["vacuum_frequency_margin_Mpl2"]
            ),
            "PX": float(kinetic["PX"]),
            "K": float(kinetic["K"]),
        }

    return {
        "schema": "apeiron-ap1-m1-exact-initial-junction-heldout-v1.0",
        "updated_utc": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "classification": (
            "M1_EXACT_INITIAL_JUNCTION_96_SHELL_PASS_SEED_RELEASED_BACKGROUND_SOLVE_OPEN"
            if passed
            else "M1_EXACT_INITIAL_JUNCTION_96_SHELL_NOT_PASS_NO_SEED"
        ),
        "input_sha256": {
            pilot_path.name: file_sha256(pilot_path),
            manifest_path.name: file_sha256(manifest_path),
            "HANDOFF/CURRENT_STATE_LATEST.npz": file_sha256(state_path),
            "AP1/CODE/ap1_m1_junction_heldout.py": file_sha256(Path(__file__)),
        },
        "heldout_was_unseen_when_tolerances_frozen": True,
        "junction_rule": manifest["junction_rule"],
        "frozen_thresholds": thresholds,
        "curvature_tail_discretization_budget": manifest[
            "curvature_tail_discretization_budget_from_2_to_1"
        ],
        "heldout_level": heldout,
        "heldout_delta_vs_64_shell_pilot": delta,
        "gates": gates,
        "all_exact_initial_junction_gates_pass": passed,
        "checkpoint_released": passed,
        "seed_released": passed,
        "accepted_seed": seed,
        "seed_use_scope": "new AP1-M1 background solver only" if passed else None,
        "mode_state_reconstruction": (
            {
                "state": "same inherited adiabatic vacuum at exact frozen index zero",
                "momentum_shells": 96,
                "decimal_digits": 80,
                "mode_reinitialization_at_later_time": False,
                "reproducible_from_hash_bound_state_and_frozen_algorithm": True,
            }
            if passed
            else None
        ),
        "new_AP1_M1_background_started": False,
        "physical_response_kernel_started": False,
        "nonpass_stored_or_used_as_seed": False,
        "equations_changed": False,
        "physics_changed": False,
        "parameters_changed": False,
        "existing_gate_thresholds_changed": False,
        "next_required": (
            "initialize the separate N-parametrized AP1-M1 background solver "
            "from this exact hash-bound seed and run a bounded non-production pilot"
            if passed
            else "improve numerical resolution without changing the exact-state rule or frozen gates"
        ),
        "claim_boundary": (
            "exact early M1 junction seed only; no late-time background, physical "
            "response kernel, observable, fit or significance"
        ),
    }


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)
    freeze_parser = subparsers.add_parser("freeze")
    freeze_parser.add_argument("pilot", type=Path)
    freeze_parser.add_argument("--output", type=Path, required=True)
    heldout_parser = subparsers.add_parser("heldout")
    heldout_parser.add_argument("apeiron_root", type=Path)
    heldout_parser.add_argument("pilot", type=Path)
    heldout_parser.add_argument("manifest", type=Path)
    heldout_parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    report = (
        build_tolerance_manifest(args.pilot)
        if args.command == "freeze"
        else build_heldout_report(args.apeiron_root, args.pilot, args.manifest)
    )
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()
