"""q256 QAT + cooldown + export, data-parallel over the GPUs torchrun gives it.

Recipe = quantize.py then cooldown.py, as in qat_nb.py, from no_burst.tch:
  QAT      b7 weights (optionally the 3 embedding slots kept at b15), 256 epochs,
           LR 2.5e-6 x cosine 10 -> 1 over the whole run, scales re-derived
           (256 activation steps, averaged across ranks by upstream's code)
  cooldown 16 epochs, LR 1e-7 x cosine 1 -> 1e-8, scales not re-derived
  export   upstream export_weights (bit-exact check on CUDA) + r5 side tensors

Data parallel: each step is 16 micro-batches of 2^16 tokens; rank r takes every
world-th one (upstream loader), gradients are averaged by DDP. Probe-measured:
same per-step losses as one GPU up to float summation order, 3.9x faster on 4.

Resumable: every `every` steps rank 0 saves model + both optimizers + step, and
all ranks re-sync parameters from rank 0 (guards against slow drift). A restart
skips the batches already done -- the loader's article order is fixed (seed 0).
A finished phase leaves its .tch, so a restart never repeats it.
"""
import argparse, contextlib, dataclasses, datetime, json, os, sys, time
from math import cos, pi
import torch
import torch.distributed as dist

ap = argparse.ArgumentParser()
ap.add_argument("--data", required=True)          # training priors dir
ap.add_argument("--eval", required=True)          # 3 MB gate slice dir
ap.add_argument("--init", required=True)          # fp32 fine-tune checkpoint
ap.add_argument("--out", required=True)           # volume dir for ckpts/results
ap.add_argument("--qat-epochs", type=int, default=256)
ap.add_argument("--cd-epochs", type=int, default=16)
ap.add_argument("--init-steps", type=int, default=256)
ap.add_argument("--accum", type=int, default=4)   # per rank; x world = 16
ap.add_argument("--every", type=int, default=100)
ap.add_argument("--embed15", type=int, default=0)
ap.add_argument("--stop-after", type=int, default=0)  # smoke test: simulate preemption
ap.add_argument("--compile", type=int, default=0)     # torch.compile + exact bf16 scale rounding
# continuation of a finished QAT run (warm restart): start model + optimizer state from that run,
# LR multiplier cosine qat_peak -> 1, no scale re-derivation (--init-steps 0), optionally no cooldown
ap.add_argument("--qat-peak", type=float, default=10.0)   # LR multiplier at the start of the QAT cosine
ap.add_argument("--opt-from", default="")                 # load Muon/AdamW state from this ckpt_*.pt
ap.add_argument("--tag", default="qat")                   # QAT phase name -> ckpt_<tag>.pt, <tag>.tch
ap.add_argument("--control-eval", type=int, default=0)    # gate-evaluate the start model before training
# widen extra slots 7 -> 15 levels on a checkpoint trained with only the --embed15 slots at 15: their weight and
# activation scales are halved (the 15-level grid then contains every 7-level grid point) before training continues
ap.add_argument("--keep15-extra", default="")               # comma list of config slot names, e.g. attention_full_rank_quantization
a = ap.parse_args()

rank = int(os.environ["RANK"]); world = int(os.environ["WORLD_SIZE"]); local = int(os.environ["LOCAL_RANK"])
dev = torch.device("cuda", local); torch.cuda.set_device(dev)
dist.init_process_group("nccl", timeout=datetime.timedelta(hours=2))
def log(*x):
    if rank == 0:
        print(f"[{time.strftime('%H:%M:%S')}]", *x, flush=True)

sys.path.insert(0, "/root/s"); os.chdir("/root/s")
# upstream QuantizeFunction rounds scales to bf16 with .to(bfloat16).type_as(fp32), and
# torch.compile fuses that round trip away (upstream's warning in quantize.py; measured:
# 0/111 weight tensors bit-identical, first-step loss +1.66%). Same RNE rounding in
# integer ops cannot be elided: 111/111 bit-identical under compile, and identical in eager.
QP = "/root/s/pysrc/quantization.py"
if not os.path.exists(QP + ".orig"):
    import shutil; shutil.copy(QP, QP + ".orig")
_src = open(QP + ".orig").read()
if a.compile:
    _old = "        scale = scale.to(quant.scale_dtype).type_as(blocks)\n"
    assert _src.count(_old) == 1
    _src = _src.replace(_old, "        assert quant.scale_dtype == torch.bfloat16\n"
                              "        scale = _bf16_round(scale.float()).type_as(blocks)\n")
    _src += '''

def _bf16_round(x):
    """float32 -> bfloat16 -> float32, round-to-nearest-even, in integer ops"""
    b = x.contiguous().view(torch.int32)
    lsb = (b >> 16) & 1
    return (((b + 0x7FFF + lsb) >> 16) << 16).view(torch.float32)
'''
open(QP + f".rank{os.environ.get('LOCAL_RANK', '0')}", "w").write(_src)   # per-rank temp, then atomic swap
os.replace(QP + f".rank{os.environ.get('LOCAL_RANK', '0')}", QP)
import pysrc.model as M
_oi = M.TransformerConfig.__init__
def _init(self, *x, **kw):
    kw.setdefault("query_key_norm_gain", False); _oi(self, *x, **kw)
M.TransformerConfig.__init__ = _init
import pysrc.export_weights as ew
from pysrc.quantization import Quantization
from pysrc.model import Transformer, MatmulQuantization
from pysrc.data import CollatedArticlesDataLoader, Batch, prefetch
from pysrc.optimizer import make_optimizers, AdamWConfig, MuonConfig, OptimizerConfig
from pysrc.train import TrainingConfig, initialize_scales
from pysrc.logger import LoggerConfig
from torch.nn.parallel import DistributedDataParallel as DDP

KEEP15_BASE = {"embedding_quantization", "prior_embedding_quantization", "unembedding_quantization"} \
    if a.embed15 else set()
EXTRA15 = {x for x in a.keep15_extra.split(",") if x}
KEEP15 = KEEP15_BASE | EXTRA15
def to7(o):
    if isinstance(o, Quantization):
        return dataclasses.replace(o, buckets=7) if o.buckets == 15 else o
    if isinstance(o, MatmulQuantization):
        return dataclasses.replace(o, weight=to7(o.weight), activation=to7(o.activation))
    return o
ref = ew.make_reference_config()
ch = {f.name: to7(getattr(ref, f.name)) for f in dataclasses.fields(ref)
      if f.name not in KEEP15 and to7(getattr(ref, f.name)) is not getattr(ref, f.name)}
raw = Transformer(dataclasses.replace(ref, **ch)).to(dev)
WIDENED = {}   # module name -> (old buckets, new buckets) for quantizers that change width vs the checkpoint
if EXTRA15:
    bad = EXTRA15 - {f.name for f in dataclasses.fields(ref)}
    assert not bad, f"unknown slot names {bad}"
    ch_old = {f.name: to7(getattr(ref, f.name)) for f in dataclasses.fields(ref)
              if f.name not in KEEP15_BASE and to7(getattr(ref, f.name)) is not getattr(ref, f.name)}
    old_mods = dict(Transformer(dataclasses.replace(ref, **ch_old)).named_modules())
    for nm, mod in raw.named_modules():
        q = getattr(mod, "quant", None); o = getattr(old_mods.get(nm), "quant", None)
        if isinstance(q, Quantization) and isinstance(o, Quantization) and q.buckets != o.buckets:
            WIDENED[nm] = (o.buckets, q.buckets)
    assert WIDENED and all(v == (7, 15) for v in WIDENED.values()), WIDENED
log(f"world {world}, b7 slots {sorted(ch)}, kept at b15: {sorted(KEEP15)}")
ddp = DDP(raw, device_ids=[local])
fwd = ddp
if a.compile:
    import torch._dynamo
    torch._dynamo.config.cache_size_limit = 64   # varlen batch layouts: allow recompiles, never fall back to eager
    fwd = torch.compile(ddp)
    log("torch.compile on (exact integer bf16 scale rounding patched in)")


def loader(d, epochs, accum):
    return CollatedArticlesDataLoader(
        tokens_filename=f"{d}/ppmd-bytes.uint8", priors_filename=f"{d}/ppmd-probs.float16",
        article_boundaries_filename=f"{d}/article-boundaries.int32", mmap=True, epochs=epochs,
        micro_batch_tokens=2 ** 16, gradient_accumulation_steps=accum, max_article_tokens=2 ** 16,
        max_merge_tokens=None, device=dev, rank=rank, world_size=world)


def opt_cfg(sched, lr):
    return OptimizerConfig(
        learning_rate_multiplier_schedule=sched,
        second_order_config=MuonConfig(base_learning_rate=lr, weight_decay=0.1, momentum=0.95,
                                       adjust_lr_fn="match_rms_adamw"),
        adamw_config=AdamWConfig(base_learning_rate=lr, weight_decay=0.01, betas=(0.95, 0.99)))


def sync_params():
    with torch.no_grad():
        for p in raw.parameters():
            dist.broadcast(p.data, 0)


def widen_scales():
    """widen quantizers 7 -> 15 levels WITHOUT changing what the model computes: set each widened weight's latent
    values to its current 7-level dequantized values (latents drift outside the +-3s clip during 7-level QAT, so any
    re-gridding of raw latents moves outliers: measured +0.0155 / +0.0194 nats for scale x3/7 / x1/2), then halve the
    scale -- the step-s/2, +-3.5s grid contains every old value k*s (|k| <= 3) exactly. The gate control must then
    reproduce the unwidened model's number."""
    if not WIDENED:
        return
    old = Transformer(dataclasses.replace(ref, **ch_old)).to(dev)
    old.load_state_dict(raw.state_dict()); old_mods = dict(old.named_modules())
    mods = dict(raw.named_modules()); n = 0; nw = 0
    with torch.no_grad():
        for nm in WIDENED:
            if nm.endswith(".quantize_weight"):
                parent = nm[: -len(".quantize_weight")]
                w = old_mods[parent].weight
                deq = old_mods[nm](w.detach(), initialize_scales_steps=None, return_ints=False)
                mods[parent].weight.copy_(deq.to(mods[parent].weight.dtype)); nw += w.numel()
            mods[nm].scale.mul_(0.5); n += mods[nm].scale.numel()
    del old
    log(f"widened {len(WIDENED)} quantizers 7->15 levels: latents snapped to their 7-level values ({nw} weights), "
        f"scales x1/2 ({n} values): {sorted(WIDENED)[:4]}...")


def phase(tag, epochs, sched, lr, init_steps, start_from, final, opt_from=""):
    if os.path.exists(final):
        raw.load_state_dict(torch.load(final, map_location=dev)); sync_params()
        log(f"[{tag}] already finished earlier: loaded {final}"); return
    data = loader(a.data, epochs, a.accum)
    opt = make_optimizers(raw, opt_cfg(sched, lr))
    ck = f"{a.out}/ckpt_{tag}.pt"; start = 0
    if os.path.exists(ck):
        st = torch.load(ck, map_location=dev, weights_only=False)
        raw.load_state_dict(st["model"])
        opt.second_order.load_state_dict(st["muon"]); opt.adamw.load_state_dict(st["adamw"])
        start = st["step"]; sync_params()
        log(f"[{tag}] RESUMED at step {start}")
    else:
        if start_from:
            miss, unexp = raw.load_state_dict(torch.load(start_from, map_location=dev), strict=False)
            log(f"[{tag}] start from {start_from}: missing {len(miss)} unexpected {len(unexp)}")
            widen_scales()
        if opt_from:
            st = torch.load(opt_from, map_location=dev, weights_only=False)
            opt.second_order.load_state_dict(st["muon"]); opt.adamw.load_state_dict(st["adamw"])
            log(f"[{tag}] optimizer state from {opt_from} (its step {st['step']})")
        sync_params()
        if init_steps:
            initialize_scales(raw_model=raw, model=ddp, data=data, cfg=TrainingConfig(
                initialize_scales_steps=init_steps, initialize_weight_scales_first=True,
                optimizer_config=opt_cfg(sched, lr), window_size_schedule=None,
                logger_config=LoggerConfig(log_to_stdout=False)))
            log(f"[{tag}] scales re-derived over {init_steps} steps")
    total = data._total_batches
    raw.train()
    if a.compile:
        # interlock: compiled fake-quant must match eager. Measured after q256 QAT: 30 of 5,868,864
        # weights sit exactly on a rounding boundary and flip by one level (division vs reciprocal);
        # that is float noise. Anything bigger than 1e-4 of the weights, or >1 level, stops the run.
        bad = 0; tot_el = 0; worst = 0
        with torch.no_grad():
            for nm, mod in raw.named_modules():
                qm = getattr(mod, "quantize_weight", None); w = getattr(mod, "weight", None)
                if qm is None or w is None:
                    continue
                e = qm(w.detach(), initialize_scales_steps=None, return_ints=True)
                c = torch.compile(lambda t, qm=qm: qm(t, initialize_scales_steps=None, return_ints=True))(w.detach())
                d = (e - c).abs(); bad += int((d > 0).sum()); tot_el += e.numel(); worst = max(worst, int(d.max()))
        if bad > 1e-4 * tot_el or worst > 1:
            log(f"[{tag}] INTERLOCK FAIL: {bad} of {tot_el} quantized weights differ, max {worst} levels")
            dist.destroy_process_group(); sys.exit(42)
        log(f"[{tag}] interlock ok: {bad} of {tot_el} quantized weights differ from eager (max {worst} level)")

    def batches():
        done = 0
        for ep in range(data.epochs):
            per = data._batches_per_epoch[ep]
            for bi, groups in enumerate(data._iter_batch_article_groups(data._permutations[ep]), 1):
                done += 1
                if done <= start:
                    continue
                yield Batch(micro_batches=[data._make_micro_batch(g) for i, g in enumerate(groups)
                                           if i % world == rank],
                            epoch=ep + bi / per, fraction_done=done / total, is_last=done == total)

    step = start; t0 = time.time(); run_loss = None
    for batch in prefetch(batches()):
        opt.update_learning_rate_schedule(fraction_done=batch.fraction_done)
        tot = torch.zeros((), device=dev); n = len(batch.micro_batches)
        for i, mb in enumerate(batch.micro_batches):
            ctx = ddp.no_sync() if i < n - 1 else contextlib.nullcontext()
            with ctx:
                loss = fwd(input_tokens=mb.input_tokens, priors=mb.priors, loss_mask=mb.loss_mask,
                           output_tokens=mb.output_tokens,
                           cumulative_sequence_lengths=mb.cumulative_sequence_lengths,
                           override_window_size=None)
                (loss / a.accum).backward()
            tot += loss.detach()
        opt.step(); opt.zero_grad()
        step += 1
        if step % 20 == 0 or step == total:
            dist.all_reduce(tot)
            l = float(tot) / (n * world)
            if not (l == l):
                raise SystemExit(f"[{tag}] NaN loss at step {step}")
            run_loss = l if run_loss is None else 0.9 * run_loss + 0.1 * l
            r = (time.time() - t0) / (step - start)
            log(f"[{tag}] step {step}/{total}  {r:.2f}s/it  loss {l:.1f} (avg {run_loss:.1f})  "
                f"eta {(total - step) * r / 3600:.2f}h")
        if step % a.every == 0 and step < total:
            sync_params()
            if rank == 0:
                torch.save({"model": raw.state_dict(), "muon": opt.second_order.state_dict(),
                            "adamw": opt.adamw.state_dict(), "step": step}, ck + ".tmp")
                os.replace(ck + ".tmp", ck)
            dist.barrier()
        if a.stop_after and step - start >= a.stop_after:
            log(f"[{tag}] --stop-after: exiting at step {step} to simulate a preemption")
            dist.destroy_process_group(); sys.exit(3)
    assert step == total, (step, total)
    sync_params()
    if rank == 0:
        torch.save(raw.state_dict(), final)
    dist.barrier()
    log(f"[{tag}] done, saved {final}")


def qsched(f):
    t = (cos(f * pi) + 1) / 2
    return t * a.qat_peak + (1 - t)
def csched(f):
    t = (cos(f * pi) + 1) / 2
    return t + (1 - t) * 1e-8

def gate_eval():
    """mean -log p(next byte) over the 3 MB gate slice, the number all our gates quote"""
    import numpy as np
    from pysrc.save_outputs import save_transformer_outputs
    raw.eval()
    idx = np.fromfile(f"{a.eval}/ppmd-bytes.uint8", dtype=np.uint8).astype(np.int64); n = len(idx)
    ev = CollatedArticlesDataLoader(
        tokens_filename=f"{a.eval}/ppmd-bytes.uint8", priors_filename=f"{a.eval}/ppmd-probs.float16",
        article_boundaries_filename=f"{a.eval}/article-boundaries.int32", mmap=True, epochs=1,
        micro_batch_tokens=2 ** 16, gradient_accumulation_steps=1, max_article_tokens=2 ** 16,
        max_merge_tokens=None, device=dev, rank=0, world_size=1)
    save_transformer_outputs(model=raw, data=ev, output_filename="/tmp/final.f16")
    arr = np.memmap("/tmp/final.f16", dtype=np.float16, mode="r", shape=(n, 205))
    tot = 0.0; cnt = 0
    for s0 in range(0, n - 1, 65536):
        e = min(s0 + 65536, n - 1)
        blk = np.asarray(arr[s0:e]).astype(np.float64); sm = blk.sum(axis=1)
        q = blk[np.arange(e - s0), idx[s0 + 1:e + 1]] / np.maximum(sm, 1e-300)
        tot += float(-np.log(np.maximum(q, 1e-38)).sum()); cnt += e - s0
    raw.train()
    return tot / cnt


control = None
if a.control_eval:
    if rank == 0:
        raw.load_state_dict(torch.load(a.init, map_location=dev), strict=False)
        widen_scales()
        control = gate_eval()
        log(f"CONTROL gate-slice loss of the start model {a.init}: {control:.6f}")
    dist.barrier()
phase(a.tag, a.qat_epochs, qsched, 2.5e-6, a.init_steps, a.init, f"{a.out}/{a.tag}.tch", a.opt_from)
if a.cd_epochs:
    phase("cd", a.cd_epochs, csched, 1e-7, 0, None, f"{a.out}/qat_cd.tch")

if rank == 0 and not os.path.exists(f"{a.out}/result.json"):
    import numpy as np
    slice_loss = gate_eval()
    log(f"FINAL gate-slice loss {slice_loss:.6f}  (q256 no_burst e15: 1.265659 before cooldown, 1.231430 after)")

    # export: upstream exporter, with the clamp taken from each weight's own quantizer
    # (upstream hard-codes the 15-level +-7), then r5 side tensors (RNE to 5 mantissa bits)
    from pysrc.weights_compress import write_tensor_file_v2, _entries_from_reference
    sd = raw.state_dict(); mods = dict(raw.named_modules())
    qmax = {sd[k].data_ptr(): mods[k[:-len(".weight")]].quantize_weight.quant.qmax
            for k in sd if k.endswith(".weight") and k[:-len(".weight")] + ".quantize_weight.scale" in sd}
    def qrows(weight, scale_fp32):
        qm = qmax[weight.data_ptr()]
        scale = scale_fp32.to(torch.bfloat16).to(torch.float32)
        q = (weight.float() / scale.unsqueeze(-1)).round().clamp(-qm, qm)
        return q.to(torch.int8), scale
    ew.quantize_weight_rows = qrows
    torch.save(raw.state_dict(), "/tmp/final.tch")
    ew.export_weights(raw, "/tmp/w.bin")
    entries = _entries_from_reference("/tmp/w.bin")
    drop = 18; half = np.uint32(1 << (drop - 1)); out = []
    for name, dt, arr_ in entries:
        if name.endswith(ew._UNQUANTIZED_FP32_SUFFIXES):
            b = np.ascontiguousarray(arr_, dtype=np.float32).view(np.uint32)
            lsb = (b >> np.uint32(drop)) & np.uint32(1)
            arr_ = (((b + half - np.uint32(1) + lsb) >> np.uint32(drop)) << np.uint32(drop)).view(np.float32).reshape(arr_.shape)
        out.append((name, dt, arr_))
    write_tensor_file_v2(f"{a.out}/r5.tfwc2", out)
    res = {"slice_loss": slice_loss, "control_slice_loss": control, "qat_peak": a.qat_peak,
           "keep15_extra": sorted(EXTRA15), "widened_quantizers": len(WIDENED),
           "init": a.init, "embed15": a.embed15, "qat_epochs": a.qat_epochs,
           "cd_epochs": a.cd_epochs, "world": world, "r5_bytes": os.path.getsize(f"{a.out}/r5.tfwc2")}
    json.dump(res, open(f"{a.out}/result.json", "w"), indent=1)
    log(json.dumps(res))
dist.barrier()
dist.destroy_process_group()