"""Read-only prompt-cluster audit of a stopped scale run's parser failures."""
from __future__ import annotations

from collections import Counter
import json
from pathlib import Path
import sys

import numpy as np
import torch


def audit(run_dir: Path, *, seed: int = 42, draws: int = 20_000) -> dict:
    manifest = json.loads((run_dir / "run_manifest.json").read_text())
    completed = json.loads((run_dir / "curves.json").read_text())
    pending = json.loads((run_dir / "pending_step.json").read_text())
    steps = len(completed) + 1
    assert pending["optimizer_step"] == steps
    selection = manifest["train_selection"]
    order = torch.randperm(len(selection), generator=torch.Generator().manual_seed(
        manifest["seed"])).tolist()[:steps * 64]
    chosen = [selection[i] for i in order]
    assert len(chosen) == steps * 64
    labels = {row["prompt_id"]: int(row["target"]) for row in chosen}
    assert len(labels) == len(chosen), "a prompt recurred during the first epoch"
    failures = [json.loads(line) for line in (run_dir / "rollout_failures.jsonl").open()]
    assert all(row["prompt_id"] in labels for row in failures)
    assert len(failures) == sum(row["verifier_failures"] for row in completed) + pending["verifier_failures"]

    by_prompt = Counter(row["prompt_id"] for row in failures)
    counts = np.asarray([by_prompt[row["prompt_id"]] for row in chosen])
    classes = np.asarray([int(row["target"]) for row in chosen])
    rates = {name: float(counts[classes == label].sum() / (8 * (classes == label).sum()))
             for name, label in (("SAFE", 0), ("VULNERABLE", 1))}
    observed = rates["VULNERABLE"] - rates["SAFE"]

    # The eight rollouts sharing a prompt are correlated. Resample/permute
    # entire prompt groups rather than treating 7,680 rollouts as IID.
    rng = np.random.default_rng(seed)
    safe = counts[classes == 0]
    vuln = counts[classes == 1]
    bootstrap = np.empty(draws)
    permutation = np.empty(draws)
    for i in range(draws):
        bootstrap[i] = (rng.choice(vuln, len(vuln), replace=True).mean() -
                        rng.choice(safe, len(safe), replace=True).mean()) / 8
        perm = rng.permutation(counts)
        permutation[i] = (perm[:len(vuln)].mean() - perm[len(vuln):].mean()) / 8
    return {
        "run_dir": str(run_dir), "seed": seed, "draws": draws,
        "steps_including_pending": steps, "prompts": len(chosen),
        "rollouts": len(chosen) * 8, "failures": len(failures),
        "class_counts": {
            name: {"prompts": int((classes == label).sum()),
                   "rollouts": int(8 * (classes == label).sum()),
                   "failures": int(counts[classes == label].sum()),
                   "prompts_with_failure": int((counts[classes == label] > 0).sum()),
                   "failure_rate": rates[name]}
            for name, label in (("SAFE", 0), ("VULNERABLE", 1))},
        "vulnerable_minus_safe_failure_rate": observed,
        "prompt_cluster_bootstrap_95_interval": np.quantile(bootstrap, [.025, .975]).tolist(),
        "prompt_label_permutation_two_sided_p": float(np.mean(np.abs(permutation) >= abs(observed))),
        "prompts_with_multiple_failures": int((counts >= 2).sum()),
        "failures_on_multiple_failure_prompts": int(counts[counts >= 2].sum()),
        "max_failures_on_one_prompt": int(counts.max()),
        "note": "Exploratory prompt-cluster comparison; no causal or equal-rate claim.",
    }


if __name__ == "__main__":
    print(json.dumps(audit(Path(sys.argv[1])), indent=2))