"""T2.7 · Export packed ternary students to llama.cpp GGUF.

Ternary weights go out as mainline llama.cpp's Q2_0: blocks of 64 values, an fp16 scale and 2-bit codes, where
the stored value q+1 means q·d (00 = −1, 01 = 0, 10 = +1). Our packing (`quant.ternary.pack_2bit`) uses the same
bit layout, and our groups are 128 wide, so each block takes 16 of our code bytes unchanged and its group's scale,
written twice per group. The export is lossless as long as every scale is an fp16 normal (≥6.1e-5; checked).

D-51's N-bit experts go out as the GGUF type whose grid they were trained on (`quant.ternary.GRIDS`): int2 as
Q2_0 (the same 2-bit packing, groups of 64, so one block per group with its signed scale), int4 as Q4_0, int3 as
Q3_K with its down projection as Q4_0 (rows of 384 cannot be Q3_K). Their packers are below; gguf-py has
Q4_0's quantiser and both types' dequantisers, and `check_round_trip` dequantises with gguf-py.

Everything else is written as the student holds it: bf16 matrices as BF16 (their bytes copied), and 1-D norms
and the router as F32, as llama.cpp's converter writes them. Metadata (tokenizer, hyperparameters, chat template)
comes from a template GGUF of the same architecture made by llama.cpp's `convert_hf_to_gguf.py`, whose tensor list
and order the export follows.
"""
from __future__ import annotations

import re
import sys
from pathlib import Path

import numpy as np
import torch

Q2_0_BLOCK, Q2_0_BYTES = 64, 18
Q4_0_BLOCK, Q4_0_BYTES = 32, 18
Q3_K_BLOCK, Q3_K_BYTES = 256, 110
IQ4_XS_BLOCK, IQ4_XS_BYTES = 256, 136
FP16_MIN_NORMAL = 6.103515625e-05


def gguf_module(llama_cpp: str | Path):
    sys.path.insert(0, str(Path(llama_cpp) / "gguf-py"))
    import gguf
    return gguf


def q2_0_blocks(codes: torch.Tensor, scales: torch.Tensor, group: int) -> np.ndarray:
    """Packed codes [..., C/4] (uint8) and scales [..., C/group] (fp16) → Q2_0 bytes [..., C/64 · 18]."""
    if group % Q2_0_BLOCK:
        raise ValueError(f"group size {group} is not a multiple of Q2_0's {Q2_0_BLOCK}")
    c = codes.contiguous().numpy()
    s = scales.to(torch.float16).contiguous().numpy()
    lead, n_bytes = c.shape[:-1], c.shape[-1]
    nb = n_bytes * 4 // Q2_0_BLOCK
    d = np.repeat(s.reshape(-1, s.shape[-1]), group // Q2_0_BLOCK, axis=1)  # [rows, nb] fp16
    out = np.empty((d.shape[0], nb, Q2_0_BYTES), dtype=np.uint8)
    out[:, :, :2] = d.view(np.uint8).reshape(d.shape[0], nb, 2)
    out[:, :, 2:] = c.reshape(-1, nb, Q2_0_BLOCK // 4)
    return out.reshape(*lead, nb * Q2_0_BYTES)


def dequant_q2_0(raw: np.ndarray) -> np.ndarray:
    """Q2_0 bytes [..., nb · 18] → float32 [..., nb · 64], as ggml's dequantize_row_q2_0."""
    lead = raw.shape[:-1]
    b = raw.reshape(-1, Q2_0_BYTES)
    d = b[:, :2].copy().view(np.float16).astype(np.float32)
    q = b[:, 2:]
    vals = np.stack([(q >> s) & 3 for s in (0, 2, 4, 6)], axis=-1).reshape(-1, Q2_0_BLOCK).astype(np.float32) - 1.0
    return (vals * d).reshape(*lead, -1)


def q4_0_blocks(codes: torch.Tensor, scales: torch.Tensor) -> np.ndarray:
    """int4 packed codes [..., C/2] (`pack_4bit` of q+8) and scales [..., C/32] (fp16) → Q4_0 bytes [..., C/32 · 18]:
    per block the fp16 d, then byte j holds element j (low nibble) and element j+16 (high), stored as q+8."""
    from lowbit.quant.ternary import unpack_4bit
    u = unpack_4bit(codes).contiguous().numpy()
    d = scales.to(torch.float16).contiguous().numpy()
    lead, nb = u.shape[:-1], u.shape[-1] // Q4_0_BLOCK
    u = u.reshape(-1, nb, 2, Q4_0_BLOCK // 2)
    out = np.empty((u.shape[0], nb, Q4_0_BYTES), dtype=np.uint8)
    out[:, :, :2] = d.reshape(-1, nb, 1).view(np.uint8)
    out[:, :, 2:] = u[:, :, 0] | (u[:, :, 1] << 4)
    return out.reshape(*lead, nb * Q4_0_BYTES)


def q3_k_blocks(codes: torch.Tensor, d: torch.Tensor, m: torch.Tensor) -> np.ndarray:
    """int3 packed codes [..., C/2] (`pack_4bit` of L = q+4), super-block scales d [..., C/256] (fp16) and sub-block
    scales m [..., C/16] (int8, -32..31) → Q3_K bytes [..., C/256 · 110], the layout of ggml's
    quantize_row_q3_K_ref: hmask[32] (bit b of byte i: L[32b+i] > 3), qs[64] (low two bits of L; per 128-half,
    byte l holds elements l, l+32, l+64, l+96 at bits 0, 2, 4, 6), scales[12] (m+32 in 6 bits: low nibbles of
    sub-blocks 0-7 then 8-15 in bytes 0-7, the high two bits of sub-block j in byte 8 + j%4 at bit 2·(j//4)), d."""
    from lowbit.quant.ternary import unpack_4bit
    L = unpack_4bit(codes).contiguous().numpy()
    lead, nb = L.shape[:-1], L.shape[-1] // Q3_K_BLOCK
    L = L.reshape(-1, nb, Q3_K_BLOCK)
    rows = L.shape[0]
    hb = ((L >> 2) & 1).reshape(rows, nb, 8, 32)
    hmask = np.zeros((rows, nb, 32), dtype=np.uint8)
    for b in range(8):
        hmask |= hb[:, :, b] << b
    low = (L & 3).reshape(rows, nb, 2, 4, 32)
    qs = np.zeros((rows, nb, 2, 32), dtype=np.uint8)
    for k in range(4):
        qs |= low[:, :, :, k] << (2 * k)
    sc = (m.to(torch.int16) + 32).to(torch.uint8).contiguous().numpy().reshape(rows, nb, 16)
    lo4 = sc & 15
    hi2 = ((sc >> 4) & 3).reshape(rows, nb, 4, 4)  # [.., j // 4, j % 4]
    scales = np.zeros((rows, nb, 12), dtype=np.uint8)
    scales[:, :, :8] = lo4[:, :, :8] | (lo4[:, :, 8:] << 4)
    for r in range(4):
        scales[:, :, 8:] |= hi2[:, :, r] << (2 * r)
    dd = d.to(torch.float16).contiguous().numpy().reshape(rows, nb, 1).view(np.uint8)
    out = np.concatenate([hmask, qs.reshape(rows, nb, 64), scales, dd], axis=-1)
    return out.reshape(*lead, nb * Q3_K_BYTES)


def iq4_xs_blocks(codes: torch.Tensor, d: torch.Tensor, m: torch.Tensor) -> np.ndarray:
    """D-57 iq4xs packed codes [..., C/2] (`pack_4bit` of table indices 0..15), super-block scales d [..., C/256]
    (fp16) and sub-block scales m [..., C/32] (int8, -32..31) → IQ4_XS bytes [..., C/256 · 136], ggml's block_iq4_xs:
    d, scales_h (uint16: bits 2j, 2j+1 = the high two bits of m_j + 32), scales_l[4] (low nibbles of m_j + 32, sub-block
    2i in the low half of byte i), then per sub-block of 32 the 16 bytes Q4_0 and IQ4_NL use (j low, j+16 high)."""
    from lowbit.quant.ternary import unpack_4bit
    u = unpack_4bit(codes).contiguous().numpy()
    lead, nb = u.shape[:-1], u.shape[-1] // IQ4_XS_BLOCK
    u = u.reshape(-1, nb, 8, 2, 16)
    rows = u.shape[0]
    qs = (u[:, :, :, 0] | (u[:, :, :, 1] << 4)).reshape(rows, nb, 128)
    ls = (m.to(torch.int16) + 32).contiguous().numpy().astype(np.uint16).reshape(rows, nb, 8)
    scales_l = ((ls[:, :, 0::2] & 15) | ((ls[:, :, 1::2] & 15) << 4)).astype(np.uint8)
    scales_h = np.zeros((rows, nb), dtype=np.uint16)
    for j in range(8):
        scales_h |= ((ls[:, :, j] >> 4) & 3) << (2 * j)
    dd = d.to(torch.float16).contiguous().numpy().reshape(rows, nb, 1).view(np.uint8)
    out = np.concatenate([dd, scales_h.astype("<u2").reshape(rows, nb, 1).view(np.uint8), scales_l, qs], axis=-1)
    return out.reshape(*lead, nb * IQ4_XS_BYTES)


def packed_blocks(p: dict, gguf) -> tuple:
    """(GGUF type, bytes) of one packed tensor in its native type: ternary and int2 → Q2_0, int4 → Q4_0, int3 → Q3_K."""
    T = gguf.GGMLQuantizationType
    if "blocks" in p:  # D-58: a codebook student's blocks are already the GGUF's
        b = p["blocks"].numpy() if torch.is_tensor(p["blocks"]) else p["blocks"]
        return getattr(T, p["format"]), np.ascontiguousarray(b.reshape(*b.shape[:-2], -1))
    method = p.get("method")
    if method == "int4":
        return T.Q4_0, q4_0_blocks(p["codes"], p["scales"])
    if method in ("int3", "int3nl"):
        return T.Q3_K, q3_k_blocks(p["codes"], p["scales"], p["sub_scales"])
    if method == "iq4nl":  # D-57: IQ4_NL's block is Q4_0's (fp16 d, 16 bytes of nibbles); the nibbles index the table
        return T.IQ4_NL, q4_0_blocks(p["codes"], p["scales"])
    if method == "iq4xs":
        return T.IQ4_XS, iq4_xs_blocks(p["codes"], p["scales"], p["sub_scales"])
    return T.Q2_0, q2_0_blocks(p["codes"], p["scales"], p["group_size"])


def dequant_native(raw: np.ndarray, qtype, gguf) -> np.ndarray:
    """A native tensor's bytes → float32: Q2_0 by `dequant_q2_0` (gguf-py has no Q2_0), the rest by gguf-py."""
    if qtype == gguf.GGMLQuantizationType.Q2_0:
        return dequant_q2_0(raw)
    return gguf.quants.dequantize(raw, qtype)


def bf16_bytes(t: torch.Tensor) -> np.ndarray:
    t = t.to(torch.bfloat16).contiguous()
    return t.view(torch.int16).numpy().view(np.uint8).reshape(*t.shape[:-1], t.shape[-1] * 2)


_LAYER = {"input_layernorm.weight": "attn_norm.weight", "post_attention_layernorm.weight": "ffn_norm.weight",
          "self_attn.q_proj.weight": "attn_q.weight", "self_attn.k_proj.weight": "attn_k.weight",
          "self_attn.v_proj.weight": "attn_v.weight", "self_attn.o_proj.weight": "attn_output.weight",
          "self_attn.q_norm.weight": "attn_q_norm.weight", "self_attn.k_norm.weight": "attn_k_norm.weight",
          "mlp.gate.weight": "ffn_gate_inp.weight", "mlp.gate_proj.weight": "ffn_gate.weight",
          "mlp.up_proj.weight": "ffn_up.weight", "mlp.down_proj.weight": "ffn_down.weight"}
_TOP = {"model.embed_tokens.weight": "token_embd.weight", "model.norm.weight": "output_norm.weight",
        "lm_head.weight": "output.weight"}


def gguf_name(hf: str) -> str:
    if hf in _TOP:
        return _TOP[hf]
    m = re.fullmatch(r"model\.layers\.(\d+)\.(.+)", hf)
    if not m or m.group(2) not in _LAYER:
        raise KeyError(f"no GGUF name for {hf}")
    return f"blk.{m.group(1)}.{_LAYER[m.group(2)]}"


def student_tensors(dense: dict[str, torch.Tensor], packed: dict[str, dict], gguf,
                    ternary_as: str = "q2_0", dense_as: str = "bf16") -> dict[str, tuple]:
    """GGUF name → (type, bytes) for a student: `dense` holds plain tensors by HF name; `packed` holds ternary
    weights by HF name (without `.weight`) as {codes, scales, group_size}. MoE experts (`...mlp.experts.gate_up_proj`,
    `...down_proj`) split the fused gate/up rows the way llama.cpp's converter does (gate first).

    `ternary_as` "q2_0" writes each packed tensor in its native type (`packed_blocks`; D-51's N-bit formats too),
    "bf16" the same weights dequantised. `dense_as` sets the other matrices (attention, embedding): "bf16" writes them exactly as trained; "q8_0"
    quantises them post hoc, as the Q8 teacher GGUF does. On CPU the bf16 attention and embedding are ~85% of the
    bytes a token reads (the ternary experts ~10%), so Q8_0 there makes the student faster than the Q8 teacher.
    Norms and the router stay F32 either way."""
    T = gguf.GGMLQuantizationType
    out = {}
    for hf, t in dense.items():
        name = gguf_name(hf)
        if t.dim() == 1 or name.endswith("ffn_gate_inp.weight"):
            out[name] = (T.F32, t.float().contiguous().numpy())
        elif dense_as == "q8_0":
            out[name] = (T.Q8_0, gguf.quants.quantize(t.float().contiguous().numpy(), T.Q8_0))
        else:
            out[name] = (T.BF16, bf16_bytes(t))
    for hf, p in packed.items():
        if ternary_as == "bf16":  # the control: the same weights dequantised, no Q2_0 kernel involved
            from lowbit.quant.ternary import dequant_packed
            w = dequant_packed(p, torch.bfloat16)
            blocks, qt = bf16_bytes(w), T.BF16
        else:
            qt, blocks = packed_blocks(p, gguf)
        m = re.fullmatch(r"model\.layers\.(\d+)\.mlp\.experts\.(gate_up|down)_proj", hf)
        if m:
            i = m.group(1)
            if m.group(2) == "gate_up":
                half = blocks.shape[1] // 2
                out[f"blk.{i}.ffn_gate_exps.weight"] = (qt, np.ascontiguousarray(blocks[:, :half]))
                out[f"blk.{i}.ffn_up_exps.weight"] = (qt, np.ascontiguousarray(blocks[:, half:]))
            else:
                out[f"blk.{i}.ffn_down_exps.weight"] = (qt, blocks)
        else:
            out[gguf_name(hf + ".weight")] = (qt, blocks)
    return out


def write_gguf(template: str | Path, out: str | Path, tensors: dict[str, tuple], gguf,
               metadata: dict | None = None) -> dict:
    """Write `tensors` (GGUF name → (type, bytes or float32 array)) in the template's tensor order, with the
    template's metadata. Every template tensor must be supplied, and nothing else. `metadata` (key → str or int)
    replaces or adds general.* keys, e.g. a release's name and its base model."""
    metadata = metadata or {}
    reader = gguf.GGUFReader(str(template))
    arch = reader.fields["general.architecture"].contents()
    names = [t.name for t in reader.tensors]
    missing, extra = set(names) - set(tensors), set(tensors) - set(names)
    if missing or extra:
        raise ValueError(f"tensor sets differ: missing {sorted(missing)[:5]}, extra {sorted(extra)[:5]}")
    writer = gguf.GGUFWriter(str(out), arch=arch)
    for field in reader.fields.values():
        if field.name == gguf.Keys.General.ARCHITECTURE or field.name.startswith("GGUF.") or field.name in metadata:
            continue
        val_type = field.types[0]
        sub_type = field.types[-1] if val_type == gguf.GGUFValueType.ARRAY else None
        value = field.contents()
        if field.name == gguf.Keys.General.FILE_TYPE:  # a label: the type most of the expert bytes are in
            T, kinds = gguf.GGMLQuantizationType, {q for q, _ in tensors.values()}
            cbk = sorted(k.name for k in kinds if k.name in ("IQ2_XXS", "IQ2_XS", "IQ2_S", "IQ3_XXS", "IQ3_S"))
            if cbk:  # D-58: label by the codebook type the gate/up experts are in
                value = int(getattr(gguf.LlamaFileType, f"MOSTLY_{cbk[0]}"))
            elif T.Q3_K in kinds:
                value = int(gguf.LlamaFileType.MOSTLY_Q3_K_M)  # D-51 int3: Q3_K gate/up, Q4_0 down (D-57: or IQ4_NL)
            elif T.IQ4_XS in kinds:
                value = int(gguf.LlamaFileType.MOSTLY_IQ4_XS)  # D-57 iq4xs: IQ4_XS gate/up, IQ4_NL down
            elif T.Q2_0 in kinds:
                value = int(gguf.LlamaFileType.MOSTLY_Q2_0)
            elif T.Q4_0 in kinds:
                value = int(gguf.LlamaFileType.MOSTLY_Q4_0)
            elif T.IQ4_NL in kinds:
                value = int(gguf.LlamaFileType.MOSTLY_IQ4_NL)
        writer.add_key_value(field.name, value, val_type, sub_type=sub_type)
    for key, value in metadata.items():
        vt = gguf.GGUFValueType.UINT32 if isinstance(value, int) else gguf.GGUFValueType.STRING
        writer.add_key_value(key, value, vt)
    for name in names:
        qtype, data = tensors[name]
        writer.add_tensor_info(name, data.shape, data.dtype, data.nbytes, raw_dtype=qtype)
    writer.write_header_to_file()
    writer.write_kv_data_to_file()
    writer.write_ti_data_to_file()
    for name in names:
        writer.write_tensor_data(tensors[name][1])
    writer.close()
    return {"tensors": len(names), "bytes": Path(out).stat().st_size}


def check_round_trip(out: str | Path, packed: dict[str, dict], gguf) -> dict:
    """Acceptance: every scale an fp16 normal, and the GGUF's packed tensors (Q2_0; D-51: also Q4_0, Q3_K) dequantise
    bit-exactly to the packed student's weights (`dequant_packed`)."""
    from lowbit.quant.ternary import dequant_packed
    reader = gguf.GGUFReader(str(out))
    by = {t.name: t for t in reader.tensors}
    cbp = {k: p for k, p in packed.items() if "blocks" in p}
    if cbp:  # D-58: codebook tensors must be the packed blocks byte for byte
        bad = []
        for hf, p in cbp.items():
            i, name = re.fullmatch(r"model\.layers\.(\d+)\.mlp\.experts\.(gate_up|down)_proj", hf).groups()
            b = p["blocks"].numpy()
            parts = {"gate": b[:, :b.shape[1] // 2], "up": b[:, b.shape[1] // 2:]} if name == "gate_up" else {"down": b}
            for k, want in parts.items():
                got = np.asarray(by[f"blk.{i}.ffn_{k}_exps.weight"].data)
                if got.tobytes() != np.ascontiguousarray(want).tobytes():
                    bad.append(f"blk.{i}.ffn_{k}_exps")
        packed = {k: p for k, p in packed.items() if "blocks" not in p}
        if not packed:
            return {"codebook_tensors": 2 * len(cbp) // 2, "mismatched_tensors": bad, "bit_exact": not bad}
    scales = torch.cat([p["scales"].float().flatten() for p in packed.values()])
    nz = scales[scales != 0]
    res = {"scale_min_nonzero": float(nz.abs().min()) if nz.numel() else None, "zero_scales": int((scales == 0).sum()),
           "scales_below_fp16_normal": int(((nz.abs() < FP16_MIN_NORMAL)).sum()), "mismatched_tensors": []}

    def deq(name: str) -> np.ndarray:
        return dequant_native(np.asarray(by[name].data), by[name].tensor_type, gguf)

    for hf, p in packed.items():
        ref = dequant_packed(p, torch.float32).numpy()
        m = re.fullmatch(r"model\.layers\.(\d+)\.mlp\.experts\.(gate_up|down)_proj", hf)
        if m and m.group(2) == "gate_up":
            got = np.concatenate([deq(f"blk.{m.group(1)}.ffn_{k}_exps.weight") for k in ("gate", "up")], axis=1)
        elif m:
            got = deq(f"blk.{m.group(1)}.ffn_down_exps.weight")
        else:
            got = deq(gguf_name(hf + ".weight"))
        if got.shape != ref.shape or not np.array_equal(got, ref):
            res["mismatched_tensors"].append(hf)
    res["bit_exact"] = not res["mismatched_tensors"] and res["scales_below_fp16_normal"] == 0
    return res