"""The paper's four arms, identical except the vocabulary representation (and C's body format).

  A  learned    tied learned embedding, 262,144 x 640 (Gemma 3 default), float body
  B  frozen     frozen 512-bit codes for all 262,144 tokens (Gemma 3 270M embedding -> centre -> top-512 principal
                directions -> sign) + fixed log-unigram prior; trainable Linear 512->640 in, Linear 640->512 out;
                float body
  C  ternary    as B, body Linears = official BitNet b1.58 BitLinear (RMSNorm in, int8 absmax activations,
                ternary absmean weights, straight-through); attention math and KV in 16-bit like BitNet
  D  A-small    learned tied table like A, but the WHOLE model sized to B's parameter count: width 280,
                FFN 896 (Gemma's 3.2x ratio), same 18 layers / 4 heads / head 256 -> 99.9M (table 73.4M + body 26.5M);
                same data, steps and recipe as A. Answers: at equal size, learned table vs frozen table.
  body          Gemma-3-270M configuration, RANDOM init (trained from scratch in every arm)
  data          prep_data.py shards; seq 1024; ~262k tokens per optimiser step; AdamW, warmup + cosine
  checkpoints   every --ckpt-every steps to ckpt/<arm>/step_XXXXX.pt and (with --hf) to the HF repo
                QLNI/embeddings-are-overpriced under <arm>/step_XXXXX.pt, never overwritten; --resume continues from the latest

    python train.py --arm A|B|C|D --data data [--tokens 2.05e9] [--micro 8] [--hf]
"""
import argparse, glob, json, math, os, pathlib, time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

V, BITS, SEQ = 262144, 512, 1024


def say(*a):
    print(time.strftime("%H:%M:%S"), *a, flush=True)


def ste(x, q):
    return x + (q - x).detach()


class TLinear(nn.Linear):
    """BitLinear as in the official BitNet b1.58 training code (Microsoft, "The Era of 1-bit LLMs: Training Tips,
    Code and FAQ"): RMSNorm on the input, int8 activations with per-token absmax scale 127/max|x| clamped to
    [-128, 127], ternary weights with absmean scale, both straight-through. Attention math and KV stay 16-bit."""

    def forward(self, x):
        xf = x.float()
        xn = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-6)
        sx = 127.0 / xn.detach().abs().amax(-1, keepdim=True).clamp_min(1e-5)
        xq = ste(xn, (xn * sx).round().clamp(-128, 127) / sx)
        w = self.weight
        sw = 1.0 / w.detach().abs().mean().clamp_min(1e-5)
        wq = ste(w, (w * sw).round().clamp(-1, 1) / sw)
        return F.linear(xq.to(x.dtype), wq.to(x.dtype), self.bias)


def random_table(counts):
    g = torch.Generator().manual_seed(1234)
    table = torch.randint(0, 2, (V, BITS), generator=g).float() * 2 - 1
    seal = torch.tensor(np.log((counts + 1) / (counts.sum() + V)), dtype=torch.float32)
    return table, seal, torch.zeros(640), torch.zeros(640, BITS)


def frozen_table(counts):
    """512-bit codes from Gemma 3 270M's embedding (section 4.1 of the paper) + log-unigram prior.
    With SAVED_TABLE=1 and the saved table file present, it is loaded, so every arm and every machine use the identical bits (a rebuild on
    another machine flips ~0.008% of near-zero bits through floating-point differences in the eigendecomposition)."""
    saved = pathlib.Path(__file__).with_name("gemma3_270m_512bit_table.npz")
    if os.environ.get("SAVED_TABLE") == "1" and saved.exists():          # opt-in
        z = np.load(saved)
        table = torch.from_numpy(np.unpackbits(z["bits"], axis=1)[:, :BITS].astype(np.float32) * 2 - 1)
        seal = torch.tensor(np.log((counts + 1) / (counts.sum() + V)), dtype=torch.float32)
        print("frozen table: loaded", saved.name, flush=True)
        return table, seal, torch.from_numpy(z["mean"]), torch.from_numpy(z["basis"])
    from safetensors import safe_open
    from huggingface_hub import hf_hub_download
    with safe_open(hf_hub_download("google/gemma-3-270m", "model.safetensors"), "pt") as f:
        E = f.get_tensor([k for k in f.keys() if k.endswith("embed_tokens.weight")][0]).float()[:V]
    mu = E.mean(0)
    evals, evecs = torch.linalg.eigh((E - mu).T @ (E - mu) / (len(E) - 1))
    W = evecs[:, evals.argsort(descending=True)][:, :BITS]
    table = torch.sign((E - mu) @ W)
    table[table == 0] = 1
    seal = torch.tensor(np.log((counts + 1) / (counts.sum() + V)), dtype=torch.float32)
    return table, seal, mu, W


class Model(nn.Module):
    def __init__(self, arm, counts, attn="sdpa", learned_bias=False):
        super().__init__()
        from transformers import AutoConfig, AutoModelForCausalLM
        cfg = AutoConfig.from_pretrained("google/gemma-3-270m")
        if arm == "D":                                  # A-small: same total size as B
            cfg.hidden_size, cfg.intermediate_size = 280, 896
        torch.manual_seed(0)
        g = AutoModelForCausalLM.from_config(cfg, attn_implementation=attn, torch_dtype=torch.float32)
        self.arm = arm
        self.body = g.model
        d = cfg.hidden_size
        if arm in ("A", "D"):
            self.embed_scale = float(self.body.embed_tokens.embed_scale) if hasattr(self.body.embed_tokens, "embed_scale") else math.sqrt(d)
        else:
            self.body.embed_tokens = None
            table, seal, mu, W = frozen_table(counts)
            self.register_buffer("table", table)
            if learned_bias:            # SHADOW-250M style: per-token bias STARTS at the log-unigram seal and is trained
                self.seal = nn.Parameter(seal)
            else:                       # SHADOW-50M style: the seal is fixed (buffer, never trained)
                self.register_buffer("seal", seal)
            self.register_buffer("table_mu", mu)
            self.register_buffer("table_W", W)
            self.inp = nn.Linear(BITS, d, bias=False)
            self.out = nn.Linear(d, BITS, bias=False)
        if arm == "C":
            for _, mod in list(self.body.layers.named_modules()):
                for name, child in list(mod.named_children()):
                    if type(child) is nn.Linear:
                        t = TLinear(child.in_features, child.out_features, bias=child.bias is not None)
                        t.weight = child.weight
                        t.bias = child.bias
                        setattr(mod, name, t)

    def hidden(self, ids):
        if self.arm in ("A", "D"):
            return self.body(input_ids=ids, use_cache=False).last_hidden_state
        return self.body(inputs_embeds=self.inp(self.table[ids]), use_cache=False).last_hidden_state

    def logits(self, h):
        if self.arm in ("A", "D"):
            return h @ self.body.embed_tokens.weight.T
        return (self.out(h) @ self.table.T) / math.sqrt(BITS) + self.seal

    fused = False

    def loss(self, ids, chunk=256):
        """Mean next-token cross-entropy, readout computed in chunks of positions to bound memory.
        fused=True: Liger fused linear + cross-entropy (never materialises the full logits; same maths)."""
        h = self.hidden(ids[:, :-1])
        tgt = ids[:, 1:]
        B, T, D = h.shape
        hf, tf = h.reshape(-1, D), tgt.reshape(-1)
        if self.fused:
            from liger_kernel.transformers.functional import liger_fused_linear_cross_entropy as flce
            if self.arm in ("A", "D"):
                return flce(hf, self.body.embed_tokens.weight, tf)
            return flce(self.out(hf), self.table / math.sqrt(BITS), tf, self.seal)
        tot = 0.0
        for s in range(0, len(hf), chunk * B):
            lg = self.logits(hf[s:s + chunk * B]).float()
            tot = tot + F.cross_entropy(lg, tf[s:s + chunk * B], reduction="sum")
        return tot / len(tf)


class Data:
    def __init__(self, root, seed):
        self.shards = [np.load(p, mmap_mode="r") for p in sorted(glob.glob(str(pathlib.Path(root) / "train_*.npy")))]
        self.sizes = np.array([len(s) for s in self.shards])
        self.rng = np.random.default_rng(seed)

    def batch(self, n):
        k = self.rng.choice(len(self.shards), n, p=self.sizes / self.sizes.sum())
        out = np.empty((n, SEQ + 1), np.int64)
        for i, j in enumerate(k):
            s = self.rng.integers(0, self.sizes[j] - SEQ - 1)
            out[i] = self.shards[j][s:s + SEQ + 1]
        return torch.from_numpy(out)

    def skip(self, n):
        """advance the random stream exactly as batch(n) would, without reading data (exact resume)"""
        k = self.rng.choice(len(self.shards), n, p=self.sizes / self.sizes.sum())
        for j in k:
            self.rng.integers(0, self.sizes[j] - SEQ - 1)


@torch.no_grad()
def val_loss(m, val, n_seq=64, micro=8):
    m.eval()
    x = torch.from_numpy(val[:n_seq * (SEQ + 1)].astype(np.int64).reshape(n_seq, SEQ + 1))
    tot = 0.0
    for i in range(0, n_seq, micro):
        with torch.autocast("cuda", dtype=torch.bfloat16):
            tot += m.loss(x[i:i + micro].cuda()).item()
    m.train()
    return tot / (n_seq // micro)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True, choices=["A", "B", "C", "D"])
    ap.add_argument("--data", default="data")
    ap.add_argument("--tokens", type=float, default=2.05e9)
    ap.add_argument("--batch-tokens", type=int, default=262144)
    ap.add_argument("--micro", type=int, default=8)
    ap.add_argument("--lr", type=float, default=6e-4)
    ap.add_argument("--ckpt-every", type=int, default=800)
    ap.add_argument("--eval-every", type=int, default=200)
    ap.add_argument("--max-steps", type=int, default=0, help="stop early (speed tests)")
    ap.add_argument("--hf", action="store_true")
    ap.add_argument("--resume", action="store_true")
    ap.add_argument("--compile", action="store_true", help="torch.compile the body")
    ap.add_argument("--fused-ce", action="store_true", help="Liger fused linear cross-entropy")
    ap.add_argument("--attn", default="sdpa", help="sdpa | flash_attention_2 | kernels-community/flash-attn ...")
    ap.add_argument("--learned-bias", action="store_true",
                    help="learned per-token output bias starting at the log-unigram seal, no weight decay (SHADOW-250M style)")
    ap.add_argument("--tag", default="", help="name for checkpoints/uploads/logs (default: the arm letter), e.g. B2")
    a = ap.parse_args()
    torch.backends.cuda.matmul.allow_tf32 = True
    # Multi-GPU (torchrun): every GPU draws the SAME sequence of micro-batches and takes every world-th one, losses are
    # scaled by the global micro-batch count, gradients are summed across GPUs -> the same maths and the same data as one
    # GPU, only the work is split. Logging, validation, checkpoints and uploads happen on rank 0.
    world = int(os.environ.get("WORLD_SIZE", 1))
    rank = int(os.environ.get("RANK", 0))
    if world > 1:
        import torch.distributed as dist
        import datetime
        torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
        dist.init_process_group("nccl", timeout=datetime.timedelta(minutes=60))
    main0 = rank == 0
    counts = np.load(pathlib.Path(a.data) / "unigram_counts.npy")
    val = np.load(pathlib.Path(a.data) / "val.npy", mmap_mode="r")
    sf = pathlib.Path(__file__).with_name("bitnet_schedule.json")       # per-arm recipe (arm C only)
    arm_cfg = json.load(open(sf)) if sf.exists() else {}
    arm_cfg = arm_cfg if arm_cfg.get("arm") == a.arm else {}
    m = Model(a.arm, counts, a.attn, learned_bias=bool(arm_cfg.get("learned_bias")) or a.learned_bias).cuda()
    tag = a.tag or a.arm
    m.fused = a.fused_ce
    if a.compile:
        m.body = torch.compile(m.body)
    say(f"attention implementation: {m.body.config._attn_implementation}, fused readout+CE: {a.fused_ce}")
    n_train = sum(p.numel() for p in m.parameters() if p.requires_grad)
    accum = a.batch_tokens // (a.micro * SEQ)
    steps = int(a.tokens // (accum * a.micro * SEQ))
    decay = [p for n, p in m.named_parameters() if p.requires_grad and n != "seal"]
    groups = [{"params": decay}]
    if isinstance(getattr(m, "seal", None), nn.Parameter):                # learned bias: never weight-decayed (as SHADOW-250M)
        groups.append({"params": [m.seal], "weight_decay": 0.0, "no_wd": True})
    opt = torch.optim.AdamW(groups, lr=a.lr, betas=(0.9, 0.95), weight_decay=0.1)
    ck = pathlib.Path("ckpt") / tag
    ck.mkdir(parents=True, exist_ok=True)
    start = 0
    if a.resume and sorted(ck.glob("step_*.pt")):
        last = sorted(ck.glob("step_*.pt"))[-1]
        st = torch.load(last, map_location="cuda", weights_only=False)
        sd = st["model"]
        if a.compile:
            sd = {(k.replace("body.", "body._orig_mod.", 1) if k.startswith("body.") else k): v for k, v in sd.items()}
        m.load_state_dict(sd)
        opt.load_state_dict(st["opt"])
        start = st["step"]
        say(f"resumed from {last}")
    # exact resume: the data stream is the one an uninterrupted run would draw - seed 1000, fast-forwarded past the
    # micro-batches of the steps already done - so a resumed arm sees exactly the same batches as an uninterrupted one
    data = Data(a.data, seed=1000)
    for _ in range(start * accum):
        data.skip(a.micro)
    log = open(ck / "log.jsonl", "a") if main0 else None
    if main0:
        say(f"arm {tag}: {n_train:,} trainable parameters, {steps} steps of {accum * a.micro * SEQ:,} tokens, {world} GPU(s)")
    warm = max(1, steps // 100)
    # BitNet b1.58 published recipe (Microsoft, "The Era of 1-bit LLMs: Training Tips, Code and FAQ", Table 2, 700M):
    # two-stage LR 1.5e-3 -> 1e-3 (linear decay; the peak drops at the midpoint), weight decay 0.1 -> 0 at the midpoint,
    # 375 warm-up steps. Applied only to the arm named in bitnet_schedule.json next to this file (arm C).
    sched = arm_cfg or None
    if sched and main0:
        say(f"BitNet schedule for arm {a.arm}: {sched}")

    def lr_wd(step):
        if sched is None:
            return a.lr * min(1, step / warm) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * step / steps))), 0.1
        half = int(steps * sched["switch_frac"])
        base = min(1, step / sched["warmup"]) * max(0.0, 1 - step / steps)      # warm-up, then linear decay to 0
        if step <= half:
            return sched["peak1"] * base, sched["wd1"]
        return sched["peak2"] * base, sched["wd2"]

    t0 = time.time()
    params = [p for p in m.parameters() if p.requires_grad]
    uploads = []
    for step in range(start + 1, steps + 1):
        lr, wd = lr_wd(step)
        for g in opt.param_groups:
            g["lr"] = lr
            g["weight_decay"] = 0.0 if g.get("no_wd") else wd
        tl = 0.0
        for i in range(accum):
            x = data.batch(a.micro)                     # every GPU draws every micro-batch: identical data stream
            if i % world != rank:
                continue
            with torch.autocast("cuda", dtype=torch.bfloat16):
                loss = m.loss(x.cuda(non_blocking=True)) / accum
            loss.backward()
            tl += loss.item()
        if world > 1:
            grads = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
            flat = torch._utils._flatten_dense_tensors(grads)
            dist.all_reduce(flat)
            for p, g in zip(params, torch._utils._unflatten_dense_tensors(flat, grads)):
                p.grad = g
            t = torch.tensor([tl], device="cuda"); dist.all_reduce(t); tl = t.item()
        gn = torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0).item()
        opt.step()
        opt.zero_grad(set_to_none=True)
        rec = {"step": step, "tokens": step * accum * a.micro * SEQ, "loss": tl, "lr": lr, "wd": wd, "grad_norm": gn,
               "tok_s": (step - start) * accum * a.micro * SEQ / (time.time() - t0)}
        if main0 and (step % a.eval_every == 0 or step == steps):
            rec["val_loss"] = val_loss(m, val)
        if main0:
            log.write(json.dumps(rec) + "\n")
            log.flush()
        if main0 and (step % 20 == 0 or "val_loss" in rec):
            say(" ".join(f"{k} {v:.4g}" if isinstance(v, float) else f"{k} {v}" for k, v in rec.items()))
        if main0 and (step % a.ckpt_every == 0 or step == steps):
            p = ck / f"step_{step:05d}.pt"
            sd = m.state_dict()
            if a.compile:                                # keep checkpoint keys identical to an uncompiled model
                sd = {k.replace("body._orig_mod.", "body."): v for k, v in sd.items()}
            torch.save({"model": sd, "opt": opt.state_dict(), "step": step, "arm": a.arm, "tag": tag, "args": vars(a)}, p)
            if a.hf:                                     # upload in the background: training does not wait for the network
                import threading, shutil
                from huggingface_hub import HfApi
                lp = ck / f"log_step_{step:05d}.jsonl"
                shutil.copy(ck / "log.jsonl", lp)
                def up(p=p, lp=lp, step=step):
                    for attempt in range(5):
                        try:
                            HfApi().upload_file(path_or_fileobj=str(p), path_in_repo=f"{tag}/{p.name}", repo_id="QLNI/embeddings-are-overpriced",
                                                commit_message=f"arm {tag} step {step}")
                            HfApi().upload_file(path_or_fileobj=str(lp), path_in_repo=f"{tag}/{lp.name}",
                                                repo_id="QLNI/embeddings-are-overpriced", commit_message=f"arm {tag} log {step}")
                            say(f"uploaded {p.name}")
                            return
                        except Exception as e:
                            say(f"upload {p.name} failed (attempt {attempt + 1}): {e}")
                            time.sleep(60)
                th = threading.Thread(target=up); th.start(); uploads.append(th)
            say(f"checkpoint {p}")
        if a.max_steps and step - start >= a.max_steps:
            peak = torch.cuda.max_memory_allocated() / 2**30
            if main0:
                say(f"stopping after {a.max_steps} steps (speed test): {rec['tok_s']:.0f} tokens/s, peak GPU memory {peak:.1f} GiB")
                with open(ck / "speed.json", "w") as f:
                    json.dump({"arm": a.arm, "steps": a.max_steps, "tokens_per_step": accum * a.micro * SEQ, "micro": a.micro,
                               "tok_s": rec["tok_s"], "peak_gib": peak, "trainable": n_train, "gpus": world,
                               "gpu": torch.cuda.get_device_name(0), "attn": a.attn, "fused_ce": a.fused_ce, "compile": a.compile}, f, indent=1)
            break
    for th in uploads:
        th.join()
    if world > 1:
        dist.barrier()
        dist.destroy_process_group()


if __name__ == "__main__":
    main()