#!/usr/bin/env python3
"""Calibrate outgoing-span confidence risk on three Extended-251 prompts."""

from __future__ import annotations

import argparse
import json
import math
import os
import sys
import tempfile
from pathlib import Path
from typing import Any


def preparse_gpu() -> str:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--gpu", default="5")
    args, _ = parser.parse_known_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
    os.environ.setdefault(
        "MPLCONFIGDIR", tempfile.mkdtemp(prefix="self_forcing_span_calibration_")
    )
    return str(args.gpu)


PHYSICAL_GPU = preparse_gpu()

import torch
from omegaconf import OmegaConf
from safetensors.torch import load_file

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from predictor_training.confidence import PredictorConfidenceHead
from scripts import evaluate_single_block_fppf as base
from scripts import generate_vbench8_extended_strategies as extended
from utils.wan_wrapper import WanTextEncoder


MAPPING_DEFAULT = REPO_ROOT / "assets/vbench8_extended_subset_mapping.json"
EXPERIMENT_DEFAULT = (
    REPO_ROOT / "confidence_experiments/layer17_stage1_step2000_20260831"
)
OUTPUT_DEFAULT = (
    REPO_ROOT
    / "confidence_experiments/layer17_stage1_step2000_spanrisk_vbench3_20260901"
)
BETAS = (0.0, 0.5, 1.0, 1.5, 2.0)
TARGETS = {
    "step12": (6, 8, 10),
    "step123": (6, 9, 12, 15),
}
CANDIDATE_STEPS = {
    "step12": (1, 2),
    "step123": (1, 2, 3),
}


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n",
        encoding="utf-8",
    )
    os.replace(temporary, path)


def quantile(values: list[float], fraction: float) -> float:
    if not values:
        raise ValueError("Cannot take a quantile of an empty list")
    ordered = sorted(values)
    position = fraction * (len(ordered) - 1)
    lower = int(math.floor(position))
    upper = int(math.ceil(position))
    if lower == upper:
        return ordered[lower]
    weight = position - lower
    return ordered[lower] * (1.0 - weight) + ordered[upper] * weight


def beta_name(beta: float) -> str:
    return str(beta).replace(".", "p")


def read_mapping_row(path: Path, global_index: int) -> dict[str, Any]:
    rows = extended.read_mapping(path)
    matches = [row for row in rows if int(row["global_index"]) == global_index]
    if len(matches) != 1:
        raise ValueError(f"Expected one mapping row for global index {global_index}")
    return matches[0]


def load_models(
    experiment_root: Path, device: torch.device
) -> tuple[Any, Any, Any, Any, Any]:
    config = OmegaConf.merge(
        OmegaConf.load(REPO_ROOT / "configs/default_config.yaml"),
        OmegaConf.load(REPO_ROOT / "configs/self_forcing_sid.yaml"),
    )
    pipeline = base.build_pipeline(
        config,
        REPO_ROOT / "checkpoints/self_forcing_dmd.pt",
        torch.nn.Identity(),
        device,
    )
    predictor = base.load_predictor(
        pipeline.generator.model,
        {
            "source_layer": 17,
            "weights": REPO_ROOT
            / "training_runs/layer17_stage1_1000p_4gpu_b16_2000steps"
            / "checkpoint_step_2000/predictor.safetensors",
            "gate_mode": "baseline",
        },
        device,
    )
    text_encoder = WanTextEncoder().to(device=device, dtype=torch.bfloat16).eval()
    text_encoder.requires_grad_(False)
    step12_head = PredictorConfidenceHead(num_steps=2).to(device=device).eval()
    step12_head.load_state_dict(
        load_file(
            str(experiment_root / "confidence_step12/confidence_best.safetensors"),
            device="cpu",
        ),
        strict=True,
    )
    step12_head.requires_grad_(False)
    step123_head = PredictorConfidenceHead(num_steps=3).to(device=device).eval()
    step123_head.load_state_dict(
        load_file(
            str(experiment_root / "confidence_step123/confidence_best.safetensors"),
            device="cpu",
        ),
        strict=True,
    )
    step123_head.requires_grad_(False)
    return pipeline, predictor, text_encoder, step12_head, step123_head


def probe(args: argparse.Namespace) -> None:
    if args.global_index is None:
        raise ValueError("--global-index is required in probe mode")
    row = read_mapping_row(args.mapping, args.global_index)
    destination = args.output_root / "probes" / f"global_{args.global_index:04d}.json"
    if destination.is_file() and not args.overwrite:
        print(f"[cached] {destination}", flush=True)
        return
    device = torch.device("cuda")
    torch.set_grad_enabled(False)
    pipeline, predictor, text_encoder, step12_head, step123_head = load_models(
        args.experiment_root, device
    )
    prompt = str(row["extended_prompt"])
    conditional = text_encoder(text_prompts=[prompt])
    probes: dict[str, Any] = {}
    for family, head in (("step12", step12_head), ("step123", step123_head)):
        config = {
            "name": f"{family}_span_probe",
            "policy": "dynamic",
            "candidate_steps": list(CANDIDATE_STEPS[family]),
            "beta": 0.0,
            "threshold": -math.inf,
            "target_accepts": 0,
            "head": family,
            "risk_mode": "outgoing_span",
        }
        _, diagnostic = extended.generate_rollout(
            pipeline=pipeline,
            conditional_dict=conditional,
            seed=args.seed,
            device=device,
            predictor=predictor,
            head=head,
            config=config,
        )
        decisions = diagnostic.pop("decisions")
        if any(decision["accepted"] for decision in decisions):
            raise RuntimeError("Probe unexpectedly accepted a Predictor result")
        probes[family] = {"generation": diagnostic, "decisions": decisions}
        print(
            f"[probe] global={args.global_index} family={family} "
            f"decisions={len(decisions)}",
            flush=True,
        )
    atomic_json(
        destination,
        {
            "status": "complete",
            "physical_gpu": PHYSICAL_GPU,
            "global_index": int(row["global_index"]),
            "prompt_suite": row["prompt_suite"],
            "suite_index": int(row["suite_index"]),
            "prompt": prompt,
            "seed": args.seed,
            "probes": probes,
        },
    )


def prepare(args: argparse.Namespace) -> None:
    paths = sorted((args.output_root / "probes").glob("global_*.json"))
    if len(paths) != 3:
        raise ValueError(f"Expected exactly three probe files, found {len(paths)}")
    probes = [json.loads(path.read_text(encoding="utf-8")) for path in paths]
    configurations: list[dict[str, Any]] = []
    span_values: dict[int, set[float]] = {1: set(), 2: set(), 3: set()}
    for family in ("step12", "step123"):
        decisions = [
            decision
            for probe_record in probes
            for decision in probe_record["probes"][family]["decisions"]
        ]
        expected = 3 * 6 * len(CANDIDATE_STEPS[family])
        if len(decisions) != expected:
            raise ValueError(
                f"Expected {expected} {family} decisions, got {len(decisions)}"
            )
        for decision in decisions:
            span_values[int(decision["step"])].add(float(decision["outgoing_span"]))
        max_calls = 6 * len(CANDIDATE_STEPS[family])
        for beta in BETAS:
            risks = [
                float(decision["predicted_local_error"])
                * float(decision["span_weight"])
                * (1.0 + beta * float(decision["chunk_alpha"]))
                for decision in decisions
            ]
            for target in TARGETS[family]:
                threshold = quantile(risks, target / max_calls)
                configurations.append(
                    {
                        "name": f"{family}_span_b{beta_name(beta)}_k{target:02d}",
                        "policy": "dynamic",
                        "candidate_steps": list(CANDIDATE_STEPS[family]),
                        "beta": beta,
                        "threshold": threshold,
                        "target_accepts": target,
                        "head": family,
                        "risk_mode": "outgoing_span",
                    }
                )
    actual_spans = {
        str(step): sorted(values) for step, values in span_values.items() if values
    }
    expected_spans = {1: 104.1666, 2: 208.3333, 3: 625.0}
    for step, expected in expected_spans.items():
        values = span_values[step]
        if values and any(abs(value - expected) > 0.02 for value in values):
            raise ValueError(f"Unexpected outgoing span at step {step}: {values}")
    atomic_json(
        args.output_root / "scan_configs.json",
        {
            "status": "ready",
            "risk_definition": (
                "predicted_local_error * (outgoing_span / 1000) * "
                "(1 + beta * chunk_alpha)"
            ),
            "probe_global_indices": [int(value["global_index"]) for value in probes],
            "outgoing_spans": actual_spans,
            "betas": list(BETAS),
            "targets": {key: list(value) for key, value in TARGETS.items()},
            "num_configs": len(configurations),
            "configs": configurations,
        },
    )
    print(f"[prepare] configurations={len(configurations)}", flush=True)


def prepare_threshold_only(args: argparse.Namespace) -> None:
    probe_root = args.probe_root or args.output_root
    paths = sorted((probe_root / "probes").glob("global_*.json"))
    if len(paths) != 3:
        raise ValueError(f"Expected exactly three probe files, found {len(paths)}")
    probes = [json.loads(path.read_text(encoding="utf-8")) for path in paths]
    configurations: list[dict[str, Any]] = []
    span_values: dict[int, set[float]] = {1: set(), 2: set(), 3: set()}
    beta = 2.0
    for family in ("step12", "step123"):
        decisions = [
            decision
            for probe_record in probes
            for decision in probe_record["probes"][family]["decisions"]
        ]
        expected = 3 * 6 * len(CANDIDATE_STEPS[family])
        if len(decisions) != expected:
            raise ValueError(
                f"Expected {expected} {family} decisions, got {len(decisions)}"
            )
        for decision in decisions:
            span_values[int(decision["step"])].add(float(decision["outgoing_span"]))
        max_calls = 6 * len(CANDIDATE_STEPS[family])
        risks = [
            float(decision["predicted_local_error"])
            * float(decision["span_weight"])
            * (1.0 + beta * float(decision["chunk_alpha"]))
            for decision in decisions
        ]
        for rank in range(1, max_calls):
            configurations.append(
                {
                    "name": f"{family}_span_beta2_thr{rank:02d}",
                    "policy": "dynamic",
                    "candidate_steps": list(CANDIDATE_STEPS[family]),
                    "beta": beta,
                    "threshold": quantile(risks, rank / max_calls),
                    "target_accepts": rank,
                    "threshold_rank": rank,
                    "head": family,
                    "risk_mode": "outgoing_span",
                }
            )
    atomic_json(
        args.output_root / "scan_configs.json",
        {
            "status": "ready",
            "risk_definition": (
                "predicted_local_error * (outgoing_span / 1000) * "
                "(1 + 2.0 * chunk_alpha)"
            ),
            "beta": beta,
            "probe_global_indices": [int(value["global_index"]) for value in probes],
            "outgoing_spans": {
                str(step): sorted(values)
                for step, values in span_values.items()
                if values
            },
            "threshold_search": "quantile rank 1..max_candidate_calls-1",
            "targets": {key: list(value) for key, value in TARGETS.items()},
            "num_configs": len(configurations),
            "configs": configurations,
        },
    )
    print(
        f"[prepare-threshold-only] beta={beta} configurations={len(configurations)}",
        flush=True,
    )


def latent_nrmse(predicted: torch.Tensor, reference: torch.Tensor) -> float:
    numerator = (predicted.float() - reference.float()).square().sum()
    denominator = reference.float().square().sum().clamp_min(1e-8)
    return float(torch.sqrt(numerator / denominator))


def scan(args: argparse.Namespace) -> None:
    if args.global_index is None:
        raise ValueError("--global-index is required in scan mode")
    row = read_mapping_row(args.mapping, args.global_index)
    config_path = args.output_root / "scan_configs.json"
    config_manifest = json.loads(config_path.read_text(encoding="utf-8"))
    configurations = list(config_manifest["configs"])
    destination = args.output_root / "scans" / f"global_{args.global_index:04d}.json"
    existing: dict[str, Any] = {}
    if destination.is_file() and not args.overwrite:
        existing = json.loads(destination.read_text(encoding="utf-8"))
        if existing.get("status") == "complete":
            print(f"[cached] {destination}", flush=True)
            return
    completed = {record["name"]: record for record in existing.get("results", [])}
    device = torch.device("cuda")
    torch.set_grad_enabled(False)
    pipeline, predictor, text_encoder, step12_head, step123_head = load_models(
        args.experiment_root, device
    )
    prompt = str(row["extended_prompt"])
    conditional = text_encoder(text_prompts=[prompt])
    reference, reference_diagnostic = extended.generate_rollout(
        pipeline=pipeline,
        conditional_dict=conditional,
        seed=args.seed,
        device=device,
        predictor=predictor,
        head=None,
        config={
            "name": "ffff",
            "policy": "ffff",
            "candidate_steps": [],
            "beta": None,
            "threshold": None,
            "target_accepts": 0,
            "head": None,
            "risk_mode": "outgoing_span",
        },
    )
    for position, config in enumerate(configurations, start=1):
        name = str(config["name"])
        if name in completed:
            continue
        head = step12_head if config["head"] == "step12" else step123_head
        latent, diagnostic = extended.generate_rollout(
            pipeline=pipeline,
            conditional_dict=conditional,
            seed=args.seed,
            device=device,
            predictor=predictor,
            head=head,
            config=config,
        )
        decisions = diagnostic.pop("decisions")
        record = {
            **config,
            "accepted_predictor_calls": diagnostic["accepted_predictor_calls"],
            "full_calls": diagnostic["full_calls"],
            "policy_latency_ms": diagnostic["policy_latency_ms"],
            "context_dit_time_ms": diagnostic["context_dit_time_ms"],
            "latent_nrmse_vs_ffff": latent_nrmse(latent, reference),
            "decisions": decisions,
        }
        completed[name] = record
        atomic_json(
            destination,
            {
                "status": "running",
                "physical_gpu": PHYSICAL_GPU,
                "global_index": int(row["global_index"]),
                "prompt_suite": row["prompt_suite"],
                "suite_index": int(row["suite_index"]),
                "prompt": prompt,
                "seed": args.seed,
                "ffff": {
                    key: value
                    for key, value in reference_diagnostic.items()
                    if key != "decisions"
                },
                "results": [completed[key] for key in sorted(completed)],
            },
        )
        print(
            f"[scan] global={args.global_index} {position}/{len(configurations)} "
            f"name={name} K={record['accepted_predictor_calls']} "
            f"latent_nrmse={record['latent_nrmse_vs_ffff']:.6f}",
            flush=True,
        )
        del latent
    atomic_json(
        destination,
        {
            "status": "complete",
            "physical_gpu": PHYSICAL_GPU,
            "global_index": int(row["global_index"]),
            "prompt_suite": row["prompt_suite"],
            "suite_index": int(row["suite_index"]),
            "prompt": prompt,
            "seed": args.seed,
            "ffff": {
                key: value
                for key, value in reference_diagnostic.items()
                if key != "decisions"
            },
            "results": [completed[key] for key in sorted(completed)],
        },
    )


def summarize(args: argparse.Namespace) -> None:
    paths = sorted((args.output_root / "scans").glob("global_*.json"))
    if len(paths) != 3:
        raise ValueError(f"Expected exactly three scan files, found {len(paths)}")
    scans = [json.loads(path.read_text(encoding="utf-8")) for path in paths]
    if any(scan.get("status") != "complete" for scan in scans):
        raise RuntimeError("At least one scan is incomplete")
    by_name = [
        {record["name"]: record for record in scan["results"]} for scan in scans
    ]
    names = sorted(set(by_name[0]))
    if any(set(records) != set(names) for records in by_name[1:]):
        raise RuntimeError("Scan configuration sets differ")
    summary: list[dict[str, Any]] = []
    for name in names:
        rows = [records[name] for records in by_name]
        first = rows[0]
        summary.append(
            {
                "name": name,
                "head": first["head"],
                "candidate_steps": first["candidate_steps"],
                "beta": first["beta"],
                "threshold": first["threshold"],
                "target_accepts": first["target_accepts"],
                "mean_accepted_predictor_calls": sum(
                    float(row["accepted_predictor_calls"]) for row in rows
                )
                / len(rows),
                "mean_full_calls": sum(float(row["full_calls"]) for row in rows)
                / len(rows),
                "mean_policy_latency_ms": sum(
                    float(row["policy_latency_ms"]) for row in rows
                )
                / len(rows),
                "mean_latent_nrmse_vs_ffff": sum(
                    float(row["latent_nrmse_vs_ffff"]) for row in rows
                )
                / len(rows),
                "per_prompt_accepts": [
                    int(row["accepted_predictor_calls"]) for row in rows
                ],
                "per_prompt_latent_nrmse": [
                    float(row["latent_nrmse_vs_ffff"]) for row in rows
                ],
            }
        )
    selected: list[dict[str, Any]] = []
    for family, targets in TARGETS.items():
        for target in targets:
            candidates = [
                row
                for row in summary
                if row["head"] == family and int(row["target_accepts"]) == target
            ]
            within_budget = [
                row
                for row in candidates
                if abs(float(row["mean_accepted_predictor_calls"]) - target) <= 0.5
            ]
            pool = within_budget or candidates
            choice = min(
                pool,
                key=lambda row: (
                    float(row["mean_latent_nrmse_vs_ffff"]),
                    abs(float(row["mean_accepted_predictor_calls"]) - target),
                    float(row["mean_policy_latency_ms"]),
                    float(row["beta"]),
                )
                if within_budget
                else (
                    abs(float(row["mean_accepted_predictor_calls"]) - target),
                    float(row["mean_latent_nrmse_vs_ffff"]),
                    float(row["mean_policy_latency_ms"]),
                    float(row["beta"]),
                ),
            )
            selected.append(choice)
    atomic_json(
        args.output_root / "calibration_summary.json",
        {
            "status": "complete",
            "global_indices": [int(scan["global_index"]) for scan in scans],
            "selection_rule": (
                "within mean K +/-0.5 choose lowest mean latent nRMSE; if no "
                "candidate is within budget, choose nearest K first"
            ),
            "summary": summary,
            "selected": selected,
        },
    )
    print("[selected]", flush=True)
    for row in selected:
        print(
            f"  {row['head']} K{row['target_accepts']:02d}: beta={row['beta']} "
            f"threshold={row['threshold']:.8f} "
            f"actual_K={row['mean_accepted_predictor_calls']:.3f} "
            f"latent_nrmse={row['mean_latent_nrmse_vs_ffff']:.6f}",
            flush=True,
        )


def summarize_threshold_only(args: argparse.Namespace) -> None:
    paths = sorted((args.output_root / "scans").glob("global_*.json"))
    if len(paths) != 3:
        raise ValueError(f"Expected exactly three scan files, found {len(paths)}")
    scans = [json.loads(path.read_text(encoding="utf-8")) for path in paths]
    if any(scan.get("status") != "complete" for scan in scans):
        raise RuntimeError("At least one scan is incomplete")
    by_name = [
        {record["name"]: record for record in scan["results"]} for scan in scans
    ]
    names = sorted(set(by_name[0]))
    if any(set(records) != set(names) for records in by_name[1:]):
        raise RuntimeError("Scan configuration sets differ")
    summary: list[dict[str, Any]] = []
    for name in names:
        rows = [records[name] for records in by_name]
        first = rows[0]
        summary.append(
            {
                "name": name,
                "head": first["head"],
                "candidate_steps": first["candidate_steps"],
                "beta": first["beta"],
                "threshold": first["threshold"],
                "threshold_rank": first.get("threshold_rank"),
                "mean_accepted_predictor_calls": sum(
                    float(row["accepted_predictor_calls"]) for row in rows
                )
                / len(rows),
                "mean_full_calls": sum(float(row["full_calls"]) for row in rows)
                / len(rows),
                "mean_policy_latency_ms": sum(
                    float(row["policy_latency_ms"]) for row in rows
                )
                / len(rows),
                "mean_latent_nrmse_vs_ffff": sum(
                    float(row["latent_nrmse_vs_ffff"]) for row in rows
                )
                / len(rows),
                "per_prompt_accepts": [
                    int(row["accepted_predictor_calls"]) for row in rows
                ],
            }
        )
    selected: list[dict[str, Any]] = []
    for family, targets in TARGETS.items():
        for target in targets:
            candidates = [row for row in summary if row["head"] == family]
            within_budget = [
                row
                for row in candidates
                if abs(float(row["mean_accepted_predictor_calls"]) - target) <= 0.5
            ]
            pool = within_budget or candidates
            choice = min(
                pool,
                key=lambda row: (
                    float(row["mean_latent_nrmse_vs_ffff"]),
                    abs(float(row["mean_accepted_predictor_calls"]) - target),
                    float(row["threshold"]),
                )
                if within_budget
                else (
                    abs(float(row["mean_accepted_predictor_calls"]) - target),
                    float(row["mean_latent_nrmse_vs_ffff"]),
                    float(row["threshold"]),
                ),
            )
            selected.append(
                {
                    **choice,
                    "target_accepts": target,
                }
            )
    atomic_json(
        args.output_root / "threshold_summary.json",
        {
            "status": "complete",
            "beta": 2.0,
            "global_indices": [int(scan["global_index"]) for scan in scans],
            "selection_rule": (
                "fixed beta=2; for each target K choose threshold with mean K "
                "within +/-0.5 and lowest latent nRMSE"
            ),
            "summary": summary,
            "selected": selected,
        },
    )
    print("[selected-threshold-only]", flush=True)
    for row in selected:
        print(
            f"  {row['head']} K{row['target_accepts']:02d}: "
            f"threshold={row['threshold']:.8f} "
            f"actual_K={row['mean_accepted_predictor_calls']:.3f} "
            f"latent_nrmse={row['mean_latent_nrmse_vs_ffff']:.6f}",
            flush=True,
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--mode",
        choices=(
            "probe", "prepare", "prepare_threshold", "scan", "summarize",
            "summarize_threshold",
        ),
        required=True,
    )
    parser.add_argument("--gpu", default=PHYSICAL_GPU)
    parser.add_argument("--global-index", type=int, default=None)
    parser.add_argument("--mapping", type=Path, default=MAPPING_DEFAULT)
    parser.add_argument("--experiment-root", type=Path, default=EXPERIMENT_DEFAULT)
    parser.add_argument("--output-root", type=Path, default=OUTPUT_DEFAULT)
    parser.add_argument("--probe-root", type=Path, default=None)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()
    args.mapping = args.mapping.resolve()
    args.experiment_root = args.experiment_root.resolve()
    args.output_root = args.output_root.resolve()
    args.probe_root = args.probe_root.resolve() if args.probe_root is not None else None
    return args


def main() -> None:
    args = parse_args()
    args.output_root.mkdir(parents=True, exist_ok=True)
    if args.mode == "probe":
        probe(args)
    elif args.mode == "prepare":
        prepare(args)
    elif args.mode == "prepare_threshold":
        prepare_threshold_only(args)
    elif args.mode == "scan":
        scan(args)
    else:
        if args.mode == "summarize_threshold":
            summarize_threshold_only(args)
        else:
            summarize(args)


if __name__ == "__main__":
    main()