# train_peer.py - from-scratch pretrain of the TERNARY + PEER DNA controller.
# Reuses the fast DNA machinery: chunk-parallel recurrence (torch.compile), codon
# memory + mutation (ZERO-SHADOW), Muon(2D) + fused AdamW + WSD + MTP aux head.
import argparse, json, math, os, time, threading, queue
from pathlib import Path
import numpy as np, torch, torch.nn.functional as F
import torch.utils.checkpoint as cp
import sys; sys.path.insert(0, '/root/dna')
from model_dna_peer import DnaPeer, peer_counts

ap = argparse.ArgumentParser()
ap.add_argument('--data-dir', default='/root/dna/data')
ap.add_argument('--out', default='/root/dna/ckpt-peer'); ap.add_argument('--resume', default='/root/dna/ckpt-peer/resume.pt')
ap.add_argument('--batch', type=int, default=16); ap.add_argument('--seq', type=int, default=512)
ap.add_argument('--steps', type=int, default=2000000); ap.add_argument('--target-tokens', type=float, default=3e9)
ap.add_argument('--lr', type=float, default=3e-4); ap.add_argument('--muonlr', type=float, default=2e-2)
ap.add_argument('--rowstep', type=float, default=1.0); ap.add_argument('--warmup', type=int, default=2000)
ap.add_argument('--ckpt-every', type=int, default=2000); ap.add_argument('--hf-every', type=int, default=40000)
ap.add_argument('--hf-repo', default='jaivial/dna-diskchat-2b-peer-v1'); ap.add_argument('--log-every', type=int, default=100)
ap.add_argument('--mtp-frac', type=float, default=0.25)
a = ap.parse_args(); dev = 'cuda'; Path(a.out).mkdir(parents=True, exist_ok=True); torch.manual_seed(1)
torch._dynamo.config.cache_size_limit = 64
torch.backends.cuda.matmul.allow_tf32 = True; torch.backends.cudnn.allow_tf32 = True; torch.backends.cudnn.benchmark = True

m = DnaPeer(chunk=int(os.environ.get('CHUNK_C', '16'))).to(dev)
muon_p = []
for blk in m.blocks: muon_p += [blk.proj.weight, blk.o.weight]
mset = set(id(p) for p in muon_p)
adam_p = [p for p in m.parameters() if id(p) not in mset]
for i in range(len(m.blocks)): m.blocks[i] = torch.compile(m.blocks[i])
ctrl, active = peer_counts(m)
print('PEER_CONFIG', json.dumps({'ctrl_params': ctrl, 'active_params_per_token': active,
      'codon_table_bytes': m.nrows * m.ncod, 'experts_per_block': m.nk * m.nk, 'batch': a.batch, 'seq': a.seq}), flush=True)

def ns5(Gm, steps=int(os.environ.get('NS_STEPS', '5'))):
    x = 3.4445; y = -4.7750; z = 2.0315; X = Gm.bfloat16(); tp = Gm.size(0) > Gm.size(1)
    if tp: X = X.t()
    X = X / (X.norm() + 1e-7)
    for _ in range(steps):
        A = X @ X.t(); B = y * A + z * (A @ A); X = x * X + B @ X
    if tp: X = X.t()
    return X.to(Gm.dtype)
class Muon(torch.optim.Optimizer):
    def __init__(self, params, lr=2e-2, momentum=0.95): super().__init__(list(params), dict(lr=lr, momentum=momentum))
    @torch.no_grad()
    def step(self):
        for grp in self.param_groups:
            for p in grp['params']:
                if p.grad is None: continue
                st = self.state[p]
                if 'm' not in st: st['m'] = torch.zeros_like(p.grad)
                buf = st['m']; buf.mul_(grp['momentum']).add_(p.grad); u = p.grad.add(buf, alpha=grp['momentum'])
                u = ns5(u); u = u.mul(max(1.0, p.size(0) / p.size(1)) ** 0.5); p.add_(u, alpha=-grp['lr'])
opt = torch.optim.AdamW(adam_p, lr=a.lr, betas=(.9, .95), weight_decay=.1, fused=True)
muon = Muon(muon_p, lr=a.muonlr); dense = muon_p + adam_p
print('MUON', sum(p.numel() for p in muon_p), 'ADAM', sum(p.numel() for p in adam_p), flush=True)

cur = None; curi = -1
def shard_list(): return sorted([p for p in Path(a.data_dir).glob('shard_*.u16') if not p.with_suffix('.tmp').exists()])
def batch():
    global cur, curi
    sh = shard_list()
    if not sh: raise SystemExit('no shards yet - start stream_data.py first')
    i = np.random.randint(0, len(sh))
    if i != curi: cur = np.memmap(sh[i], np.uint16, 'r'); curi = i
    n = len(cur) - a.seq - 1; ix = np.random.randint(0, n, size=a.batch)
    return torch.from_numpy(np.stack([np.asarray(cur[j:j+a.seq+1], np.int64) for j in ix])).to(dev)
def cosf(s):
    if s < a.warmup: return s / a.warmup
    ds = int(a.steps * 0.8)
    if s < ds: return 1.0
    p = (s - ds) / max(1, a.steps - ds); return 1.0 - 0.9 * p
def feats(ids):
    B, T = ids.shape; x = m.embed(ids)
    for blk in m.blocks: x = cp.checkpoint(blk, x, m.chunk, use_reentrant=False)
    mem, bal, rows = m.route(x.reshape(B * T, m.d)); return m.norm(x + mem.reshape(B, T, m.d)), bal, rows
def save(step):
    tmp = a.resume + '.tmp'
    torch.save({'model': m.state_dict(), 'opt': opt.state_dict(), 'muon': muon.state_dict(), 'step': step, 'config': m.config()}, tmp)
    os.replace(tmp, a.resume)
def hf_push(step):
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=os.environ.get('HUGGING_FACE')); api.create_repo(a.hf_repo, exist_ok=True)
        p = f'{a.out}/ckpt_{step}.pt'
        torch.save({'model': m.state_dict(), 'config': m.config(), 'step': step}, p)
        api.upload_file(path_or_fileobj=p, path_in_repo=f'checkpoints/ckpt_{step}.pt', repo_id=a.hf_repo)
        os.remove(p); print(f'HF_PUSH step {step}', flush=True)
    except Exception as e: print('HF_PUSH_FAIL', str(e)[:140], flush=True)

@torch.no_grad()
def valid_loss():
    m.eval(); ids = batch()
    with torch.autocast('cuda', dtype=torch.bfloat16):
        feat, _, _ = feats(ids)
        ce = F.cross_entropy(F.linear(feat[:, :-1].reshape(-1, m.d), m.embed.weight), ids[:, 1:].reshape(-1))
    m.train(); return ce.item()

start = 0
if os.path.exists(a.resume):
    ck = torch.load(a.resume, map_location=dev, weights_only=False)
    m.load_state_dict(ck['model']); opt.load_state_dict(ck['opt']); muon.load_state_dict(ck['muon']); start = int(ck['step'])
    print('RESUME', start, flush=True)

_bq = queue.Queue(maxsize=4)
def _worker():
    while True: _bq.put(batch())
for _ in range(2): threading.Thread(target=_worker, daemon=True).start()
def nextbatch(): return _bq.get()

m.train(); t0 = time.time(); seen = 0; step = start
print('LOOP start', start, 'steps', a.steps, flush=True)
for step in range(start + 1, a.steps + 1):
    f = cosf(step)
    for g in opt.param_groups: g['lr'] = a.lr * f
    for g in muon.param_groups: g['lr'] = a.muonlr * f
    ids = nextbatch(); opt.zero_grad(set_to_none=True)
    with torch.autocast('cuda', dtype=torch.bfloat16):
        feat, bal, rows = feats(ids)
        fl = feat[:, :-1].reshape(-1, m.d); tg = ids[:, 1:].reshape(-1)
        l1 = F.cross_entropy(F.linear(fl, m.embed.weight), tg)
        f2 = feat[:, :-2].reshape(-1, m.d); t2 = ids[:, 2:].reshape(-1)
        sel = torch.randint(0, f2.size(0), (int(a.mtp_frac * f2.size(0)),), device=dev)
        l2 = F.cross_entropy(F.linear(f2[sel], m.head2.weight), t2[sel])
        loss = l1 + 0.5 * l2 + 1e-3 * bal
    loss.backward(); torch.nn.utils.clip_grad_norm_(dense, 1.0); opt.step()
    if step % int(os.environ.get('MUON_EVERY', '2')) == 0: muon.step()
    with torch.no_grad():
        leaf = getattr(m, '_leaf', None)
        if leaf is not None and leaf.grad is not None:
            flat = rows.reshape(-1); g = leaf.grad
            gn = g.norm(dim=-1, keepdim=True).clamp_min(1e-8); qstep = m.codebook.detach().std()
            rstep = a.rowstep * qstep * f * (m.d ** 0.5); newrows = leaf.detach() - rstep * (g / gn)
            uid, inv = torch.unique(flat, return_inverse=True)
            acc = torch.zeros(uid.size(0), m.d, device=newrows.device, dtype=newrows.dtype); acc.index_copy_(0, inv, newrows)
            if not hasattr(m, '_pu'): m._pu = []; m._pa = []
            m._pu.append(uid); m._pa.append(acc)
            if step % int(os.environ.get('MUT_EVERY', '16')) == 0:
                pu = torch.cat(m._pu); pa = torch.cat(m._pa)
                u2, inv2 = torch.unique(pu, return_inverse=True)
                a2 = torch.zeros(u2.size(0), m.d, device=pa.device, dtype=pa.dtype); a2.index_copy_(0, inv2, pa)
                before = m.codons[u2].clone(); m.mutate_rows(u2, a2)
                m._mut = (m.codons[u2] != before).float().mean().item(); m._pu = []; m._pa = []
            m._leaf = None
    seen += a.batch * (a.seq - 1); toks = step * a.batch * (a.seq - 1)
    if step % a.log_every == 0:
        torch.cuda.synchronize(); dt = time.time() - t0; ts = seen / dt
        print(f'step={step} tokens={toks} ce={l1.item():.4f} ppl={math.exp(min(20, l1.item())):.1f} '
              f'bal={bal.item():.3f} mut={getattr(m,"_mut",0):.3f} tok_s={ts:,.0f}', flush=True)
        t0 = time.time(); seen = 0
    if step % a.ckpt_every == 0:
        save(step); vl = valid_loss(); print(f'CKPT step={step} valid_ce={vl:.4f} valid_ppl={math.exp(min(20, vl)):.1f}', flush=True)
    if step % a.hf_every == 0: hf_push(step)
    if toks >= a.target_tokens:
        print(f'TARGET_TOKENS reached {toks}', flush=True); break
save(step); torch.save({'model': m.state_dict(), 'config': m.config()}, Path(a.out) / 'base.pt'); hf_push(step)
print('PEER_BASE_DONE', json.dumps({'valid_ce': valid_loss(), 'steps': step}), flush=True)