"""Report the row-level Wilson gate and pair-cluster VALID accuracy interval."""
from __future__ import annotations

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

import numpy as np


def wilson(k: int, n: int, z: float = 1.96) -> list[float]:
    p = k / n
    denom = 1 + z * z / n
    center = p + z * z / (2 * n)
    half = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n))
    return [(center - half) / denom, (center + half) / denom]


def audit(predictions_path: Path, *, replicates: int = 20_000, seed: int = 42) -> dict:
    pairs = defaultdict(list)
    with predictions_path.open() as stream:
        for line in stream:
            row = json.loads(line)
            pairs[row["pair_id"]].append(row)
    if not pairs or any(len(rows) != 2 for rows in pairs.values()):
        raise ValueError("expected exactly two validation rows per pair")

    pair_values = []
    class_counts = {"SAFE": {"correct": 0, "n": 0},
                    "VULNERABLE": {"correct": 0, "n": 0}}
    outcome_counts = {"both_correct": 0, "safe_only": 0,
                      "vulnerable_only": 0, "neither": 0}
    for rows in pairs.values():
        by_target = {int(row["target"]): row for row in rows}
        if set(by_target) != {0, 1}:
            raise ValueError("each pair must contain one SAFE and one VULNERABLE row")
        safe = by_target[0]
        vuln = by_target[1]
        safe_ok = bool(safe["format_ok"] and safe["pred_label"] == "SAFE")
        vuln_ok = bool(vuln["format_ok"] and vuln["pred_label"] == "VULNERABLE")
        pair_values.append((int(safe_ok), int(vuln_ok)))
        class_counts["SAFE"]["correct"] += int(safe_ok)
        class_counts["SAFE"]["n"] += 1
        class_counts["VULNERABLE"]["correct"] += int(vuln_ok)
        class_counts["VULNERABLE"]["n"] += 1
        key = ("both_correct" if safe_ok and vuln_ok else
               "safe_only" if safe_ok else
               "vulnerable_only" if vuln_ok else "neither")
        outcome_counts[key] += 1

    array = np.asarray(pair_values, dtype=np.float64)
    total_pairs = int(len(array))
    total_rows = int(array.size)
    total_correct = int(array.sum())
    rng = np.random.default_rng(seed)
    indices = rng.integers(0, len(array), size=(replicates, len(array)))
    pair_bootstrap = array[indices].mean(axis=(1, 2))
    threshold = next(k for k in range(total_rows + 1)
                     if wilson(k, total_rows)[0] > 0.5)
    classes = {}
    for label, values in class_counts.items():
        k, n = values["correct"], values["n"]
        classes[label] = {"correct": k, "n": n, "accuracy": k / n,
                          "wilson_95": wilson(k, n)}
    digest = hashlib.sha256(predictions_path.read_bytes()).hexdigest()
    return {
        "source": str(predictions_path),
        "source_sha256": digest,
        "pairs": total_pairs,
        "rows": total_rows,
        "row_accuracy": total_correct / total_rows,
        "row_correct": total_correct,
        "row_wilson_95": wilson(total_correct, total_rows),
        "minimum_row_correct_for_wilson_lower_above_0_50": threshold,
        "additional_correct_needed_from_observed": max(0, threshold - total_correct),
        "per_class": classes,
        "pair_outcome_counts": outcome_counts,
        "pair_cluster_bootstrap_accuracy_95": [
            float(x) for x in np.quantile(pair_bootstrap, [0.025, 0.975])
        ],
        "pair_bootstrap_replicates": replicates,
        "pair_bootstrap_seed": seed,
    }


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("--output", type=Path, default=Path(
        "results/phase6/scale_20261007_v5/valid_gate_precision_audit.json"))
    parser.add_argument("--replicates", type=int, default=20_000)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()
    report = audit(args.predictions, replicates=args.replicates, seed=args.seed)
    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()