# tern_qat.py - Experiment step 1: warm-start ternary QAT of the 97M controller.
# 1) load base.pt  2) measure valid CE (fp baseline)  3) ternarize -> measure raw
# quant CE  4) short QAT to recover  5) report CE + per-token byte budget.
import argparse, json, math, os, time, sys
from pathlib import Path
import numpy as np, torch, torch.nn.functional as F
sys.path.insert(0, '/root/dna')
from model_dna import DnaChat, counts
from model_dna_tern import ternarize_, controller_byte_budget

ap = argparse.ArgumentParser()
ap.add_argument('--base', default='/root/dna/ckpt/base.pt')
ap.add_argument('--data-dir', default='/root/dna/data')
ap.add_argument('--out', default='/root/dna/ckpt-tern')
ap.add_argument('--batch', type=int, default=16)
ap.add_argument('--seq', type=int, default=512)
ap.add_argument('--steps', type=int, default=1500)
ap.add_argument('--lr', type=float, default=1e-4)
ap.add_argument('--warmup', type=int, default=100)
ap.add_argument('--valid-iters', type=int, default=20)
ap.add_argument('--smoke', action='store_true')
a = ap.parse_args()
dev = 'cuda' if torch.cuda.is_available() else 'cpu'
Path(a.out).mkdir(parents=True, exist_ok=True); torch.manual_seed(0)

shards = sorted(Path(a.data_dir).glob('shard_*.u16'))
_cur = {'i': -1, 'arr': None}
def batch():
    i = np.random.randint(0, len(shards))
    if i != _cur['i']: _cur['arr'] = np.memmap(shards[i], np.uint16, 'r'); _cur['i'] = i
    arr = _cur['arr']; n = len(arr) - a.seq - 1
    ix = np.random.randint(0, n, size=a.batch)
    return torch.from_numpy(np.stack([np.asarray(arr[j:j+a.seq+1], np.int64) for j in ix])).to(dev)

@torch.no_grad()
def valid_ce(m, iters):
    m.eval(); tot = 0.0
    for _ in range(iters):
        ids = batch()
        with torch.autocast(dev, dtype=torch.bfloat16, enabled=(dev=='cuda')):
            feat, _, _ = m.features(ids)
            ce = F.cross_entropy(F.linear(feat[:, :-1].reshape(-1, m.d), m.embed.weight),
                                 ids[:, 1:].reshape(-1))
        tot += ce.item()
    m.train(); return tot / iters

ck = torch.load(a.base, map_location=dev, weights_only=False)
m = DnaChat(**ck['config']).to(dev)
m.load_state_dict({k.replace('_orig_mod.', ''): v for k, v in ck['model'].items()}, strict=True)
ctrl, _ = counts(m)
if a.smoke:
    a.steps = 3; a.valid_iters = 2

ce_fp = valid_ce(m, a.valid_iters)
n_swapped = ternarize_(m)
bud = controller_byte_budget(m)
ce_q0 = valid_ce(m, a.valid_iters)
print('TERN_INIT', json.dumps({
    'ctrl_params': ctrl, 'ternarized_linears': n_swapped,
    'valid_ce_fp16': round(ce_fp, 4), 'valid_ce_ternary_raw': round(ce_q0, 4),
    'bytes_per_token_bf16': round(bud['bytes_per_token_bf16']/1e6, 1),
    'bytes_per_token_tern': round(bud['bytes_per_token_tern']/1e6, 1),
    'tern_params': bud['tern_params'], 'fp_linear_params': bud['fp_linear_params'],
}), flush=True)

opt = torch.optim.AdamW(m.parameters(), lr=a.lr, betas=(.9, .95), weight_decay=0.0)
def lrf(s): return s/a.warmup if s < a.warmup else max(0.1, 0.5*(1+math.cos(math.pi*(s-a.warmup)/max(1, a.steps-a.warmup))))
t0 = time.time()
for step in range(1, a.steps + 1):
    for g in opt.param_groups: g['lr'] = a.lr * lrf(step)
    ids = batch(); opt.zero_grad(set_to_none=True)
    with torch.autocast(dev, dtype=torch.bfloat16, enabled=(dev=='cuda')):
        feat, bal, _ = m.features(ids)
        ce = F.cross_entropy(F.linear(feat[:, :-1].reshape(-1, m.d), m.embed.weight), ids[:, 1:].reshape(-1))
        loss = ce + 1e-4 * bal
    loss.backward(); torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0); opt.step()
    if step % 100 == 0 or step == a.steps:
        print(f'qat step={step}/{a.steps} ce={ce.item():.4f} ppl={math.exp(min(20,ce.item())):.1f} '
              f'tok_s={step*a.batch*a.seq/(time.time()-t0):,.0f}', flush=True)

ce_qat = valid_ce(m, a.valid_iters)
torch.save({'model': m.state_dict(), 'config': m.config(), 'ternary': True}, Path(a.out)/'tern.pt')
print('TERN_DONE', json.dumps({
    'valid_ce_fp16': round(ce_fp, 4), 'valid_ce_ternary_raw': round(ce_q0, 4),
    'valid_ce_ternary_qat': round(ce_qat, 4),
    'ppl_fp16': round(math.exp(min(20, ce_fp)), 1),
    'ppl_ternary_qat': round(math.exp(min(20, ce_qat)), 1),
    'bytes_per_token_bf16': round(bud['bytes_per_token_bf16']/1e6, 1),
    'bytes_per_token_tern': round(bud['bytes_per_token_tern']/1e6, 1),
}), flush=True)