#!/usr/bin/env python3
"""
Native KOLM vs transformer twin — rung 1 of the scaling ladder.

Two ~25M models sharing one skeleton (tied embeddings, learned positions,
pre-LN, causal attention in BOTH — attention does routing), differing only
in the per-token processing block:

    --arch transformer   attention + MLP (d -> 4d -> d)
    --arch kolm          attention + KuramotoBlock (kuramoto_torch.py)

Same tokenizer, same data, same step budget => the val-loss curves are a
fair architecture comparison. Trained from scratch on TinyStories: this is
the native test of "dynamics shaped from step one", not a retrofit.

    .venv/bin/python native_kolm.py --prep                 # tokenizer + memmap
    .venv/bin/python native_kolm.py --arch kolm --steps N
    .venv/bin/python native_kolm.py --arch transformer --steps N
"""

import argparse
import math
import os
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from kuramoto_torch import KuramotoBlock

DEV = "mps" if torch.backends.mps.is_available() else "cpu"
TOK_JSON = "tiny8k.json"
BIN = "tiny_train.bin"
VAL_BIN = "tiny_val.bin"


# ---------------------------------------------------------------- data prep
def prep(vocab_size=8192, max_bytes=500_000_000):
    from huggingface_hub import hf_hub_download
    from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
    path = hf_hub_download(repo_id="roneneldan/TinyStories", repo_type="dataset",
                           filename="TinyStoriesV2-GPT4-train.txt")
    text = open(path, encoding="utf-8", errors="ignore").read(max_bytes)

    tok = Tokenizer(models.BPE(unk_token=None))
    tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True)
    tok.decoder = decoders.ByteLevel()
    trainer = trainers.BpeTrainer(vocab_size=vocab_size, special_tokens=[])
    step = 1_000_000
    tok.train_from_iterator((text[i:i + step] for i in range(0, len(text), step)),
                            trainer=trainer)
    tok.save(TOK_JSON)
    print(f"tokenizer: {tok.get_vocab_size()} merges saved to {TOK_JSON}", flush=True)

    parts, total = [], 0
    for i in range(0, len(text), step):
        parts.append(np.array(tok.encode(text[i:i + step]).ids, dtype=np.uint16))
        total += len(parts[-1])
        if i % 50_000_000 < step:
            print(f"  tokenized {i / 1e6:.0f}MB -> {total / 1e6:.1f}M tokens",
                  flush=True)
    arr = np.concatenate(parts)
    n_val = 2_000_000
    arr[:-n_val].tofile(BIN)
    arr[-n_val:].tofile(VAL_BIN)
    print(f"train {len(arr) - n_val:,} tokens -> {BIN} | val {n_val:,} -> {VAL_BIN}",
          flush=True)


# ---------------------------------------------------------------- model
class Block(nn.Module):
    def __init__(self, d, n_head, arch, osc_h=320, frustrated=False,
                 grad_steps=0):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, n_head, batch_first=True)
        self.ln2 = nn.LayerNorm(d)
        if arch == "kolm":
            # KuramotoBlock is residual + zero-init internally
            self.ffn = KuramotoBlock(d, H=osc_h, N=4, groups=32, steps=4,
                                     frustrated=frustrated,
                                     grad_steps=grad_steps)
            self.residual_ffn = False
        else:
            self.ffn = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(),
                                     nn.Linear(4 * d, d))
            self.residual_ffn = True

    def forward(self, x, mask):
        h = self.ln1(x)
        a, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
        x = x + a
        h = self.ln2(x)
        return x + self.ffn(h) - h if not self.residual_ffn else x + self.ffn(h)


class TinyLM(nn.Module):
    def __init__(self, vocab, d=384, n_layer=8, n_head=6, ctx=512, arch="kolm",
                 osc_h=320, frustrated=False, grad_steps=0):
        super().__init__()
        self.ctx = ctx
        self.emb = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(ctx, d)
        self.blocks = nn.ModuleList(
            Block(d, n_head, arch, osc_h, frustrated, grad_steps)
            for _ in range(n_layer))
        self.ln_f = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        self.head.weight = self.emb.weight        # tied
        mask = torch.triu(torch.full((ctx, ctx), float("-inf")), diagonal=1)
        self.register_buffer("mask", mask)
        self.apply(self._init)
        for b in self.blocks:
            if not b.residual_ffn:      # kolm: restore identity-at-init
                nn.init.zeros_(b.ffn.out.weight)
                nn.init.zeros_(b.ffn.out.bias)

    @staticmethod
    def _init(m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.zeros_(m.bias)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        x = self.emb(idx) + self.pos.weight[:T]
        m = self.mask[:T, :T]
        for b in self.blocks:
            x = b(x, m)
        logits = self.head(self.ln_f(x))
        if targets is None:
            return logits, None
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss


# ---------------------------------------------------------------- training
def batches(bin_path, B, T, seed=0):
    data = np.memmap(bin_path, dtype=np.uint16, mode="r")
    rng = np.random.default_rng(seed)
    while True:
        ix = rng.integers(0, len(data) - T - 1, B)
        x = np.stack([data[i:i + T] for i in ix]).astype(np.int64)
        y = np.stack([data[i + 1:i + T + 1] for i in ix]).astype(np.int64)
        yield torch.from_numpy(x), torch.from_numpy(y)


@torch.no_grad()
def val_loss(model, B, T, iters=20, path=VAL_BIN):
    model.eval()
    it = batches(path, B, T, seed=1)
    tot = 0.0
    for _ in range(iters):
        x, y = next(it)
        _, loss = model(x.to(DEV), y.to(DEV))
        tot += loss.item()
    model.train()
    return tot / iters


def main():
    ap = argparse.ArgumentParser(description="native KOLM / transformer trainer")
    ap.add_argument("--prep", action="store_true")
    ap.add_argument("--arch", choices=["kolm", "transformer"], default="kolm")
    ap.add_argument("--steps", type=int, default=2000)
    ap.add_argument("--batch", type=int, default=24)
    ap.add_argument("--ctx", type=int, default=512)
    ap.add_argument("--lr", type=float, default=6e-4)
    ap.add_argument("--warmup", type=int, default=200)
    ap.add_argument("--val-every", type=int, default=250)
    ap.add_argument("--d-model", type=int, default=384)
    ap.add_argument("--n-layer", type=int, default=8)
    ap.add_argument("--n-head", type=int, default=6)
    ap.add_argument("--osc-h", type=int, default=320)
    ap.add_argument("--frustrated", action="store_true")
    ap.add_argument("--grad-steps", type=int, default=0,
                    help="settle steps to backprop through (0 = all)")
    ap.add_argument("--init", default=None, help="warm-start state dict")
    ap.add_argument("--resume", default=None, help="checkpoint to resume")
    ap.add_argument("--ckpt-every", type=int, default=500)
    ap.add_argument("--train-bin", default=BIN)
    ap.add_argument("--val-bin", default=VAL_BIN)
    ap.add_argument("--name", default=None, help="run name for curve/save files")
    args = ap.parse_args()

    if args.prep:
        prep()
        return

    from tokenizers import Tokenizer
    vocab = Tokenizer.from_file(TOK_JSON).get_vocab_size()
    torch.manual_seed(0)
    name = args.name or args.arch
    model = TinyLM(vocab, d=args.d_model, n_layer=args.n_layer,
                   n_head=args.n_head, ctx=args.ctx, arch=args.arch,
                   osc_h=args.osc_h, frustrated=args.frustrated,
                   grad_steps=args.grad_steps).to(DEV)
    if args.init:
        model.load_state_dict(torch.load(args.init, map_location=DEV))
    n = sum(p.numel() for p in model.parameters())
    print(f"{name} ({args.arch}) | {n:,} params | vocab {vocab} | {DEV}",
          flush=True)

    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.1,
                            betas=(0.9, 0.95))
    start = 1
    if args.resume:
        ck = torch.load(args.resume, map_location=DEV)
        model.load_state_dict(ck["model"])
        opt.load_state_dict(ck["opt"])
        start = ck["step"] + 1
        print(f"resumed from {args.resume} at step {start}", flush=True)
    sched = lambda s: min(s / args.warmup, 1.0) * \
        (0.5 * (1 + math.cos(math.pi * s / args.steps)))
    it = batches(args.train_bin, args.batch, args.ctx)
    csv = open(f"curve_{name}.csv", "a")
    model.train()
    t0 = time.time()
    for step in range(start, args.steps + 1):
        for g in opt.param_groups:
            g["lr"] = args.lr * sched(step)
        x, y = next(it)
        _, loss = model(x.to(DEV), y.to(DEV))
        opt.zero_grad()
        loss.backward()
        gn = nn.utils.clip_grad_norm_(model.parameters(), 1.0).item()
        opt.step()
        if step % 25 == 0 or step == 1 or not math.isfinite(gn):
            li = loss.item()
            print(f"step {step:5d} | loss {li:.4f} | gnorm {gn:.2f} | "
                  f"{(time.time() - t0):.0f}s", flush=True)
            if not math.isfinite(gn):
                print(f"ABORT: non-finite grad norm at step {step}", flush=True)
                return
            if not math.isfinite(li):
                print("ABORT: non-finite loss", flush=True)
                return
        if step % args.val_every == 0 or step == args.steps:
            vl = val_loss(model, args.batch, args.ctx, path=args.val_bin)
            toks = step * args.batch * args.ctx
            print(f"  val {vl:.4f} | ppl {math.exp(vl):.2f} | {toks / 1e6:.0f}M tokens",
                  flush=True)
            csv.write(f"{step},{toks},{vl}\n")
            csv.flush()
        if args.ckpt_every and step % args.ckpt_every == 0:
            torch.save({"model": model.state_dict(),
                        "opt": opt.state_dict(), "step": step},
                       f"ckpt_{name}.pt")
    torch.save(model.state_dict(), f"native_{name}.pt")
    print(f"saved native_{name}.pt", flush=True)


if __name__ == "__main__":
    main()