"""Capture pinned CPU reference attention inputs into a private directory."""
from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Any, Callable, cast

import torch


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--verify-root", required=True)
    ap.add_argument("--captures", required=True)
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--official", required=True)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--keys", default="r000008_b0,r000011_b0")
    a = ap.parse_args()
    sys.path.insert(0, a.verify_root)
    from verify.captures import CaptureSet
    from verify.native_ref import NativeRef
    from verify.run_ref import WindowCaches, prepare_window
    torch.set_num_threads(12)
    cs = CaptureSet(a.captures)
    keys = set(a.keys.split(","))
    samples = [s for s in cs.samples if s.key in keys]
    if {s.key for s in samples} != keys:
        raise ValueError("missing requested reference sample")
    ref = NativeRef(a.ckpt, a.official, mode="mirror", max_seq_len=max(s.P for s in samples) + 80)
    caches = cast(Callable[[], Any], WindowCaches)()
    module = cast(Any, ref.M)
    original = module.sparse_attn
    a.out.mkdir(parents=True, exist_ok=False)
    for sample in samples:
        _, _, mx, _ = prepare_window(ref, cs, sample, caches)
        count = 0

        def traced(q: torch.Tensor, kv: torch.Tensor, sink: torch.Tensor, idx: torch.Tensor,
                   scale: float) -> torch.Tensor:
            nonlocal count
            result = cast(torch.Tensor, original(q, kv, sink, idx, scale))
            torch.save(dict(q=q.clone(), kv=kv.clone(), sink=sink.clone(), indices=idx.clone(),
                            scale=scale, out=result.clone()), a.out / f"{sample.key}-s{count}.pt")
            count += 1
            return result

        module.sparse_attn = traced
        ref.draft(sample.P, sample.anchor, mx, keep_full_logits=True)
        module.sparse_attn = original
        print(f"saved {sample.key}: {count} stages", flush=True)


if __name__ == "__main__":
    main()