"""Shared pieces for the decoupled (capture -> offline-train) MLP compression
pipeline.

`Bank`, `collect`, and `build_buffer` are vendored VERBATIM (modulo this
docstring) from the in-context recipe scripts so this directory is a
self-contained, committable, re-runnable pipeline:

  * Bank          <- /tmp/train_compress2.py   (subset-of-neurons compressed MLP)
  * collect       <- /tmp/train_compare.py      (per-source streaming w/ retry)
  * build_buffer  <- /tmp/train_compare.py      (reproducible corpus buffer)

Keeping build_buffer byte-identical to train_compare.py guarantees the held-out
eval set used by train_offline.py is the SAME fixed eval set as the in-context
baseline (52.4 ppl @ 5M tokens), so the two regimes are directly comparable.
"""
import itertools, time, random
import torch, torch.nn as nn, torch.nn.functional as F
from datasets import load_dataset

MODEL = "bigcode/starcoder2-3b"


class Bank(nn.Module):
    """A compressed MLP that keeps only E of the original intermediate neurons.
    Initialized from the source MLP's own weights (a subset of rows/cols)."""
    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 collect(name, cfg, field, n_docs, attempts=5):
    """Stream up to n_docs texts; retry the whole source on transient errors
    (e.g. HF CDN 408) until we have a healthy fraction."""
    texts = []
    for attempt in range(attempts):
        texts = []
        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)
            if len(texts) >= n_docs * 0.5:
                print(f"  {name}: {len(texts)} docs (attempt {attempt+1})", flush=True)
                return texts
            print(f"  {name}: short ({len(texts)}), retrying", flush=True)
        except Exception as e:
            print(f"  {name}: attempt {attempt+1} failed: {str(e)[:90]}", flush=True)
        time.sleep(5)
    print(f"  {name}: giving up with {len(texts)} docs", flush=True)
    return texts


def build_buffer(tok, target_tokens, seed=0):
    """Build a reproducible token buffer. The first K tokens are deterministic
    given (collected corpus, seed), so a prefix slice is a fixed eval set."""
    texts = collect("codeparrot/codeparrot-clean", None, "content", 9000)
    texts += collect("HuggingFaceFW/fineweb-edu", "sample-10BT", "text", 9000)
    random.seed(seed); random.shuffle(texts)
    eos = tok.eos_token_id or 0
    buf = []
    for t in texts:
        buf.extend(tok(t).input_ids + [eos])
        if len(buf) >= target_tokens: break
    return torch.tensor(buf[:target_tokens], dtype=torch.long)


# ------------------------------------------------------------------ tiny model
def tiny_starcoder2(seed=0):
    """A minuscule, randomly-initialized StarCoder2 with the SAME module
    structure (Starcoder2MLP: c_fc / act / c_proj) as the 3B. Used to validate
    the capture/train/eval code paths on CPU with no download and no GPU."""
    from transformers import Starcoder2Config, AutoModelForCausalLM
    torch.manual_seed(seed)
    cfg = Starcoder2Config(
        vocab_size=256, hidden_size=64, intermediate_size=128,
        num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2,
        max_position_embeddings=256, use_bias=True, tie_word_embeddings=True,
    )
    return AutoModelForCausalLM.from_config(cfg)