"""Estimate J_l = E_{prompt, t, t'>=t}[dh_final,t'/dh_l,t] over a pretraining-like corpus."""

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

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import load_model, prompt_jacobian_pairsum


def corpus_texts(n, min_chars=400):
    from datasets import load_dataset
    ds = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT",
                      split="train", streaming=True)
    got = 0
    for ex in ds:
        t = ex["text"].strip()
        if len(t) >= min_chars:
            yield t
            got += 1
            if got >= n:
                return


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n-prompts", type=int, default=256)
    ap.add_argument("--seq-len", type=int, default=64)
    ap.add_argument("--chunk", type=int, default=128)
    ap.add_argument("--out", default="results/jbar.pt")
    ap.add_argument("--dtype", default="float32")
    args = ap.parse_args()

    model, tok = load_model(dtype=getattr(torch, args.dtype))
    L = len(model.model.language_model.layers)
    d = model.config.get_text_config().hidden_size

    Jsum = torch.zeros(L, d, d, device="cuda", dtype=torch.float32)
    pairs_total = 0
    t0 = time.time()
    for i, text in enumerate(corpus_texts(args.n_prompts)):
        ids = tok(text, return_tensors="pt", truncation=True,
                  max_length=args.seq_len)["input_ids"].cuda()
        if ids.shape[1] < args.seq_len:
            continue
        J, n_pairs = prompt_jacobian_pairsum(model, ids, chunk=args.chunk)
        Jsum += J
        pairs_total += n_pairs
        if i % 5 == 0 or i == args.n_prompts - 1:
            el = time.time() - t0
            print(f"[{i+1}/{args.n_prompts}] {el:.0f}s ({el/(i+1):.1f}s/prompt)",
                  flush=True)
        if i % 50 == 49:  # checkpoint
            torch.save({"Jbar": (Jsum / pairs_total).cpu(),
                        "n_prompts": i + 1, "seq_len": args.seq_len},
                       args.out + ".ckpt")

    Jbar = (Jsum / pairs_total).cpu()
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    torch.save({"Jbar": Jbar, "n_prompts": args.n_prompts,
                "seq_len": args.seq_len, "pairs": pairs_total}, args.out)
    print("saved", args.out)


if __name__ == "__main__":
    main()