#!/usr/bin/env python3
"""
Phase 3 Frozen-Expert Q2_0 Dequantizer + Byte-Level GGUF Dense/Router/Shared Transplant
=======================================================================================
Implements `PLAN.md` §2.2, §2.4, and §7.1 (with all B1 & B4 fixes from `audits/opus_audit_phase3_prep.md`):
1. Exact `qwen4exp.cpp` + `llama-arch.cpp` tensor classification across trunk layers `0..47` and MTP `48`:
   - `frozen_routed_experts`: `blk.{l}.ffn_{gate,up,down,gate_up}_exps.{weight,bias,scale}` (`Q2_0` type 42)
   - `frozen_ple_and_embd`: `per_layer_token_embd.weight`, `token_embd.weight`, `output.weight`, `blk.48.*`, `blk.{l}.ssm_a`
   - `trainable_dense_router_shared`: trunk (`0..47`) residual writers (`ssm_out`, `attn_out`/`attn_output`,
     `ffn_down_shexp`), router gates (`ffn_gate_inp`, `ffn_gate_inp_shexp`), shared experts (`ffn_{gate,up}_shexp`),
     QSA/GDN projections & norms (`attn_{q,k,v,qkv,gate,q_norm,k_norm}`, `indexer.{q_proj,k_proj,q_norm,k_norm}`,
     `ssm_{conv1d,dt,beta,alpha,norm}`), PLE layer projections (`ple_{key,value,norm_key,norm_query,norm_conv,conv1d}`),
     and hyper-connection mixers (`hc_{attn,ffn}_{norm,down,up,inject}`, `output_hc_{norm,down,up}`).
2. Exact `GGML_TYPE_Q2_0` (type 42, `QK2_0 = 64` weights in 18 bytes) quantize/dequantize matching
   `quantize_row_q2_0_ref` and `dequantize_row_q2_0` in `ggml-quants.c` (`d = amax`, `q = clamp(roundf(w/d) + 1, 0, 3)`,
   `byte_index = j // 4`, `bit_offset = (j % 4) * 2`, `y = (q - 1) * d` in `{-d, 0, +d, +2d}`), returning
   contiguous `np.ndarray` or `array.array('f')` buffers (never boxed Python float lists) and verified against
   compiled C `ggml-quants.c` in `--self-test`.
3. Whitelist-only byte-level GGUF tensor transplant (`transplant_gguf_tensors`):
   - Refuses ANY tensor not in `trainable_dense_router_shared` (including `other_frozen`, `frozen_routed_experts`,
     and `frozen_ple_and_embd`).
   - Verifies both `ggml_type` and exact byte size for every replacement tensor, checks disk free space,
     and cleans up `.tmp` on any failure.
"""

import argparse
from array import array
import ctypes
import math
import os
import re
import shutil
import struct
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Set, Tuple, Union

# `slice_qwen4exp_gguf` stayed at the repo's top-level scripts/ when the persona scripts moved into
# persona/scripts/, so this folder alone no longer resolves the import.
sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "scripts"))
from slice_qwen4exp_gguf import (
    GGML_TYPE_TRAITS,
    GGUF_MAGIC,
    GGUF_TYPE_STRING,
    GGUF_TYPE_UINT32,
    align_offset,
    compute_tensor_nbytes,
    parse_gguf_header,
    pwrite_all,
)

try:
    import numpy as np
except ImportError:
    np = None  # type: ignore

GGML_TYPE_F32 = 0
GGML_TYPE_F16 = 1
GGML_TYPE_BF16 = 30
GGML_TYPE_Q2_0 = 42  # block_q2_0: QK2_0 = 64 weights, 2B fp16 scale d + 16B qs = 18 bytes/block

# Match individual or fused routed expert tensors (including any bias/scale companions) across all layers
FROZEN_EXPERT_RE = re.compile(
    r"^blk\.(\d+)\.ffn_(?:gate|up|down|gate_up)_exps(?:\.(?:weight|bias|scale))?$"
)
FROZEN_PLE_EMBD_NAMES: Set[str] = {
    "per_layer_token_embd.weight",
    "token_embd.weight",
    "output.weight",
}
FROZEN_SSM_A_RE = re.compile(r"^blk\.\d+\.ssm_a(?:\.weight)?$")

# Trunk layers 0..47 only (blk.48.* is the MTP head and stays frozen unless explicitly tuned separately)
TRUNK_LAYER_PAT = r"(?:[0-9]|[1-3][0-9]|4[0-7])"

# Strict whitelist of trainable trunk dense / router / shared-expert tensor names from qwen4exp.cpp & llama-arch.cpp
TRAINABLE_TENSOR_RE = re.compile(
    rf"^(?:"
    rf"output_hc_(?:norm|down|up)\.weight"
    rf"|blk\.{TRUNK_LAYER_PAT}\.(?:"
    rf"ffn_gate_inp(?:_shexp)?\.weight"
    rf"|ffn_(?:gate|up|down)_shexp\.weight"
    rf"|attn_(?:q|k|v|qkv|gate|out|output|q_norm|k_norm)\.weight"
    rf"|indexer\.(?:q_proj|k_proj|q_norm|k_norm)\.weight"
    rf"|ssm_(?:conv1d|beta|alpha|norm|out)\.weight"
    rf"|ssm_dt\.bias"
    rf"|hc_(?:attn|ffn)_(?:norm|down|up|inject)\.weight"
    rf"|ple_(?:key|value|norm_key|norm_query|norm_conv|conv1d)\.weight"
    rf")"
    rf")$"
)

# Minimal Swift-1.5 residual-writer + router subset (96 trunk writers + 96 router gates = 192 tensors)
SWIFT_WRITER_ROUTER_RE = re.compile(
    rf"^blk\.{TRUNK_LAYER_PAT}\.(?:"
    rf"ssm_out\.weight"
    rf"|attn_(?:out|output)\.weight"
    rf"|ffn_down_shexp\.weight"
    rf"|ffn_gate_inp(?:_shexp)?\.weight"
    rf")$"
)

# Pure 96 residual-writer subset (keeps router gates frozen to protect MoCE/LRU hit rate)
SWIFT_WRITERS_ONLY_RE = re.compile(
    rf"^blk\.{TRUNK_LAYER_PAT}\.(?:"
    rf"ssm_out\.weight"
    rf"|attn_(?:out|output)\.weight"
    rf"|ffn_down_shexp\.weight"
    rf")$"
)


def classify_qwen4exp_tensor(name: str, scope: str = "all_trunk_dense_router_shared") -> str:
    """
    Classifies a tensor name from a Qwen3.8-Flash-Next (`qwen4exp`) GGUF into:
      - "frozen_routed_experts"
      - "frozen_ple_and_embd"
      - "trainable_dense_router_shared"
      - "other_frozen"
    `scope` may be `"all_trunk_dense_router_shared"` (default) or `"swift_writers_plus_router"`.
    """
    if FROZEN_EXPERT_RE.fullmatch(name):
        return "frozen_routed_experts"
    if (
        name in FROZEN_PLE_EMBD_NAMES
        or name.startswith("mtp.")
        or name.startswith("nextn.")
        or name.startswith("blk.48.")
        or ".nextn." in name
        or FROZEN_SSM_A_RE.fullmatch(name)
    ):
        return "frozen_ple_and_embd"
    if scope == "swift_writers":
        if SWIFT_WRITERS_ONLY_RE.fullmatch(name):
            return "trainable_dense_router_shared"
        return "other_frozen"
    if scope == "swift_writers_plus_router":
        if SWIFT_WRITER_ROUTER_RE.fullmatch(name):
            return "trainable_dense_router_shared"
        return "other_frozen"
    if TRAINABLE_TENSOR_RE.fullmatch(name):
        return "trainable_dense_router_shared"
    return "other_frozen"


def f32_to_f16_scalar(val: float) -> float:
    """Round-trips a Python float through IEEE 754 binary16 (`<e`), matching `GGML_FP32_TO_FP16`."""
    return float(struct.unpack("<e", struct.pack("<e", float(val)))[0])


def c_roundf(val: float) -> float:
    """Matches C99 `roundf(x)` (round half-way cases away from zero, unlike Python's round-to-even)."""
    return math.copysign(math.floor(abs(val) + 0.5), val)


def dequantize_q2_0_bytes(raw_bytes: bytes, n_elements: int) -> Union["np.ndarray", array]:
    """
    Exact dequantization of `GGML_TYPE_Q2_0` (type 42) matching `dequantize_row_q2_0` in `ggml-quants.c` (L439):
      - Block size `QK2_0 = 64`, 18 bytes per block:
        - `d`: `ggml_half` (IEEE FP16, 2 bytes, `<e`)
        - `qs`: `uint8_t[16]` (16 bytes = 64 2-bit values, 4 consecutive elements per byte)
      - For `j in 0..63`:
        - `byte_index = j // 4` (`j >> 2`)
        - `bit_offset = (j % 4) * 2` (`(j & 3) * 2`)
        - `q = (qs[byte_index] >> bit_offset) & 0x03`
        - `y[i*64 + j] = (q - 1) * d`  (`00 -> -d, 01 -> 0, 10 -> +d, 11 -> +2d`)
    Returns a contiguous `np.ndarray` (`float32`) if NumPy is installed, else `array('f')`.
    """
    if n_elements % 64 != 0:
        raise ValueError(f"Q2_0 element count ({n_elements}) must be divisible by 64")
    n_blocks = n_elements // 64
    expected_nbytes = n_blocks * 18
    if len(raw_bytes) != expected_nbytes:
        raise ValueError(f"Q2_0 byte length mismatch: got {len(raw_bytes)}, expected {expected_nbytes}")

    if np is not None:
        raw_arr = np.frombuffer(raw_bytes, dtype=np.uint8).reshape(n_blocks, 18)
        d = raw_arr[:, :2].copy().view(np.float16).astype(np.float32)  # (n_blocks, 1)
        qs = raw_arr[:, 2:]  # (n_blocks, 16)
        # Interleave 4 consecutive 2-bit lanes per byte: j = 4*byte_index + lane
        lanes = np.empty((n_blocks, 16, 4), dtype=np.float32)
        lanes[:, :, 0] = ((qs >> 0) & 0x03).astype(np.float32) - 1.0
        lanes[:, :, 1] = ((qs >> 2) & 0x03).astype(np.float32) - 1.0
        lanes[:, :, 2] = ((qs >> 4) & 0x03).astype(np.float32) - 1.0
        lanes[:, :, 3] = ((qs >> 6) & 0x03).astype(np.float32) - 1.0
        return (lanes.reshape(n_blocks, 64) * d).reshape(-1)

    out = array("f", [0.0]) * n_elements
    for b in range(n_blocks):
        boff = b * 18
        d_val = float(struct.unpack_from("<e", raw_bytes, boff)[0])
        base = b * 64
        for byte_idx in range(16):
            byte_v = raw_bytes[boff + 2 + byte_idx]
            j0 = base + (byte_idx << 2)
            out[j0 + 0] = float(((byte_v >> 0) & 0x03) - 1) * d_val
            out[j0 + 1] = float(((byte_v >> 2) & 0x03) - 1) * d_val
            out[j0 + 2] = float(((byte_v >> 4) & 0x03) - 1) * d_val
            out[j0 + 3] = float(((byte_v >> 6) & 0x03) - 1) * d_val
    return out


def quantize_q2_0_f32(values: Union[Iterable[float], "np.ndarray", array]) -> bytes:
    """
    Exact `GGML_TYPE_Q2_0` reference quantizer matching `quantize_row_q2_0_ref` in `ggml-quants.c` (L74):
      - `amax = max(|x[j]| for j in 0..63)`
      - `d = amax` (stored as FP16 in `y[i].d`; `quantize_row_q2_0_ref` uses `const float d = amax; id = d > 0 ? 1.0f/d : 0.0f`)
      - `q = clamp((int)roundf(w * id) + 1, 0, 3)`
      - `y[i].qs[j / 4] |= (q << ((j % 4) * 2))`
    """
    vals = values if isinstance(values, (list, tuple, array)) else list(values)
    n_elements = len(vals)
    if n_elements % 64 != 0:
        raise ValueError(f"Input size {n_elements} must be divisible by 64")
    n_blocks = n_elements // 64
    buf = bytearray(n_blocks * 18)

    for b in range(n_blocks):
        base = b * 64
        amax = 0.0
        for j in range(64):
            # Cast through 32-bit float to match C `const float w = x[i*qk + j]`
            w = struct.unpack("<f", struct.pack("<f", float(vals[base + j])))[0]
            a = abs(w)
            if a > amax:
                amax = a
        d = amax
        id_val = struct.unpack("<f", struct.pack("<f", (1.0 / d) if d > 0.0 else 0.0))[0]
        boff = b * 18
        struct.pack_into("<e", buf, boff, d)

        for j in range(64):
            w = struct.unpack("<f", struct.pack("<f", float(vals[base + j])))[0]
            prod = struct.unpack("<f", struct.pack("<f", w * id_val))[0]
            q = int(c_roundf(prod)) + 1
            if q < 0:
                q = 0
            elif q > 3:
                q = 3
            byte_index = j >> 2
            bit_offset = (j & 3) * 2
            buf[boff + 2 + byte_index] |= (q & 0x03) << bit_offset
    return bytes(buf)


def load_writers_bin(bin_path: str) -> Dict[str, Tuple[int, bytes]]:
    """
    Loads a binary writer bundle written by `export_merged_writers_bin` or `heal_k256_q2_0`:
    `u32 count`, then for each entry: `u32 name_len, name_bytes, i32 ggml_type, u64 nbytes, raw_bytes`.
    """
    with open(bin_path, "rb") as fh:
        data = fh.read()
    pos = 0
    (count,) = struct.unpack_from("<I", data, pos)
    pos += 4
    out: Dict[str, Tuple[int, bytes]] = {}
    for _ in range(count):
        (nlen,) = struct.unpack_from("<I", data, pos)
        pos += 4
        name = data[pos : pos + nlen].decode("utf-8")
        pos += nlen
        gtype, nbytes = struct.unpack_from("<iQ", data, pos)
        pos += 12
        raw = data[pos : pos + nbytes]
        if len(raw) != nbytes:
            raise ValueError(f"Truncated writers_bin {bin_path} at {name!r}")
        pos += nbytes
        out[name] = (int(gtype), raw)
    return out


def transplant_gguf_tensors(
    in_gguf_path: str,
    out_gguf_path: str,
    replacements: Dict[str, Tuple[int, bytes]],
    scope: str = "all_trunk_dense_router_shared",
    allow_type_change: bool = False,
) -> Dict:
    """
    Whitelist-only byte-level GGUF transplant (`PLAN.md` §2.2 L64, SC117 method):
      - Refuses in-place overwrite.
      - Strictly whitelists `classify_qwen4exp_tensor(name, scope=scope) == "trainable_dense_router_shared"`
        (refuses `frozen_routed_experts`, `frozen_ple_and_embd`, AND `other_frozen`).
      - Requires every value in `replacements` to be a `(expected_ggml_type, raw_bytes)` tuple and
        verifies `expected_ggml_type == t_info["ggml_type"]` as well as exact byte length.
      - Checks disk free space before copying, writes via atomic `.tmp` + rename + directory fsync, and removes `.tmp` on error.
    """
    in_resolved = os.path.realpath(os.path.abspath(in_gguf_path))
    out_resolved = os.path.realpath(os.path.abspath(out_gguf_path))
    if in_resolved == out_resolved:
        raise ValueError(f"Refusing in-place overwrite: {in_gguf_path}")

    file_size = os.path.getsize(in_resolved)
    fd_in = os.open(in_resolved, os.O_RDONLY)
    try:
        hdr = parse_gguf_header(fd_in, file_size)
    finally:
        os.close(fd_in)

    tensor_by_name = {t["name"]: t for t in hdr["tensors"]}
    data_off = hdr["data_section_offset"]

    validated_writes: List[Tuple[str, int, bytes]] = []
    any_type_changed = False
    for t_name, item in replacements.items():
        if t_name not in tensor_by_name:
            raise KeyError(f"Replacement tensor {t_name!r} not found in {in_gguf_path}")
        cat = classify_qwen4exp_tensor(t_name, scope=scope)
        if cat != "trainable_dense_router_shared":
            raise ValueError(
                f"Refusing to overwrite non-whitelisted tensor {t_name!r} (category={cat}, scope={scope})"
            )
        t_info = tensor_by_name[t_name]
        if not (isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], int) and isinstance(item[1], (bytes, bytearray))):
            raise ValueError(
                f"Replacement for {t_name!r} must be a (ggml_type: int, raw_bytes: bytes) tuple"
            )
        given_type, raw_buf = item[0], bytes(item[1])
        if given_type != t_info["ggml_type"]:
            if not allow_type_change:
                raise ValueError(
                    f"ggml_type mismatch for {t_name!r}: replacement has type {given_type}, "
                    f"GGUF header expects type {t_info['ggml_type']}"
                )
            any_type_changed = True

        expected_nbytes = compute_tensor_nbytes(t_info["dims"], given_type)
        if len(raw_buf) != expected_nbytes:
            raise ValueError(
                f"Byte length mismatch for {t_name!r}: replacement has {len(raw_buf)} bytes, "
                f"expected {expected_nbytes} bytes (dims={t_info['dims']}, type={given_type})"
            )
        abs_off = data_off + t_info["rel_offset"]
        validated_writes.append((t_name, abs_off, raw_buf))

    out_dir = os.path.dirname(out_resolved) or "."
    os.makedirs(out_dir, exist_ok=True)
    extra_bytes = sum(len(item[1]) for item in replacements.values()) if any_type_changed else 0
    free_bytes = shutil.disk_usage(out_dir).free
    if free_bytes < file_size + extra_bytes + (64 * 1024 * 1024):
        raise OSError(
            f"Insufficient disk space in {out_dir}: need {file_size + extra_bytes} bytes, have {free_bytes} bytes free"
        )

    tmp_out = out_resolved + ".tmp"
    try:
        if not any_type_changed:
            shutil.copyfile(in_resolved, tmp_out)
            fd_out = os.open(tmp_out, os.O_RDWR)
            try:
                for _t_name, abs_off, raw_buf in validated_writes:
                    pwrite_all(fd_out, raw_buf, abs_off)
                os.fsync(fd_out)
            finally:
                os.close(fd_out)
        else:
            alignment = hdr["alignment"]
            new_tensors = []
            rel_off = 0
            for t in hdr["tensors"]:
                nm = t["name"]
                orig_nbytes = compute_tensor_nbytes(t["dims"], t["ggml_type"])
                is_present = (data_off + t["rel_offset"] + orig_nbytes) <= file_size
                if nm in replacements:
                    gtype, rbuf = replacements[nm]
                    nbytes = len(rbuf)
                else:
                    gtype = t["ggml_type"]
                    rbuf = None
                    nbytes = orig_nbytes
                new_tensors.append({
                    "name": nm,
                    "raw_name": t["raw_name"],
                    "n_dims": t["n_dims"],
                    "dims": t["dims"],
                    "ggml_type": gtype,
                    "new_rel_offset": rel_off,
                    "orig_abs_start": data_off + t["rel_offset"],
                    "orig_nbytes": orig_nbytes,
                    "new_nbytes": nbytes,
                    "is_present": is_present,
                    "rbuf": rbuf,
                })
                rel_off = align_offset(rel_off + nbytes, alignment)

            hb = bytearray(struct.pack("<IIQQ", GGUF_MAGIC, hdr["version"], hdr["n_tensors"], hdr["n_kv"]))
            for kv in hdr["kvs"]:
                hb += kv["raw_key"] + kv["raw_vtype"] + kv["raw_val"]
            for nt in new_tensors:
                hb += nt["raw_name"]
                hb += struct.pack(f"<I{nt['n_dims']}QIQ", nt["n_dims"], *nt["dims"], nt["ggml_type"], nt["new_rel_offset"])
            new_data_off = align_offset(len(hb), alignment)
            hb += b"\x00" * (new_data_off - len(hb))

            present_tensors = [nt for nt in new_tensors if nt["is_present"]]
            final_size = new_data_off + (
                align_offset(present_tensors[-1]["new_rel_offset"] + present_tensors[-1]["new_nbytes"], alignment)
                if present_tensors else 0
            )
            fd_in = os.open(in_resolved, os.O_RDONLY)
            fd_out = os.open(tmp_out, os.O_CREAT | os.O_RDWR | os.O_TRUNC, 0o644)
            try:
                os.ftruncate(fd_out, final_size)
                pwrite_all(fd_out, bytes(hb), 0)
                for nt in present_tensors:
                    dst_off = new_data_off + nt["new_rel_offset"]
                    if nt["rbuf"] is not None:
                        pwrite_all(fd_out, bytes(nt["rbuf"]), dst_off)
                    else:
                        rem = nt["orig_nbytes"]
                        src_off = nt["orig_abs_start"]
                        while rem > 0:
                            step = min(rem, 16 * 1024 * 1024)
                            chunk = os.pread(fd_in, step, src_off)
                            if not chunk:
                                raise IOError(f"Unexpected EOF reading {nt['name']}")
                            pwrite_all(fd_out, chunk, dst_off)
                            src_off += len(chunk)
                            dst_off += len(chunk)
                            rem -= len(chunk)
                os.fsync(fd_out)
            finally:
                os.close(fd_out)
                os.close(fd_in)
            file_size = final_size

        os.replace(tmp_out, out_resolved)
        in_exp = in_resolved + ".experts"
        out_exp = out_resolved + ".experts"
        if os.path.exists(in_exp) and not os.path.exists(out_exp):
            os.symlink(in_exp, out_exp)
        dir_fd = os.open(out_dir, os.O_RDONLY)
        try:
            os.fsync(dir_fd)
        finally:
            os.close(dir_fd)
    except Exception:
        if os.path.exists(tmp_out):
            try:
                os.unlink(tmp_out)
            except OSError:
                pass
        raise

    return {
        "in_path": in_gguf_path,
        "out_path": out_gguf_path,
        "file_size": file_size,
        "total_tensors": len(hdr["tensors"]),
        "transplanted_tensors_count": len(validated_writes),
        "transplanted_tensor_names": [x[0] for x in validated_writes],
    }


def run_self_test() -> None:
    # 1. Exact qwen4exp.cpp + llama-arch.cpp tensor classification checks (including fused gate_up_exps & dotted indexer)
    for frozen_exp in (
        "blk.0.ffn_gate_exps.weight",
        "blk.47.ffn_down_exps.weight",
        "blk.12.ffn_gate_up_exps.weight",
        "blk.3.ffn_down_exps.scale",
    ):
        assert classify_qwen4exp_tensor(frozen_exp) == "frozen_routed_experts", frozen_exp

    for frozen_other in (
        "per_layer_token_embd.weight",
        "token_embd.weight",
        "output.weight",
        "blk.0.ssm_a",
        "blk.48.attn_qkv.weight",
        "blk.48.nextn.eh_proj.weight",
    ):
        assert classify_qwen4exp_tensor(frozen_other) == "frozen_ple_and_embd", frozen_other

    for trainable_name in (
        "output_hc_norm.weight",
        "output_hc_down.weight",
        "output_hc_up.weight",
        "blk.0.ffn_gate_inp.weight",
        "blk.0.ffn_gate_inp_shexp.weight",
        "blk.0.ffn_down_shexp.weight",
        "blk.0.ssm_out.weight",
        "blk.0.ssm_dt.bias",
        "blk.3.attn_out.weight",
        "blk.3.attn_q_norm.weight",
        "blk.3.attn_k_norm.weight",
        "blk.3.indexer.q_proj.weight",
        "blk.3.indexer.k_proj.weight",
        "blk.3.indexer.q_norm.weight",
        "blk.3.indexer.k_norm.weight",
        "blk.2.ple_norm_key.weight",
        "blk.2.ple_conv1d.weight",
        "blk.0.hc_attn_up.weight",
    ):
        assert classify_qwen4exp_tensor(trainable_name) == "trainable_dense_router_shared", trainable_name

    # Unknown tensor names classify as other_frozen and are REFUSED by transplant_gguf_tensors
    assert classify_qwen4exp_tensor("blk.0.unknown_tensor.weight") == "other_frozen"

    # 2. Verify Q2_0 codec on all 4 codes {0->-d, 1->0, 2->+d, 3->+2d} and against compiled C ggml-quants.c
    d_test = f32_to_f16_scalar(0.25)
    # Construct a raw 18-byte Q2_0 block where byte 0 has codes (0, 1, 2, 3) -> 0 | (1<<2) | (2<<4) | (3<<6) = 0xE4
    raw_blk = struct.pack("<e", d_test) + bytes([0xE4] + [0x55] * 15)  # 0x55 = codes (1,1,1,1) -> 0.0
    deq_blk = list(dequantize_q2_0_bytes(raw_blk, 64))
    assert deq_blk[:4] == [-0.25, 0.0, 0.25, 0.50], f"Wrong Q2_0 lane decoding: {deq_blk[:4]}"
    assert all(v == 0.0 for v in deq_blk[4:]), f"Non-zero tail: {deq_blk[4:]}"

    # Compile upstream C quantize_row_q2_0_ref & dequantize_row_q2_0 from ggml-quants.c and compare on 16,384 floats
    c_src = r"""
    #include <stdint.h>
    #include <math.h>
    typedef struct {
        _Float16 d;
        uint8_t qs[16];
    } block_q2_0;

    void quantize_row_q2_0_ref(const float * x, block_q2_0 * y, int64_t k) {
        const int qk = 64;
        const int nb = k / qk;
        for (int i = 0; i < nb; i++) {
            float amax = 0.0f;
            for (int j = 0; j < qk; j++) {
                const float a = fabsf(x[i*qk + j]);
                if (a > amax) amax = a;
            }
            const float d = amax;
            const float id = d > 0.0f ? 1.0f / d : 0.0f;
            y[i].d = (_Float16)d;
            for (int j = 0; j < qk / 4; ++j) y[i].qs[j] = 0;
            for (int j = 0; j < qk; ++j) {
                const float w = x[i*qk + j];
                int q = (int)roundf(w * id) + 1;
                if (q < 0) q = 0;
                if (q > 3) q = 3;
                y[i].qs[j / 4] |= ((uint8_t)q << ((j % 4) * 2));
            }
        }
    }

    void dequantize_row_q2_0(const block_q2_0 * x, float * y, int64_t k) {
        const int qk = 64;
        const int nb = k / qk;
        for (int i = 0; i < nb; i++) {
            const float d = (float)(x[i].d);
            for (int j = 0; j < qk; ++j) {
                const uint8_t q = (x[i].qs[j / 4] >> ((j % 4) * 2)) & 0x03;
                y[i*qk + j] = ((int)q - 1) * d;
            }
        }
    }
    """
    with tempfile.TemporaryDirectory() as tmpd:
        c_path = os.path.join(tmpd, "q2_ref.c")
        so_path = os.path.join(tmpd, "q2_ref.so")
        Path(c_path).write_text(c_src)
        cc_bin = shutil.which("clang") or shutil.which("gcc") or shutil.which("cc")
        if not cc_bin:
            raise RuntimeError("No C compiler (clang/gcc/cc) found for Q2_0 golden self-test")
        subprocess.run([cc_bin, "-O2", "-shared", "-fPIC", c_path, "-o", so_path], check=True)
        lib = ctypes.CDLL(so_path)

        n_test = 4096
        import random
        rng = random.Random(20261014)
        test_floats = [rng.gauss(0.0, 1.2) for _ in range(n_test)]
        c_in = (ctypes.c_float * n_test)(*test_floats)
        c_qbuf = (ctypes.c_ubyte * ((n_test // 64) * 18))()
        c_deq = (ctypes.c_float * n_test)()

        lib.quantize_row_q2_0_ref(c_in, c_qbuf, ctypes.c_int64(n_test))
        # Also inject explicit code-3 (0xFF) bytes in block 0 to verify GSQ full-4-code dequantization
        c_qbytes = bytearray(bytes(c_qbuf))
        c_qbytes[2] = 0xE4  # codes 0, 1, 2, 3
        c_qbytes[3] = 0xFF  # codes 3, 3, 3, 3
        c_qbuf_mod = (ctypes.c_ubyte * len(c_qbytes))(*c_qbytes)
        lib.dequantize_row_q2_0(c_qbuf_mod, c_deq, ctypes.c_int64(n_test))

        py_qbytes = quantize_q2_0_f32(test_floats)
        assert py_qbytes == bytes(c_qbuf), "Python quantize_q2_0_f32 bytes differed from C quantize_row_q2_0_ref!"

        py_deq = dequantize_q2_0_bytes(bytes(c_qbytes), n_test)
        max_err = max(abs(float(a) - float(b)) for a, b in zip(py_deq, c_deq))
        assert max_err == 0.0, f"Python dequantize_q2_0_bytes differed from C dequantize_row_q2_0: max_err={max_err}"

        # 3. Synthetic GGUF byte-level transplant test (verifies whitelist + type check + bit-identical experts)
        def _pack_str(s: str) -> bytes:
            b = s.encode("utf-8")
            return struct.pack("<Q", len(b)) + b

        in_p = os.path.join(tmpd, "in.gguf")
        out_p = os.path.join(tmpd, "out.gguf")

        router_orig = [float(i) for i in range(64 * 4)]
        router_new = [float(i) + 10.0 for i in range(64 * 4)]
        t0_bytes = struct.pack(f"<{len(router_orig)}f", *router_orig)
        t0_new_bytes = struct.pack(f"<{len(router_new)}f", *router_new)
        t1_bytes = quantize_q2_0_f32([0.5] * (64 * 2 * 4))

        kv1 = _pack_str("general.architecture") + struct.pack("<I", GGUF_TYPE_STRING) + _pack_str("qwen4exp")
        kv2 = _pack_str("qwen4exp.expert_count") + struct.pack("<I", GGUF_TYPE_UINT32) + struct.pack("<I", 4)
        off0 = 0
        off1 = align_offset(len(t0_bytes), 32)
        off2 = align_offset(off1 + len(t1_bytes), 32)

        ti0 = _pack_str("blk.0.ffn_gate_inp.weight") + struct.pack("<I2QIQ", 2, 64, 4, GGML_TYPE_F32, off0)
        ti1 = _pack_str("blk.0.ffn_gate_up_exps.weight") + struct.pack("<I3QIQ", 3, 64, 2, 4, GGML_TYPE_Q2_0, off1)
        ti2 = _pack_str("blk.0.unknown_other.weight") + struct.pack("<I2QIQ", 2, 64, 4, GGML_TYPE_F32, off2)
        hdr_raw = struct.pack("<IIQQ", GGUF_MAGIC, 3, 3, 2) + kv1 + kv2 + ti0 + ti1 + ti2
        data_off = align_offset(len(hdr_raw), 32)

        with open(in_p, "wb") as f:
            f.write(hdr_raw + b"\x00" * (data_off - len(hdr_raw)))
            f.seek(data_off + off0)
            f.write(t0_bytes)
            f.seek(data_off + off1)
            f.write(t1_bytes)
            f.seek(data_off + off2)
            f.write(t0_bytes)

        # Verify refusal on fused frozen expert tensor, other_frozen tensor, wrong ggml_type, AND bare bytes
        for bad_rep in (
            {"blk.0.ffn_gate_up_exps.weight": (GGML_TYPE_Q2_0, t1_bytes)},
            {"blk.0.unknown_other.weight": (GGML_TYPE_F32, t0_new_bytes)},
            {"blk.0.ffn_gate_inp.weight": (GGML_TYPE_F16, t0_new_bytes)},
            {"blk.0.ffn_gate_inp.weight": t0_new_bytes},  # bare bytes without type tuple must fail
        ):
            try:
                transplant_gguf_tensors(in_p, out_p, bad_rep)
                raise AssertionError(f"Should have refused {bad_rep.keys()}")
            except ValueError:
                pass

        rep = transplant_gguf_tensors(
            in_p, out_p, {"blk.0.ffn_gate_inp.weight": (GGML_TYPE_F32, t0_new_bytes)}
        )
        assert rep["transplanted_tensors_count"] == 1
        in_all = Path(in_p).read_bytes()
        out_all = Path(out_p).read_bytes()
        assert in_all[:data_off] == out_all[:data_off]
        assert in_all[data_off + off1 :] == out_all[data_off + off1 :]
        assert out_all[data_off : data_off + len(t0_bytes)] == t0_new_bytes

    print("scripts/phase3_lora_and_transplant.py self-test: PASS")


def main() -> None:
    parser = argparse.ArgumentParser(description="Phase 3 Frozen-Expert Q2_0 Dequant & GGUF Transplant")
    parser.add_argument("--self-test", action="store_true", help="Run unit self-tests and exit")
    args = parser.parse_args()
    if args.self_test:
        run_self_test()
        return


if __name__ == "__main__":
    main()