"""
Train HOARD (or a matched baseline) on one Mac.

    python train.py --config hoard_small --data data/fineweb --steps 3000 --batch_tokens 65536
    python train.py --config transformer_small --data data/fineweb --steps 3000 --batch_tokens 65536

Optimiser: Muon for 2-D hidden matrices, AdamW for embeddings / head / norms /
gates / hoard block weights. LR: linear warmup, flat, linear cooldown.
"""
import argparse, json, math, os, sys, time
import numpy as np
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from mlx.utils import tree_flatten

import configs
from model import HOARD, HoardConfig


# --------------------------------------------------------------------------- #
class TokenLoader:
    def __init__(self, path, seq_len, batch, seed=0):
        self.data = np.memmap(path, dtype=np.uint16, mode="r")
        self.T, self.B = seq_len, batch
        self.rng = np.random.default_rng(seed)
        assert len(self.data) > seq_len + 1, "dataset too small for seq_len"

    def next(self):
        ix = self.rng.integers(0, len(self.data) - self.T - 1, size=self.B)
        x = np.stack([self.data[i: i + self.T + 1] for i in ix]).astype(np.int32)
        return mx.array(x)

    def sequential(self, n_batches):
        """Deterministic contiguous batches for validation. Shrinks the batch
        when the split is smaller than one full B x (T+1) span (tiny datasets
        like shakespeare's 6.7K-token val split) instead of yielding nothing."""
        B = self.B
        if B * (self.T + 1) > len(self.data):
            B = max(1, len(self.data) // (self.T + 1))
        span = B * (self.T + 1)
        for j in range(n_batches):
            st = j * span
            if st + span > len(self.data):
                break
            x = np.asarray(self.data[st: st + span]).astype(np.int32).reshape(B, self.T + 1)
            yield mx.array(x)


def lr_at(step, total, base, warmup, cooldown_frac):
    if step < warmup:
        return base * (step + 1) / warmup
    cool_start = int(total * (1 - cooldown_frac))
    if step >= cool_start:
        return base * max(0.0, (total - step) / max(1, total - cool_start))
    return base


def fp32_master(opt_cls):
    """Wrap an MLX optimizer so bf16 parameters train with fp32 master weights
    and fp32 optimizer state. bf16-only AdamW updates diverge (measured: nan by
    step 10 on shakespeare); compute stays bf16, the update math moves to fp32.
    The master copy is created lazily in apply_single because Optimizer.init
    receives the gradient tree, not the parameters."""

    class FP32Master(opt_cls):
        def apply_single(self, gradient, parameter, state):
            if parameter.dtype == mx.float32:
                return super().apply_single(gradient, parameter, state)
            if "master" not in state:
                state["master"] = parameter.astype(mx.float32)
                for k in list(state):
                    if k != "master" and isinstance(state[k], mx.array):
                        state[k] = state[k].astype(mx.float32)
            new = super().apply_single(gradient.astype(mx.float32), state["master"], state)
            state["master"] = new
            return new.astype(parameter.dtype)

    return FP32Master


class NorMuonLite(optim.Muon):
    """Muon plus per-row second-moment normalization of the orthogonalized
    update (NorMuon-style), scale-preserving: rows are rebalanced, the global
    update RMS is unchanged so learning-rate semantics carry over."""

    def apply_single(self, gradient, parameter, state):
        if parameter.ndim < 2:
            return super().apply_single(gradient, parameter, state)
        if self.weight_decay != 0:
            gradient = gradient + self.weight_decay * parameter
        v = self.momentum * state["v"] + (1 - self.momentum) * gradient
        state["v"] = v
        update = gradient * (1 - self.momentum) + v * self.momentum if self.nesterov else v
        original_shape = update.shape
        reshape_needed = update.ndim > 2
        if reshape_needed:
            update = mx.reshape(update, (update.shape[0], -1))
        update = self._zeropower_via_newtonschulz5(update, steps=self.ns_steps)
        row_ms = mx.mean(mx.square(update), axis=-1, keepdims=True)
        vn = state.get("vn")
        vn = row_ms if vn is None else 0.95 * vn + 0.05 * row_ms
        state["vn"] = vn
        pre_rms = mx.sqrt(mx.mean(mx.square(update)) + 1e-12)
        update = update / (mx.sqrt(vn) + 1e-8)
        update = update * (pre_rms / (mx.sqrt(mx.mean(mx.square(update))) + 1e-12))
        if reshape_needed:
            update = mx.reshape(update, original_shape)
        lr = self.learning_rate.astype(gradient.dtype)
        lr *= max(1, update.shape[-2] / update.shape[-1]) ** 0.5
        return parameter - lr * update


def is_hoard_param(path: str, p: mx.array) -> bool:
    return "W_up" in path or "W_down" in path


def is_muon_param(path: str, p: mx.array) -> bool:
    if p.ndim != 2:
        return False
    for bad in ("embed", "head", "W_up", "W_down", "norm", "loop_embed", "conv"):
        if bad in path:
            return False
    return True


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="hoard_small")
    ap.add_argument("--data", default="data/fineweb")
    ap.add_argument("--out", default=None)
    ap.add_argument("--steps", type=int, default=2000)
    ap.add_argument("--seq_len", type=int, default=1024)
    ap.add_argument("--batch_tokens", type=int, default=65536)
    ap.add_argument("--lr", type=float, default=3e-3, help="AdamW lr (embeddings etc.)")
    ap.add_argument("--muon_lr", type=float, default=0.02)
    ap.add_argument("--wd", type=float, default=0.0)
    ap.add_argument("--warmup", type=int, default=100)
    ap.add_argument("--cooldown_frac", type=float, default=0.3)
    ap.add_argument("--grad_clip", type=float, default=1.0)
    ap.add_argument("--eval_every", type=int, default=200)
    ap.add_argument("--eval_batches", type=int, default=8)
    ap.add_argument("--log_every", type=int, default=10)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--loops", type=int, default=None, help="override n_loops")
    ap.add_argument("--loop_schedule", default=None, choices=[None, "fixed", "random"])
    ap.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"])
    ap.add_argument("--resume", default=None)
    ap.add_argument("--vocab", type=int, default=None, help="override vocab_size (e.g. 4096 for synthetic data)")
    ap.add_argument("--tie", action="store_true", help="tie input/output embeddings")
    ap.add_argument("--window", type=int, default=None, help="override attention window")
    ap.add_argument("--gate", action="store_true", help="silu-gated hoard output")
    ap.add_argument("--hoard_lr", type=float, default=None,
                    help="separate AdamW lr for hoard W_up/W_down (memory-layer recipe)")
    ap.add_argument("--normuon", action="store_true", help="NorMuon-lite instead of Muon")
    ap.add_argument("--ckpt", action="store_true", help="gradient-checkpoint each cell")
    ap.add_argument("--compile_step", action="store_true", help="mx.compile the fwd+bwd step")
    a = ap.parse_args()

    mx.random.seed(a.seed)
    cfg: HoardConfig = configs.get(a.config)
    if a.loops:
        cfg.n_loops = a.loops
    if a.loop_schedule:
        cfg.loop_schedule = a.loop_schedule
    if a.vocab:
        cfg.vocab_size = a.vocab
    if a.tie:
        cfg.tie_embeddings = True
    if a.window:
        cfg.window = a.window
    if a.gate:
        cfg.hoard_gate = True
    out = a.out or os.path.join("runs", f"{a.config}_{int(time.time())}")
    os.makedirs(out, exist_ok=True)
    json.dump({"config": cfg.to_dict(), "args": vars(a)}, open(os.path.join(out, "config.json"), "w"), indent=2)

    model = HOARD(cfg)
    if a.dtype == "bfloat16":
        # parameters in bf16; the GDN core casts to fp32 internally, state is fp32
        model.set_dtype(mx.bfloat16)
    if a.resume:
        model.load_weights(a.resume)
    counts = model.param_counts()
    print(f"[{a.config}] params: total {counts['total']/1e6:.1f}M | hoard {counts['hoard']/1e6:.1f}M | "
          f"active/token/loop {counts['active_per_token_per_loop']/1e6:.1f}M | active x loops {counts['active_x_loops']/1e6:.1f}M")

    B = max(1, a.batch_tokens // a.seq_len)
    train_loader = TokenLoader(os.path.join(a.data, "train.bin"), a.seq_len, B, a.seed)
    val_loader = TokenLoader(os.path.join(a.data, "val.bin"), a.seq_len, B, a.seed)

    muon_cls = NorMuonLite if a.normuon else optim.Muon
    muon = fp32_master(muon_cls)(learning_rate=a.muon_lr, weight_decay=a.wd)
    adamw = fp32_master(optim.AdamW)(learning_rate=a.lr, betas=(0.9, 0.95), weight_decay=a.wd)
    sched = [(muon, a.muon_lr), (adamw, a.lr)]
    if a.hoard_lr:
        hoard_opt = fp32_master(optim.AdamW)(learning_rate=a.hoard_lr, betas=(0.9, 0.95), weight_decay=a.wd)
        sched.insert(1, (hoard_opt, a.hoard_lr))
        opt = optim.MultiOptimizer([muon, hoard_opt, adamw], [is_muon_param, is_hoard_param])
    else:
        opt = optim.MultiOptimizer([muon, adamw], [is_muon_param])

    if a.ckpt:
        from mlx.nn.utils import checkpoint as _checkpoint
        model._ckpt_cells = [_checkpoint(c) for c in model.cells]

    def loss_fn(model, batch, n_loops):
        x, y = batch[:, :-1], batch[:, 1:]
        logits = model(x, n_loops=n_loops)
        ce = nn.losses.cross_entropy(logits.astype(mx.float32), y, reduction="mean")
        return ce + model.balance_loss(), ce

    grad_fn = nn.value_and_grad(model, loss_fn)

    compiled = None
    if a.compile_step and cfg.loop_schedule == "fixed":
        _state = [model.state, mx.random.state]

        def _compute(batch):
            return grad_fn(model, batch, cfg.n_loops)

        compiled = mx.compile(_compute, inputs=_state, outputs=_state)

    def train_step(batch, n_loops):
        if compiled is not None and n_loops == cfg.n_loops:
            (loss, ce), grads = compiled(batch)
        else:
            (loss, ce), grads = grad_fn(model, batch, n_loops)
        if a.grad_clip > 0:
            grads, _ = optim.clip_grad_norm(grads, a.grad_clip)
        opt.update(model, grads)
        return loss, ce

    def evaluate(n_loops):
        model.eval()
        losses = []
        for batch in val_loader.sequential(a.eval_batches):
            _, ce = loss_fn(model, batch, n_loops)
            losses.append(ce.item())
        model.train()
        return float(np.mean(losses)) if losses else float("nan")

    log = open(os.path.join(out, "log.txt"), "a")
    t0 = time.time(); tok_seen = 0
    for step in range(a.steps):
        lr_scale = lr_at(step, a.steps, 1.0, a.warmup, a.cooldown_frac)
        for o, base_lr in sched:
            o.learning_rate = base_lr * lr_scale
        n_loops = cfg.n_loops
        if cfg.loop_schedule == "random":
            n_loops = int(np.random.randint(1, cfg.n_loops + 1))
        batch = train_loader.next()
        loss, ce = train_step(batch, n_loops)
        mx.eval(model.parameters(), opt.state, loss)
        tok_seen += B * a.seq_len
        if step % a.log_every == 0 or step == a.steps - 1:
            dt = time.time() - t0
            msg = (f"step {step:5d} | loss {loss.item():.4f} ce {ce.item():.4f} | loops {n_loops} | "
                   f"lr x{lr_scale:.2f} | {tok_seen/dt:,.0f} tok/s | {dt/60:.1f} min")
            print(msg); log.write(msg + "\n"); log.flush()
        if (step + 1) % a.eval_every == 0 or step == a.steps - 1:
            vl = evaluate(cfg.n_loops)
            msg = f"  val loss {vl:.4f} @ step {step+1} ({tok_seen/1e6:.1f}M tokens)"
            print(msg); log.write(msg + "\n"); log.flush()
            model.save_weights(os.path.join(out, "model.safetensors"))
    print(f"done. checkpoint in {out}")


if __name__ == "__main__":
    main()