"""Pack / unpack native 1.58-bit ternary weights (5 trits per byte)."""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch

from model import Boopit, BoopitConfig, ternary_and_scale


MAX_PACKED_BYTES = 7 * 1024 * 1024


def pack_trits(t: np.ndarray) -> tuple[bytes, int]:
    """Map {-1,0,1} -> {0,1,2} and pack 5 trits into each byte (3^5=243)."""
    t = t.astype(np.int8).ravel()
    coded = (t + 1).astype(np.uint8)
    pad = (5 - (coded.size % 5)) % 5
    if pad:
        coded = np.concatenate([coded, np.full(pad, 1, dtype=np.uint8)])  # 0-trit
    packed = (
        coded[0::5].astype(np.uint16)
        + coded[1::5].astype(np.uint16) * 3
        + coded[2::5].astype(np.uint16) * 9
        + coded[3::5].astype(np.uint16) * 27
        + coded[4::5].astype(np.uint16) * 81
    ).astype(np.uint8)
    return packed.tobytes(), pad


def unpack_trits(data: bytes, n: int, pad: int) -> np.ndarray:
    packed = np.frombuffer(data, dtype=np.uint8).astype(np.uint16)
    out = np.empty(packed.size * 5, dtype=np.int8)
    v = packed
    for i in range(5):
        out[i::5] = (v % 3).astype(np.int8)
        v //= 3
    if pad:
        out = out[: out.size - pad]
    return (out[:n] - 1).astype(np.int8)


def extract_packed(model: Boopit) -> dict:
    tensors = {}
    seen: set[int] = set()
    for name, param in model.named_parameters():
        if id(param) in seen:
            continue
        seen.add(id(param))
        w = param.detach()
        if w.ndim == 2:
            t, scale = ternary_and_scale(w)
            packed, pad = pack_trits(t.cpu().numpy())
            tensors[name] = {
                "kind": "ternary",
                "shape": list(w.shape),
                "scale": float(scale.cpu()),
                "pad": pad,
                "data": packed,
            }
        else:
            tensors[name] = {
                "kind": "fp16",
                "shape": list(w.shape),
                "data": w.detach().to(torch.float16).cpu().numpy().tobytes(),
            }
    return tensors


def packed_nbytes(tensors: dict) -> int:
    return sum(len(v["data"]) for v in tensors.values())


def save_packed(model: Boopit, path: Path, extra: dict | None = None) -> int:
    path = Path(path)
    tensors = extract_packed(model)
    nbytes = packed_nbytes(tensors)
    payload = {
        "format": "boopit-1.58",
        "config": model.config.to_dict(),
        "tied": ["tok_emb.weight", "lm_head.weight"],
        "extra": extra or {},
        "tensors": {},
    }
    blob_parts = []
    offset = 0
    for name, spec in tensors.items():
        data = spec["data"]
        entry = {k: v for k, v in spec.items() if k != "data"}
        entry["offset"] = offset
        entry["nbytes"] = len(data)
        payload["tensors"][name] = entry
        blob_parts.append(data)
        offset += len(data)
    header = json.dumps(payload, separators=(",", ":")).encode("utf-8")
    header_len = len(header).to_bytes(8, "little")
    blob = b"".join(blob_parts)
    raw = header_len + header + blob
    if len(raw) >= MAX_PACKED_BYTES:
        raise RuntimeError(f"packed model is {len(raw)} bytes, limit is {MAX_PACKED_BYTES}")
    path.write_bytes(raw)
    return len(raw)


def load_packed(path: Path, device: torch.device | str = "cpu") -> Boopit:
    raw = Path(path).read_bytes()
    header_len = int.from_bytes(raw[:8], "little")
    payload = json.loads(raw[8 : 8 + header_len].decode("utf-8"))
    blob = raw[8 + header_len :]
    cfg = BoopitConfig.from_dict(payload["config"])
    model = Boopit(cfg)
    named = dict(model.named_parameters())
    loaded: set[int] = set()
    with torch.no_grad():
        for name, spec in payload["tensors"].items():
            param = named[name]
            chunk = blob[spec["offset"] : spec["offset"] + spec["nbytes"]]
            if spec["kind"] == "ternary":
                t = unpack_trits(chunk, int(np.prod(spec["shape"])), spec["pad"])
                w = torch.from_numpy(t.astype(np.float32).reshape(spec["shape"])) * float(spec["scale"])
                param.copy_(w.to(dtype=param.dtype))
            else:
                arr = np.frombuffer(chunk, dtype=np.float16).reshape(spec["shape"])
                param.copy_(torch.from_numpy(arr.copy()).to(dtype=param.dtype))
            loaded.add(id(param))
        # tied lm_head
        if id(model.lm_head.weight) not in loaded:
            model.lm_head.weight = model.tok_emb.weight
    return model.to(device)