"""Audit eligible step-200 VALID R2 scores by gold label (CPU-only)."""
from __future__ import annotations

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


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def audit(predictions_path: Path, annotations_path: Path) -> dict:
    annotations = {}
    with annotations_path.open() as stream:
        for line in stream:
            row = json.loads(line)
            annotations[row["sample_id"]] = row

    by_label: dict[str, list[float]] = defaultdict(list)
    both_empty: dict[str, int] = defaultdict(int)
    eligible_correct: dict[str, int] = defaultdict(int)
    correct_but_ineligible: dict[str, int] = defaultdict(int)
    total = 0
    with predictions_path.open() as stream:
        for line in stream:
            row = json.loads(line)
            total += 1
            label = "VULNERABLE" if int(row["target"]) == 1 else "SAFE"
            if row["pred_label"] != label:
                continue
            # Mirrors score_sample.r2_eligible: all SAFE rows are eligible;
            # vulnerable rows require at least one annotated L1 line.
            eligible = label == "SAFE" or bool(
                annotations.get(row["sample_id"], {}).get("l1_lines")
            )
            if not eligible:
                correct_but_ineligible[label] += 1
                continue
            pred = set(row["pred_lines"])
            gold = set(row["gold_lines"])
            union = pred | gold
            iou = 1.0 if not union else len(pred & gold) / len(union)
            by_label[label].append(iou)
            eligible_correct[label] += 1
            both_empty[label] += int(not pred and not gold)

    return {
        "predictions": str(predictions_path),
        "predictions_sha256": sha256(predictions_path),
        "annotations": str(annotations_path),
        "annotations_sha256": sha256(annotations_path),
        "total_valid_rows": total,
        "correct_but_r2_ineligible": dict(correct_but_ineligible),
        "eligible_correct_predictions": {
            label: {
                "n": eligible_correct[label],
                "mean_r2": statistics.mean(scores) if scores else None,
                "both_empty_n": both_empty[label],
                "both_empty_fraction": (
                    both_empty[label] / len(scores) if scores else None
                ),
            }
            for label, scores in sorted(by_label.items())
        },
        "scope": "fixed-output R2 replay; no inference, training, or policy counterfactual",
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--predictions", type=Path, default=Path(
        "results/phase6/scale_20261007_v5/valid_predictions.jsonl"))
    parser.add_argument("--annotations", type=Path, default=Path(
        "data/processed/annotations.jsonl"))
    parser.add_argument("--output", type=Path, default=Path(
        "results/phase6/scale_20261007_v5/valid_empty_set_r2_audit.json"))
    args = parser.parse_args()
    report = audit(args.predictions, args.annotations)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()