#!/usr/bin/env python3
"""Run the official DCF pipeline on the released MATH graphs and audit released results."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import sys
import types
from pathlib import Path

import torch


ROOT = Path(__file__).resolve().parent
SOURCE = ROOT / "source_current"
COMMIT = "0b4d5487a9868a18c4f9aa5b3d96cdccc705ca97"


def load(relative: str) -> object:
    return json.loads((SOURCE / relative).read_text(encoding="utf-8"))


def write(path: Path, payload: object) -> None:
    path.write_text(json.dumps(payload, indent=2, sort_keys=True, allow_nan=False) + "\n")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)

    torch.manual_seed(260420098)
    torch.use_deterministic_algorithms(True)

    # The audited path does not call soft_quantile, but the official module imports
    # torchsort at module import time. This exact-sort shim is therefore an explicit
    # fail-closed import adapter, not a replacement used in any reported result.
    torchsort = types.ModuleType("torchsort")
    torchsort.soft_sort = lambda values, regularization_strength=1e-4: torch.sort(values, dim=-1).values
    sys.modules.setdefault("torchsort", torchsort)
    sys.path.insert(0, str(SOURCE))

    from src.differentiable_conformal_factuality import (  # type: ignore
        compute_nonconformity_score,
        predict,
    )
    from src.models import ForwardScorer, WarmStartLogisticClaimScorer  # type: ignore
    from src.reasonining_graph_dataset import Reasoning_Graph_Dataset  # type: ignore

    data_path = SOURCE / "data/MATH_open_subclaims_with_scores_and_semantic_eval.json"
    dataset = Reasoning_Graph_Dataset(str(data_path), ["frequency-score"])
    examples = len(dataset)
    claims = sum(len(labels) for labels in dataset.y)
    edges = int(sum(int(x["adj"].sum().item()) for x in dataset.x))
    scorer = ForwardScorer(0)
    nonconformity = []
    for x, y in dataset[:]:
        value = compute_nonconformity_score(
            x, y, 0.0, scorer,
            C=6.0, beta_mix=0.5, margin=2.0, temp=0.2, beta=8.0,
            gamma=1.0, lambda_=1.0, violation_mode="exponential", squash_temp=0.8,
        )
        nonconformity.append(float(value.detach().item()))
    if len(nonconformity) != 50 or not all(math.isfinite(v) for v in nonconformity):
        raise RuntimeError("official nonconformity pipeline failed")
    threshold = sorted(nonconformity)[47]
    predicted = predict(
        dataset.x, [0.0] * examples, scorer, torch.tensor(threshold),
        C=6.0, beta_mix=0.5, margin=2.0, temp=0.2, beta=1.0,
        gamma=1.0, cutoff_temp=0.1,
    )
    no_ancestors = [{**x, "ancestors": torch.zeros_like(x["ancestors"])} for x in dataset.x]
    predicted_without_ancestors = predict(
        no_ancestors, [0.0] * examples, scorer, torch.tensor(threshold),
        C=6.0, beta_mix=0.5, margin=2.0, temp=0.2, beta=1.0,
        gamma=1.0, cutoff_temp=0.1,
    )
    ancestor_l1 = sum(
        float((with_graph - without_graph).abs().sum().item())
        for with_graph, without_graph in zip(predicted, predicted_without_ancestors)
    )
    all_predictions = torch.cat(predicted)
    if all_predictions.numel() != claims or not torch.isfinite(all_predictions).all():
        raise RuntimeError("official prediction pipeline failed")
    if ancestor_l1 <= 1.0:
        raise RuntimeError("ancestor destructive control did not change predictions")

    # End-to-end gradient witness through score, risk, keep, coherence, validity,
    # and soft supremum on an actual released MATH reasoning graph.
    learned = WarmStartLogisticClaimScorer(1)
    objective = compute_nonconformity_score(
        dataset.x[0], dataset.y[0], 0.0, learned,
        C=6.0, beta_mix=0.5, margin=2.0, temp=0.2, beta=8.0,
        gamma=1.0, lambda_=1.0, violation_mode="exponential", squash_temp=0.8,
    )
    objective.backward()
    gradient = learned.linear.weight.grad
    gradient_finite = gradient is not None and bool(torch.isfinite(gradient).all())
    gradient_norm = float(gradient.norm().item()) if gradient is not None else 0.0
    if not gradient_finite or gradient_norm <= 0.0:
        raise RuntimeError("end-to-end DCF gradient vanished")

    native = {
        "paper_id": "XfndtVLIub",
        "official_commit": COMMIT,
        "dataset": "released MATH_open_subclaims_with_scores_and_semantic_eval.json",
        "examples": examples,
        "claims": claims,
        "dependency_edges": edges,
        "nonconformity": {
            "evaluations": len(nonconformity),
            "minimum": min(nonconformity),
            "maximum": max(nonconformity),
            "mean": sum(nonconformity) / len(nonconformity),
            "finite": True,
        },
        "prediction": {
            "node_probabilities": int(all_predictions.numel()),
            "minimum": float(all_predictions.min().item()),
            "maximum": float(all_predictions.max().item()),
            "finite": True,
            "calibration_threshold_from_native_nonconformity": threshold,
        },
        "joint_graph_destructive_control": {
            "ancestor_removal_l1_difference": ancestor_l1,
            "changed": True,
        },
        "gradient_witness": {
            "actual_graph_nodes": len(dataset.y[0]),
            "finite": gradient_finite,
            "nonzero_norm": gradient_norm,
        },
        "soft_sort_shim_used_in_reported_path": False,
    }

    math_results = load("results/math_best_results.json")
    felm_results = load("results/felm_best_optimization_results.json")
    confusion = load("results/confusion_matrices_cv/confusion_matrices_results.json")
    math03 = math_results["best_results"]["0.03"]
    felm01 = felm_results["results"]["0.01"]
    math_improvement = (
        (math03["hard_avg_claims_retained_mean"] - math03["baseline_retention"])
        / math03["baseline_retention"] * 100.0
    )
    felm_improvement = (
        (felm01["learned"]["avg_claims_retained"] - felm01["baseline"]["avg_claims_retained"])
        / felm01["baseline"]["avg_claims_retained"] * 100.0
    )
    agreements = []
    for row in confusion["results_table"]:
        recomputed = (row["TP"] + row["TN"]) / row["total"]
        if abs(recomputed - row["agreement"]) > 1e-15:
            raise RuntimeError("confusion matrix arithmetic drift")
        agreements.append(recomputed)
    release = {
        "paper_id": "XfndtVLIub",
        "official_commit": COMMIT,
        "claim1_math": {
            "dcf_retained": math03["hard_avg_claims_retained_mean"],
            "frequency_baseline_retained": math03["baseline_retention"],
            "relative_improvement_percent": math_improvement,
            "dcf_coverage": math03["hard_coverage_mean"],
            "target_coverage": 0.97,
            "coverage_shortfall_percentage_points": (0.97 - math03["hard_coverage_mean"]) * 100.0,
            "literal_verdict": "falsified_as_composite_reliability_claim",
        },
        "claim2_felm": {
            "dcf_retained": felm01["learned"]["avg_claims_retained"],
            "frequency_baseline_retained": felm01["baseline"]["avg_claims_retained"],
            "relative_improvement_percent": felm_improvement,
            "dcf_coverage": felm01["learned"]["coverage"],
            "target_coverage": 0.99,
            "target_met": felm01["learned"]["coverage"] >= 0.99,
        },
        "claim3_calibration": {
            "released_convergence_suites": 5,
            "released_trials_per_suite": 20,
            "native_actual_graph_nonconformity_evaluations": len(nonconformity),
            "theorem_scope": "soft nonconformity-score convergence; quantile recovery is a separate source statement",
        },
        "claim4_prediction": {
            "released_convergence_suites": 2,
            "released_trials_per_suite": 20,
            "native_actual_graph_prediction_nodes": claims,
        },
        "claim5_agreement": {
            "cv_folds": confusion["n_folds"],
            "rows": len(agreements),
            "predictions_per_row": confusion["results_table"][0]["total"],
            "minimum_agreement": min(agreements),
            "maximum_agreement": max(agreements),
            "all_between_90_and_100_percent": all(0.90 <= value <= 1.0 for value in agreements),
        },
        "claim6_joint_pipeline": {
            "actual_graphs": examples,
            "actual_claim_nodes": claims,
            "ancestor_removal_l1_difference": ancestor_l1,
            "end_to_end_gradient_finite_and_nonzero": gradient_finite and gradient_norm > 0.0,
        },
        "all_passed": bool(
            math_improvement >= 141.0
            and math03["hard_coverage_mean"] < 0.97
            and felm_improvement >= 60.0
            and felm01["learned"]["coverage"] >= 0.99
            and min(agreements) >= 0.90
            and ancestor_l1 > 1.0
            and gradient_finite and gradient_norm > 0.0
        ),
    }
    if not release["all_passed"]:
        raise RuntimeError(release)
    write(args.output_dir / "native_pipeline.json", native)
    write(args.output_dir / "official_release_audit.json", release)
    print(json.dumps({
        "status": "pass",
        "examples": examples,
        "claims": claims,
        "ancestor_l1": ancestor_l1,
        "gradient_norm": gradient_norm,
        "math_literal_falsification": True,
        "felm_verified": True,
        "agreement_rows": len(agreements),
    }, sort_keys=True))


if __name__ == "__main__":
    main()