"""Low-bit weight storage for the char Transformer: matrices -> b-bit (4 or 8) symmetric integers with one fp16
scale per group of G consecutive input weights (per output row); vectors -> fp16. De-quantised to fp32 at load."""
import io, json
import numpy as np, torch
from bhc.xlit_model import XlitTransformer


def _q(v, bits, G):
    r, c = v.shape; pad = (-c) % G; vp = np.pad(v, ((0, 0), (0, pad))).reshape(r, -1, G)
    L = 2 ** (bits - 1) - 1; base = np.abs(vp).max(2, keepdims=True) / L; base[base == 0] = 1
    best_err = None
    for al in ((1.0,) if bits >= 8 else (1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65)):   # MSE-optimal clipping per group
        s_ = (base * al).astype(np.float16).astype(np.float32); q_ = np.clip(np.round(vp / s_), -L, L)
        e = ((q_ * s_ - vp) ** 2).sum(2, keepdims=True)
        if best_err is None: best_err, sc, q = e, s_, q_
        else: m = e < best_err; best_err = np.where(m, e, best_err); sc = np.where(m, s_, sc); q = np.where(m, q_, q)
    q = q.astype(np.int8)
    return q.reshape(r, -1)[:, :c], sc[..., 0].astype(np.float16)


def export_q(model, path, bits=4, G=32, emb_bits=8):
    """bits for Linear weights; emb_bits for the (small, outlier-prone) embedding/position tables."""
    sd = model.state_dict(); meta = {"cfg": model.cfg, "bits": bits, "G": G, "tensors": []}; buf = io.BytesIO()
    for k, v in sd.items():
        v = v.float().cpu().numpy()
        if v.ndim == 2:
            b_ = emb_bits if ("emb" in k or "pos" in k) else bits
            q, sc = _q(v, b_, G); flat = q.reshape(-1)
            if b_ == 4:
                u = (flat + 8).astype(np.uint8); u = np.r_[u, np.zeros(len(u) % 2, np.uint8)]; data = (u[0::2] << 4 | u[1::2]).astype(np.uint8).tobytes()
            else: data = flat.astype(np.int8).tobytes()
            meta["tensors"].append([k, list(v.shape), "q%d" % b_, buf.tell(), len(data)]); buf.write(data); buf.write(sc.tobytes())
        else:
            meta["tensors"].append([k, list(v.shape), "f16", buf.tell(), 0]); buf.write(v.astype(np.float16).tobytes())
    hdr = json.dumps(meta, separators=(",", ":")).encode()
    with open(path, "wb") as f: f.write(len(hdr).to_bytes(4, "little")); f.write(hdr); f.write(buf.getvalue())


def load_q(path):
    b = open(path, "rb").read(); n = int.from_bytes(b[:4], "little"); meta = json.loads(b[4:4 + n]); data = b[4 + n:]
    bits, G = meta["bits"], meta["G"]; m = XlitTransformer(**meta["cfg"]); sd = {}
    for k, shape, kind, off, ln in meta["tensors"]:
        cnt = int(np.prod(shape))
        if kind.startswith("q"):
            bits = int(kind[1:]) if len(kind) > 1 else meta["bits"]
            raw = np.frombuffer(data, np.uint8 if bits == 4 else np.int8, ln, off)
            if bits == 4: q = (np.stack([raw >> 4, raw & 15], 1).reshape(-1)[:cnt].astype(np.int16) - 8)
            else: q = raw.astype(np.int16)
            r, c = shape; ng = (c + G - 1) // G
            sc = np.frombuffer(data, np.float16, r * ng, off + ln).astype(np.float32).reshape(r, ng)
            w = q.reshape(r, c).astype(np.float32) * np.repeat(sc, G, 1)[:, :c]
            sd[k] = torch.from_numpy(w)
        else:
            sd[k] = torch.from_numpy(np.frombuffer(data, np.float16, cnt, off).astype(np.float32).reshape(shape).copy())
    m.load_state_dict(sd); m.eval(); return m


def load_any(path):
    if path.endswith(".int8"):
        from bhc.xlit_model import load_int8; return load_int8(path)
    if path.endswith(".pt"):
        ck = torch.load(path, weights_only=False, map_location="cpu"); m = XlitTransformer(**ck["cfg"]); m.load_state_dict(ck["model"]); m.eval(); return m
    return load_q(path)