#!/usr/bin/env python3
"""How much does full BPTT through the latent chain actually cost at depth 15/20?

`backprop_depth: null` (every config) means full BPTT: gradients flow through all
D latent recurrence steps. coconut.Coconut also implements a truncated window --
if backprop_depth=W, passes older than W steps are detached, capping backward
compute and activation memory at O(W) instead of O(D).

This times one forward+backward at batch_size_training on real prompts, sweeping
the number of latents and the BPTT window, so the depth cost can be separated
from the dataset-size cost.

Usage:
  PYTHONPATH=. python scripts/bench_bptt.py --train data/star_2arm_L20_100k_train_fo_bfs.json
"""
import argparse
import json
import time

import torch
from transformers import AutoModelForCausalLM, AutoConfig

from stokenizer import STokenizer
from coconut import Coconut


def build_batch(data, tok, latent_id, n_latent, bs, device):
    seqs = []
    for s in data[:bs]:
        q = "<eos> " + "|".join([f" {e[0]} {e[1]} " for e in s["edges"]]).strip() + " [Q] "
        q += f'{s["target"]} {s["neg_target"]} [R] {s["root"]}'
        ids = tok.encode(q, add_special_tokens=False) + [latent_id] * n_latent
        # supervise a single node token, as dataset.expand_data does
        hop = min(max(1, n_latent + 1), max(int(k) for k in s["neighbor_k"]))
        tgt = s["neighbor_k"][str(hop)][0]
        ids = ids + tok.encode(str(tgt), add_special_tokens=False)
        seqs.append(ids)
    n = min(len(x) for x in seqs)
    seqs = [x[:n] for x in seqs]
    ids = torch.tensor(seqs, device=device)
    labels = ids.clone()
    labels[:, :-1] = -100
    return ids, labels


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--train", required=True)
    ap.add_argument("--model_id", default="configs/symbol-2layer-8head-768dim-L20.json")
    ap.add_argument("--bs", type=int, default=128)
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--iters", type=int, default=8)
    args = ap.parse_args()

    tok = STokenizer()
    latent_id = tok.convert_tokens_to_ids("<|latent|>")
    data = json.load(open(args.train))[: args.bs]

    print(f"data: {args.train}   edges/graph: {len(data[0]['edges'])}   bs={args.bs}\n")
    print(f"{'n_latent':>9} {'bptt_W':>7} {'fwd+bwd ms':>12} {'peak MiB':>10}")

    for n_latent in (0, 5, 10, 15, 20):
        for W in (None, 1, 2, 5):
            if W is not None and n_latent <= W:
                continue
            base = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(args.model_id))
            model = Coconut(base, latent_id,
                            tok.convert_tokens_to_ids("<|start-latent|>"),
                            tok.convert_tokens_to_ids("<|end-latent|>"),
                            tok.eos_token_id, backprop_depth=W).to(args.device)
            ids, labels = build_batch(data, tok, latent_id, n_latent, args.bs, args.device)
            pos = torch.arange(ids.shape[1], device=args.device).unsqueeze(0).expand(len(ids), -1)
            attn = torch.ones_like(ids)
            torch.cuda.reset_peak_memory_stats(args.device)
            for i in range(args.iters):
                if i == 2:
                    torch.cuda.synchronize(); t0 = time.time()
                out = model.forward(ids, attn, labels, pos)
                out.loss.backward()
                model.zero_grad(set_to_none=True)
            torch.cuda.synchronize()
            ms = (time.time() - t0) / (args.iters - 2) * 1e3
            peak = torch.cuda.max_memory_allocated(args.device) / 2**20
            print(f"{n_latent:>9} {str(W):>7} {ms:>12.1f} {peak:>10.0f}")
            del model, ids, labels
            torch.cuda.empty_cache()


if __name__ == "__main__":
    main()