#!/usr/bin/env python
"""Stage C: counterfactual (paired intervention) label collection with a FIXED editor checkpoint.

For each batch: encode; reach a state s by applying r_pre ~ U{0..max_prefix} random forced actions; evaluate
the 4 logged candidate actions (policy preference if a controller is given, competitive alternative, 2 random;
singles/pairs balanced) from the SAME state and store g(s,a;y) with log-propensities. On every `audit_every`-th
batch all 24 actions are evaluated (oracle audit). The dataset stores what is needed to REGENERATE each state
deterministically (tokens, prefix actions) rather than the slot tensors. Partitions never share a state.

Usage: python -m training.collect_interventions --config configs/experiments/labels_pilot.yaml [--device mps]
"""
from __future__ import annotations

import argparse, json, sys, time
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from prisme.counterfactual import all_action_gains, paired_gains, sample_logged_actions, state_loss  # noqa: E402
from prisme.flops import action_cost_proxy  # noqa: E402
from training.common import Run, TokenBin, load_checkpoint, load_yaml, pick_device, seed_all, sync  # noqa: E402
from training.train_editor import build_model, sample_forced_actions  # noqa: E402
from data.synthetic_dataset import AnyTokenizer, SyntheticBatches, load_items  # noqa: E402


@torch.no_grad()
def collect(model, cfg, batches, n_batches, rng, device, max_prefix, audit_every, use_policy, cost):
    recs = []
    for bi in range(n_batches):
        x, y = next(batches)
        B, T = x.shape
        ctx = model.encode(x)
        r_pre = int(rng.integers(0, max_prefix + 1))
        R_total = int(rng.choice([1, 2, 4, 8]))
        R_total = max(R_total, r_pre + 1)
        prefix, _ = sample_forced_actions(cfg, B, T, [r_pre], [1.0], rng, device) if r_pre > 0 else (torch.full((B, T, cfg.r_max), cfg.stop_action, device=device, dtype=torch.long), None)
        S = ctx.S0
        edited = torch.zeros(B, T, dtype=torch.bool, device=device)
        for r in range(r_pre):
            a = prefix[:, :, r]
            S, _ = model.apply_action(S, ctx, a, r, R_total - r - 1)
            edited |= a != cfg.stop_action
        step = torch.full((B, T), r_pre, device=device, dtype=torch.long)
        remaining = torch.full((B, T), R_total - r_pre - 1, device=device, dtype=torch.long)
        pref = None
        if use_policy:
            preds = model.controller(ctx.ctrl_in, S, step, remaining, ctx.entropy)
            q = model.controller.q_values(preds, cost)
            pref = q.argmax(-1)
        gen = torch.Generator(device="cpu").manual_seed(int(rng.integers(0, 2**31)))
        acts, logp = sample_logged_actions(cfg, B, T, pref.cpu() if pref is not None else None, generator=gen, device="cpu")
        acts = acts.to(device); logp = logp.to(device)
        gains = paired_gains(model, S, ctx, y, edited, acts, step, remaining)
        base = state_loss(model, S, ctx, y, edited)
        rec = {"tokens": x.cpu().to(torch.int32), "targets": y.cpu().to(torch.int32), "prefix": prefix[:, :, :max(r_pre, 1)].cpu().to(torch.int16),
               "r_pre": r_pre, "R_total": R_total, "cand": acts.cpu().to(torch.int16), "gain": gains.cpu(), "logp": logp.cpu(),
               "base_loss": base.cpu(), "valid": (y != -100).cpu()}
        if audit_every and bi % audit_every == 0:
            rec["all_gains"] = all_action_gains(model, S, ctx, y, edited, step, remaining).cpu()
        recs.append(rec)
        if bi % 10 == 0:
            v = rec["valid"]
            print(f"batch {bi}/{n_batches} r_pre={r_pre} mean gain {float(gains[v].mean()):.5f} frac>0 {float((gains[v] > 0).float().mean()):.3f}", flush=True)
    return recs


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--config", required=True)
    ap.add_argument("--device", default="auto")
    ap.add_argument("--seed", type=int, default=None)
    ap.add_argument("--run-root", default="experiments/runs")
    args = ap.parse_args(argv)
    exp = load_yaml(args.config)
    seed = args.seed if args.seed is not None else exp.get("seed", 1)
    seed_all(seed); device = pick_device(args.device)
    cfg, model = build_model(exp)
    model = model.to(device).eval()
    meta = load_checkpoint(exp["checkpoint"], model, strict=False)
    ck_sha = meta.get("weights_sha256")
    tr = exp["collect"]
    data = exp["data"]
    run = Run(exp.get("name", "labels"), {"experiment": exp, "model": cfg.to_dict()}, root=args.run_root, seed=seed)
    run.write_manifest(device=str(device), checkpoint_parent=exp["checkpoint"], checkpoint_sha256=ck_sha)
    cost = torch.tensor(action_cost_proxy(cfg, tr["seq_len"]), device=device)
    parts = {}
    for part, spec in exp["partitions"].items():   # e.g. train: {split: train, n_batches: 200}, calib: {split: calib, ...}
        rng = np.random.default_rng(seed + hash(part) % 1000)
        if data["kind"] == "tokens":
            tb = TokenBin(spec["bin"], tr["seq_len"]); it = tb.batches(tr["batch_size"], seed + spec.get("stream_seed", 0), device)
        else:
            tok = AnyTokenizer(data.get("tokenizer", "char"))
            sb = SyntheticBatches(load_items(Path(data["splits_dir"]) / f"{spec['split']}.jsonl"), tok, tr["seq_len"])
            it = sb.random_batches(tr["batch_size"], seed + spec.get("stream_seed", 0), device, data.get("loss_span", "answer"))
        t0 = time.time()
        recs = collect(model, cfg, it, spec["n_batches"], rng, device, tr.get("max_prefix", 3), tr.get("audit_every", 5),
                       exp.get("use_policy", False), cost)
        sync(device)
        path = run.dir / f"labels_{part}.pt"
        torch.save({"records": recs, "checkpoint_sha256": ck_sha, "config": cfg.to_dict(), "part": part, "seed": seed}, path)
        g = torch.cat([r["gain"][r["valid"]] for r in recs]); n_states = sum(int(r["valid"].sum()) for r in recs)
        parts[part] = {"file": str(path), "n_states": n_states, "n_labels": int(g.numel()), "mean_gain": float(g.mean()),
                       "frac_positive": float((g > 0).float().mean()), "frac_harm": float((g < 0).float().mean()),
                       "seconds": round(time.time() - t0, 1), "n_audit_batches": sum("all_gains" in r for r in recs)}
        print(part, parts[part], flush=True)
    run.finish(partitions=parts)
    print("run", run.run_id)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())