#!/usr/bin/env python3
"""Decompose paired VALID IoU changes by SFT-to-GRPO label transition.

Run from the repository root. This is descriptive analysis of the frozen
chronological VALID outputs; it does not access the TEST split.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import random
from collections import defaultdict
from pathlib import Path

from src.rewards.checker import statement_iou


def load_rows(path: Path) -> dict[str, dict]:
    rows = {}
    with path.open() as handle:
        for line in handle:
            if not line.strip():
                continue
            row = json.loads(line)
            sid = row["sample_id"]
            if sid in rows:
                raise ValueError(f"duplicate sample_id: {sid}")
            rows[sid] = row
    return rows


def predicted_label(row: dict) -> str:
    return row.get("pred_label") if row.get("format_ok") else "INVALID"


def row_iou(row: dict) -> float:
    # Match validation evaluation: invalid/unparsed rows contribute zero;
    # parsed rows are scored regardless of their predicted class.
    if not row.get("format_ok"):
        return 0.0
    return statement_iou(set(row.get("pred_lines") or []),
                         set(row.get("gold_lines") or []))


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def bootstrap_mean(values: list[float], *, seed: int, reps: int = 20_000) -> list[float] | None:
    if not values:
        return None
    rng = random.Random(seed)
    n = len(values)
    means = sorted(sum(rng.choice(values) for _ in range(n)) / n
                   for _ in range(reps))
    return [means[int(0.025 * reps)], means[int(0.975 * reps) - 1]]


def analyze(sft_path: Path, grpo_path: Path, *, seed: int = 20261008) -> dict:
    sft, grpo = load_rows(sft_path), load_rows(grpo_path)
    if sft.keys() != grpo.keys():
        raise ValueError("SFT and GRPO sample_id sets differ")

    grouped: dict[tuple[str, str], list[tuple[float, float]]] = defaultdict(list)
    overall_sft, overall_grpo = [], []
    class_parity = True
    for sid, a in sft.items():
        b = grpo[sid]
        if int(a["target"]) != int(b["target"]):
            class_parity = False
            raise ValueError(f"target mismatch for {sid}")
        if set(a.get("gold_lines") or []) != set(b.get("gold_lines") or []):
            raise ValueError(f"gold line mismatch for {sid}")
        if int(a["target"]) != 1:
            continue
        sa, gb = row_iou(a), row_iou(b)
        overall_sft.append(sa)
        overall_grpo.append(gb)
        grouped[(predicted_label(a), predicted_label(b))].append((sa, gb))

    n = len(overall_sft)
    transitions = {}
    for (sft_label, grpo_label), values in sorted(grouped.items()):
        ds = [b - a for a, b in values]
        transitions[f"{sft_label}->{grpo_label}"] = {
            "n": len(values),
            "sft_mean_iou": sum(a for a, _ in values) / len(values),
            "grpo_mean_iou": sum(b for _, b in values) / len(values),
            "paired_mean_delta": sum(ds) / len(ds),
            "contribution_to_overall_mean_delta": sum(ds) / n,
        }

    both_vulnerable = grouped.get(("VULNERABLE", "VULNERABLE"), [])
    common_deltas = [b - a for a, b in both_vulnerable]
    overall_deltas = [b - a for a, b in zip(overall_sft, overall_grpo)]
    return {
        "analysis": "gold-vulnerable VALID localization by paired label transition",
        "n_valid_rows_paired": len(sft),
        "n_gold_vulnerable": n,
        "gold_target_and_lines_parity_verified": class_parity,
        "test_touched": False,
        "overall": {
            "sft_mean_iou": sum(overall_sft) / n,
            "grpo_step200_mean_iou": sum(overall_grpo) / n,
            "paired_mean_delta": sum(overall_deltas) / n,
            "paired_row_bootstrap_95_ci": bootstrap_mean(overall_deltas, seed=seed),
        },
        "transitions": transitions,
        "both_predicted_vulnerable": {
            "n": len(common_deltas),
            "sft_mean_iou": transitions.get("VULNERABLE->VULNERABLE", {}).get("sft_mean_iou"),
            "grpo_mean_iou": transitions.get("VULNERABLE->VULNERABLE", {}).get("grpo_mean_iou"),
            "paired_mean_delta": (sum(common_deltas) / len(common_deltas)
                                  if common_deltas else None),
            "paired_row_bootstrap_95_ci": bootstrap_mean(common_deltas, seed=seed + 1),
        },
        "bootstrap": {"unit": "gold-vulnerable sample (one per pair)",
                      "replicates": 20_000, "seed": seed},
        "input_sha256": {str(sft_path): sha256(sft_path),
                         str(grpo_path): sha256(grpo_path)},
        "interpretation_caveat": (
            "Transitions condition on model predictions and are descriptive, not causal. "
            "An overall IoU increase can reflect more vulnerable rows receiving a "
            "VULNERABLE prediction and nonempty localization; it does not alone show "
            "better localization precision on rows both models detect."
        ),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--sft", type=Path,
                        default=Path("results/phase4/contract_v2/sft_val_step_500.jsonl"))
    parser.add_argument("--grpo", type=Path,
                        default=Path("results/phase6/scale_20261007_v5/valid_predictions.jsonl"))
    parser.add_argument("--output", type=Path,
                        default=Path("results/phase6/scale_20261007_v5/localization_transition_decomposition_step200.json"))
    args = parser.parse_args()
    result = analyze(args.sft, args.grpo)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()