"""Recipe-finder v2 for StarCoder2 MLP compression.

Improvements over v1: (1) sequence PACKING + batching for ~8x token throughput,
(2) FEATURE distillation -- match each compressed MLP's output to what the
ORIGINAL MLP produces on the same input (the "mimic the weights" / doc-to-LoRA
signal), computed by recomputing the frozen orig MLP on the student's own
activations (cheap, no full second forward), (3) optional logit KL with
temperature, (4) tuned layernorms, (5) expandable_segments to fit bigger banks.

Objectives: loss = NTP + feat_w * mean_layer relMSE(bank_out, orig_mlp(in)) [+ kl].

Usage: python train_compress2.py --tag r_feat --E 2048 --feat-w 1 --tune-norms 1 \
         --batch 2 --ctx 512 --steps 4000
"""
import argparse, itertools, math, time
import torch, torch.nn as nn, torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

DEV = 0

class Bank(nn.Module):
    def __init__(self, src_mlp, E, init):
        super().__init__()
        if init == "topnorm":
            idx = src_mlp.c_proj.weight.data.float().norm(dim=0).topk(E).indices
        else:
            idx = torch.randperm(src_mlp.c_fc.weight.shape[0])[:E]
        self.down = nn.Parameter(src_mlp.c_fc.weight.data[idx].clone().float())
        self.up = nn.Parameter(src_mlp.c_proj.weight.data[:, idx].t().clone().float())
        self.b = nn.Parameter(src_mlp.c_fc.bias.data[idx].clone().float()
                              if src_mlp.c_fc.bias is not None else torch.zeros(E))
        self.obias = nn.Parameter(src_mlp.c_proj.bias.data.clone().float()
                                  if src_mlp.c_proj.bias is not None
                                  else torch.zeros(src_mlp.c_proj.weight.shape[0]))
        self.last_in = None; self.last_out = None
    def forward(self, x):
        self.last_in = x
        act = F.gelu(x.float() @ self.down.t() + self.b, approximate="tanh")
        out = (act @ self.up + self.obias).to(x.dtype)
        self.last_out = out
        return out


def token_stream(tok, ctx, n_docs=8000):
    texts = []
    for name, cfg, field in [("codeparrot/codeparrot-clean", None, "content"),
                             ("HuggingFaceFW/fineweb-edu", "sample-10BT", "text")]:
        try:
            ds = (load_dataset(name, cfg, split="train", streaming=True) if cfg
                  else load_dataset(name, split="train", streaming=True))
            for ex in itertools.islice(ds, n_docs):
                t = ex.get(field) or ""
                if len(t) > 60: texts.append(t)
        except Exception as e:
            print("FAIL", name, str(e)[:100], flush=True)
    import random; random.shuffle(texts)
    buf = []
    while True:
        for t in texts:
            buf.extend(tok(t).input_ids + [tok.eos_token_id or 0])
            while len(buf) >= ctx:
                yield torch.tensor(buf[:ctx]); buf = buf[ctx:]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tag", default="run")
    ap.add_argument("--E", type=int, default=2048)
    ap.add_argument("--init", default="random")
    ap.add_argument("--feat-w", type=float, default=1.0)
    ap.add_argument("--kl-w", type=float, default=0.0)
    ap.add_argument("--kl-temp", type=float, default=2.0)
    ap.add_argument("--tune-norms", type=int, default=1)
    ap.add_argument("--batch", type=int, default=2)
    ap.add_argument("--ctx", type=int, default=512)
    ap.add_argument("--steps", type=int, default=4000)
    ap.add_argument("--lr", type=float, default=5e-4)
    ap.add_argument("--warmup", type=int, default=150)
    args = ap.parse_args()
    print(f"=== {args.tag}: E={args.E} feat_w={args.feat_w} kl_w={args.kl_w} "
          f"norms={args.tune_norms} batch={args.batch} ctx={args.ctx} ===", flush=True)

    tok = AutoTokenizer.from_pretrained("bigcode/starcoder2-3b")
    m = AutoModelForCausalLM.from_pretrained("bigcode/starcoder2-3b", dtype=torch.bfloat16,
                                             device_map={"": DEV})
    m.config.use_cache = False
    for p in m.parameters(): p.requires_grad_(False)
    layers = m.model.layers
    orig_mlps = [l.mlp for l in layers]
    banks = [Bank(l.mlp, args.E, args.init).to(DEV) for l in layers]
    def use_banks(on):
        for l, om, bk in zip(layers, orig_mlps, banks):
            l.mlp = bk if on else om

    gen = token_stream(tok, args.ctx)
    def get_batch():
        return torch.stack([next(gen) for _ in range(args.batch)]).to(DEV)

    @torch.no_grad()
    def ppl(on, n=15):
        m.eval(); use_banks(on); tot = 0.0
        for _ in range(n):
            ids = get_batch(); tot += m(ids, labels=ids).loss.item()
        return math.exp(tot / n)
    op = ppl(False); ip = ppl(True)
    print(f"[ref] ORIGINAL {op:.2f} | INIT {ip:.1f} ppl", flush=True)

    params = [p for bk in banks for p in bk.parameters()]
    if args.tune_norms:
        for l in layers:
            for mod in (l.input_layernorm, l.post_attention_layernorm):
                for p in mod.parameters(): p.requires_grad_(True); params.append(p)
    for p in params: p.requires_grad_(True)
    m.gradient_checkpointing_enable()
    opt = torch.optim.AdamW(params, lr=args.lr, betas=(0.9, 0.95))
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s:
        s / args.warmup if s < args.warmup else
        0.5 * (1 + math.cos(math.pi * min(1.0, (s - args.warmup) / max(1, args.steps - args.warmup)))))

    t0 = time.time(); ema = None
    for step in range(1, args.steps + 1):
        ids = get_batch()
        if args.kl_w > 0:
            use_banks(False); m.eval()
            with torch.no_grad(): t_logits = m(ids).logits
        use_banks(True); m.train()
        out = m(ids, labels=ids)
        loss = out.loss; ce = out.loss.item()
        if args.feat_w > 0:                       # mimic each MLP's output
            fl = 0.0
            for om, bk in zip(orig_mlps, banks):
                with torch.no_grad(): tgt = om(bk.last_in)
                fl = fl + ((bk.last_out - tgt).float().pow(2).mean()
                           / tgt.float().pow(2).mean().clamp_min(1e-6))
            loss = loss + args.feat_w * fl / len(banks)
        if args.kl_w > 0:
            T = args.kl_temp
            kl = F.kl_div(F.log_softmax(out.logits / T, -1), F.softmax(t_logits / T, -1),
                          reduction="batchmean") * (T * T) / out.logits.shape[1]
            loss = loss + args.kl_w * kl
        opt.zero_grad(set_to_none=True); loss.backward()
        torch.nn.utils.clip_grad_norm_(params, 1.0); opt.step(); sched.step()
        ema = ce if ema is None else 0.98 * ema + 0.02 * ce
        if step % 25 == 0:
            tps = step * args.batch * args.ctx / (time.time() - t0)
            print(f"[{args.tag}] step {step}/{args.steps} ce {ce:.3f} ema {ema:.3f} "
                  f"ppl {math.exp(ema):.1f} (orig {op:.1f}) {tps/1000:.0f}k tok/s", flush=True)
    fp = ppl(True, 30)
    print(f"\n[result {args.tag}] ORIGINAL {op:.2f} | INIT {ip:.1f} | FINAL {fp:.1f} ppl  "
          f"(tokens {args.steps*args.batch*args.ctx/1e6:.1f}M)", flush=True)
    torch.save({"E": args.E, "init": args.init, "final_ppl": fp, "orig_ppl": op,
                "banks": [bk.state_dict() for bk in banks]}, f"/tmp/banks_{args.tag}.pt")
    print("DONE " + args.tag, flush=True)


if __name__ == "__main__":
    main()