#!/usr/bin/env python3
"""Resident-QUANTIZED any-k base-3 (ubit) — the cheat code generalized past k5.

Cody, 2026-09-17: "remember the k5 cheat code. we don't need to unpack." grouped_place's --out
path DEQUANTED every placed linear to a dense bf16 checkpoint (54 GB, and the 248320-row head's
unpack OOMs). That threw the packing away. Here the placed PackedKTritLinear buffers (packed
bytes + group scales) are saved AS-IS and reloaded resident, decoding per-forward one linear at a
time — never a global unpack, ~10 GB on disk and in VRAM. Mirrors k5_resident.load_k5_resident.
"""
import os, json, torch, torch.nn as nn
import rotation as ROT
from base3_pack import PackedKTritLinear
from ternary_gptq import _RotatedInput


def save_packed(model, out_dir, rot_map, grp_map, log=print):
    """Save the PLACED model without dequant: each PackedKTritLinear's bytes+scale + a manifest of
    (k,g,rot,in,out,bias) per linear, plus the dense small tensors (embed/norms/conv/biases)."""
    from safetensors.torch import save_file
    os.makedirs(out_dir, exist_ok=True)
    packed_stems = set()
    tensors, manifest = {}, {}
    wrapped = {id(m.lin) for _, m in model.named_modules() if isinstance(m, _RotatedInput)}
    for name, mod in model.named_modules():
        if isinstance(mod, _RotatedInput):
            inner, blk = mod.lin, int(getattr(mod, "blk", 0))
        elif isinstance(mod, PackedKTritLinear) and id(mod) not in wrapped:
            inner, blk = mod, 0                       # bare (unrotated) linear
        else:
            continue                                  # skip the .lin child of a wrapper (dup)
        if isinstance(inner, PackedKTritLinear):
            tensors[f"{name}.bytes"] = inner.bytes.cpu().contiguous()
            tensors[f"{name}.scale"] = inner.scale.cpu().contiguous()
            if inner.bias is not None:
                tensors[f"{name}.bias"] = inner.bias.detach().to(torch.bfloat16).cpu().contiguous()
            manifest[name] = {"k": int(inner.k), "g": int(inner.group_size), "blk": blk,
                              "in": int(inner.in_features), "out": int(inner.out_features),
                              "bias": inner.bias is not None}
            packed_stems.add(name)
    # dense remainder: every state_dict tensor not owned by a packed linear
    for k, v in model.state_dict().items():
        if any(k == f"{s}.bytes" or k == f"{s}.scale" or k == f"{s}.bias" or
               k.startswith(f"{s}.") for s in packed_stems):
            continue
        tensors[f"__dense__{k}"] = (v.to(torch.bfloat16) if v.is_floating_point() else v).cpu().contiguous()
    save_file(tensors, os.path.join(out_dir, "packed.safetensors"))
    json.dump(manifest, open(os.path.join(out_dir, "ubit_manifest.json"), "w"))
    log(f"[save-packed] {len(manifest)} packed linears + {sum(1 for k in tensors if k.startswith('__dense__'))} "
        f"dense tensors -> {out_dir}/packed.safetensors (NO dequant)")
    return len(manifest)


def load_ktrit_resident(ck, device, dtype=torch.bfloat16, log=print):
    """Rebuild the placed model resident: each manifest linear as a PackedKTritLinear (rotation-
    wrapped if blk>1), decode-per-forward. No global unpack."""
    from transformers import AutoConfig, AutoModelForCausalLM
    import transformers
    from safetensors.torch import load_file
    manifest = json.load(open(os.path.join(ck, "ubit_manifest.json")))
    blob = load_file(os.path.join(ck, "packed.safetensors"))
    config = AutoConfig.from_pretrained(ck)
    dk = "dtype" if int(transformers.__version__.split(".")[0]) >= 5 else "torch_dtype"
    is_vl = getattr(config, "vision_config", None) is not None      # retrofitted vision+MTP -> VL class
    with torch.device("meta"):
        if is_vl:
            from transformers import AutoModelForImageTextToText
            model = AutoModelForImageTextToText.from_config(config, **{dk: dtype})
        else:
            model = AutoModelForCausalLM.from_config(config, **{dk: dtype})
    model.eval()
    # 1) swap in the packed linears
    def _set(root, dotted, module):
        *ps, leaf = dotted.split("."); p = root
        for x in ps: p = getattr(p, x)
        setattr(p, leaf, module)
    skipped = []
    for name, m in manifest.items():
        bias = blob.get(f"{name}.bias")
        pk = PackedKTritLinear(blob[f"{name}.bytes"], blob[f"{name}.scale"], m["in"], m["k"],
                               bias.float() if bias is not None else None, m["out"], m["g"]).to(device)
        try:
            _set(model, name, _RotatedInput(pk, m["blk"]) if m["blk"] and m["blk"] > 1 else pk)
        except AttributeError:
            skipped.append(name)                 # no module in this arch (e.g. MTP head absent in this transformers)
    if skipped:
        log(f"[load-packed] {len(skipped)} packed keys have no module in this arch (kept in file, not attached): e.g. {skipped[:3]}")
    # 2) load the dense remainder (embed/norms/conv/dt_bias/…) onto meta
    dense = {k[len("__dense__"):]: v for k, v in blob.items() if k.startswith("__dense__")}
    model.load_state_dict(dense, strict=False, assign=True)
    # 3) materialize meta rotary/pos buffers (recomputed from config, never in the checkpoint) by
    # re-instantiating their owning module -- LM via text_config, vision via its dim. Any leftover
    # meta tensor makes model.to(device) throw "Cannot copy out of meta tensor".
    rope = getattr(config, "text_config", None) or config
    vcfg = getattr(config, "vision_config", None)
    def _set(root, dotted, module):
        *ps, leaf = dotted.split("."); p = root
        for x in ps: p = getattr(p, x)
        setattr(p, leaf, module)
    for mn, mod in list(model.named_modules()):
        metab = [bn for bn, b in mod._buffers.items() if b is not None and b.is_meta]
        if not metab:
            continue
        cfg_for = vcfg if "visual" in mn else rope
        ib = mod._buffers.get("inv_freq")
        trials = [lambda: type(mod)(config=cfg_for, device=device), lambda: type(mod)(config=cfg_for), lambda: type(mod)(cfg_for)]
        if ib is not None:
            dim = 2 * ib.shape[0]
            trials += [lambda: type(mod)(dim), lambda: type(mod)(dim=dim), lambda: type(mod)(dim, device=device)]
        made = None
        for t in trials:
            try: made = t(); break
            except Exception: continue
        if made is not None:
            _set(model, mn, made.to(device))
        else:
            for bn in metab:
                mod._buffers[bn] = torch.zeros(mod._buffers[bn].shape, dtype=torch.float32, device=device)
            log(f"[load-packed] WARN zeroed meta {metab} in {mn} (rope may be off)")
    model = model.to(device)
    log(f"[load-packed] {len(manifest)} resident PackedKTritLinear + {len(dense)} dense tensors")
    return model


if __name__ == "__main__":   # toy round-trip smoke: packed save/load must be bit-exact
    import tempfile
    from base3_pack import quantize_ktrit, pack_ktrit
    torch.manual_seed(0)
    for k, g, blk in [(1, 64, 0), (2, 64, 64), (3, 128, 128)]:
        lin = nn.Linear(256, 512)
        codes, scale = quantize_ktrit(lin.weight.data, k, g)
        b, _, _ = pack_ktrit(codes, k)
        pk = PackedKTritLinear(b, scale, 256, k, lin.bias.data.float(), 512, g)
        wrapped = _RotatedInput(pk, blk) if blk > 1 else pk
        x = torch.randn(4, 256)
        y0 = wrapped(x)
        # emulate save_packed's per-module capture + load reconstruction
        st = {"bytes": pk.bytes.clone(), "scale": pk.scale.clone(), "bias": pk.bias.clone()}
        pk2 = PackedKTritLinear(st["bytes"], st["scale"], 256, k, st["bias"], 512, g)
        w2 = _RotatedInput(pk2, blk) if blk > 1 else pk2
        y1 = w2(x)
        assert torch.equal(y0, y1), f"packed round-trip MISMATCH k{k} blk{blk}"
        print(f"k{k} g{g} blk{blk}: packed save/load forward EXACT")
    print("========== ktrit_resident round-trip OK ==========")