#!/usr/bin/env python3
"""SDQ MoE kernel 的 NumPy 對照實作（單元測試）。

用途：拿 C++ 端 `epq-chat --selftest IL:IE:SEED:OUT.bin` 算出來的結果，
用純 NumPy 從 GGUF 的同一段位元組重算一次，確認分頁 + 量化點積 + SiLU + 加權
這條路徑的數學完全正確。

實作完全照抄 ggml 的定義：
  * dequantize_row_q4_K / dequantize_row_q5_K / dequantize_row_q6_K
  * quantize_row_q8_K_ref（Q8_K activation）
  * SiLU(x) = x * sigmoid(x)
"""

from __future__ import annotations

import argparse
import json
import struct
from pathlib import Path

import numpy as np

QK_K = 256


def get_scale_min_k4(j: int, q: np.ndarray) -> tuple[int, int]:
    if j < 4:
        return int(q[j] & 63), int(q[j + 4] & 63)
    return ((int(q[j + 4] & 0xF) | ((int(q[j - 4]) >> 6) << 4)),
            ((int(q[j + 4]) >> 4) | ((int(q[j]) >> 6) << 4)))


def fp16_to_fp32(raw: bytes) -> float:
    return struct.unpack("<e", raw)[0]


def dequant_q4_K(raw: bytes, n: int) -> np.ndarray:
    nb = n // QK_K
    out = np.empty(n, dtype=np.float32)
    for i in range(nb):
        d = fp16_to_fp32(raw[i * 144:i * 144 + 2])
        dmin = fp16_to_fp32(raw[i * 144 + 2:i * 144 + 4])
        scales = np.frombuffer(raw, dtype=np.uint8, count=12, offset=i * 144 + 4)
        qs = np.frombuffer(raw, dtype=np.uint8, count=128, offset=i * 144 + 16)
        base = i * QK_K
        is_ = 0
        for j in range(0, QK_K, 64):
            sc1, m1 = get_scale_min_k4(is_ + 0, scales)
            sc2, m2 = get_scale_min_k4(is_ + 1, scales)
            d1, mm1 = d * sc1, dmin * m1
            d2, mm2 = d * sc2, dmin * m2
            q = qs[j // 2: j // 2 + 32]
            out[base + j: base + j + 32] = d1 * (q & 0xF).astype(np.float32) - mm1
            out[base + j + 32: base + j + 64] = d2 * (q >> 4).astype(np.float32) - mm2
            is_ += 2
    return out


def dequant_q5_K(raw: bytes, n: int) -> np.ndarray:
    nb = n // QK_K
    out = np.empty(n, dtype=np.float32)
    for i in range(nb):
        o = i * 176
        d = fp16_to_fp32(raw[o:o + 2])
        dmin = fp16_to_fp32(raw[o + 2:o + 4])
        scales = np.frombuffer(raw, dtype=np.uint8, count=12, offset=o + 4)
        qh = np.frombuffer(raw, dtype=np.uint8, count=32, offset=o + 16)
        qs = np.frombuffer(raw, dtype=np.uint8, count=128, offset=o + 48)
        base = i * QK_K
        is_ = 0
        u1, u2 = 1, 2
        for j in range(0, QK_K, 64):
            sc1, m1 = get_scale_min_k4(is_ + 0, scales)
            sc2, m2 = get_scale_min_k4(is_ + 1, scales)
            d1, mm1 = d * sc1, dmin * m1
            d2, mm2 = d * sc2, dmin * m2
            q = qs[j // 2: j // 2 + 32]
            # ggml 的迴圈只推進 ql，qh 永遠從 l=0 開始讀
            h1 = ((qh[0:32] & u1) != 0).astype(np.float32) * 16
            h2 = ((qh[0:32] & u2) != 0).astype(np.float32) * 16
            out[base + j: base + j + 32] = d1 * ((q & 0xF).astype(np.float32) + h1) - mm1
            out[base + j + 32: base + j + 64] = d2 * ((q >> 4).astype(np.float32) + h2) - mm2
            is_ += 2
            u1 <<= 2
            u2 <<= 2
    return out


def dequant_q6_K(raw: bytes, n: int) -> np.ndarray:
    """逐位元組對照 ggml/src/ggml-quants.c 的 dequantize_row_q6_K。

    block_q6_K（210 bytes）：ql[128] qh[64] scales[16(int8)] d(fp16)
    對每個 128 元素半區塊（n 迴圈，ql += 64 / qh += 32 / sc += 8）：
        q1 = (ql[l]    & 0xF) | ((qh[l] >> 0 & 3) << 4) → y[l]
        q2 = (ql[l+32] & 0xF) | ((qh[l] >> 2 & 3) << 4) → y[l+32]
        q3 = (ql[l]    >> 4)  | ((qh[l] >> 4 & 3) << 4) → y[l+64]
        q4 = (ql[l+32] >> 4)  | ((qh[l] >> 6 & 3) << 4) → y[l+96]
        縮放：y[l] 用 sc[is+0]、y[l+32] 用 sc[is+2]、y[l+64] 用 sc[is+4]、y[l+96] 用 sc[is+6]，
              其中 is = l / 16（l<16 用偶數、l>=16 用緊接的奇數）→ 等價於「每 16 個元素一個 scale」。
    """
    nb = n // QK_K
    out = np.empty(n, dtype=np.float32)
    for i in range(nb):
        o = i * 210
        ql = np.frombuffer(raw, dtype=np.uint8, count=128, offset=o)
        qh = np.frombuffer(raw, dtype=np.uint8, count=64, offset=o + 128)
        sc = np.frombuffer(raw, dtype=np.int8, count=16, offset=o + 192).astype(np.float32)
        d = fp16_to_fp32(raw[o + 208:o + 210])
        base = i * QK_K
        for half in range(2):
            qa = ql[half * 64: half * 64 + 64].astype(np.int32)   # ql[l], ql[l+32]
            ha = qh[half * 32: half * 32 + 32].astype(np.int32)   # qh[l]
            sa = sc[half * 8: half * 8 + 8]                       # 8 個 scale，每 16 元素一個
            scale = np.repeat(sa, 16)                             # 長度 128，對應整個 128 元素半區塊

            lo = qa[:32] & 0xF        # ql[l]   低 4 bit
            hi = qa[32:] & 0xF        # ql[l+32] 低 4 bit
            blo = qa[:32] >> 4        # ql[l]   高 4 bit
            bhi = qa[32:] >> 4        # ql[l+32] 高 4 bit

            q1 = lo | ((ha & 0x3) << 4)
            q2 = hi | (((ha >> 2) & 0x3) << 4)
            q3 = blo | (((ha >> 4) & 0x3) << 4)
            q4 = bhi | (((ha >> 6) & 0x3) << 4)

            p = base + half * 128
            out[p + 0: p + 32] = d * (q1 - 32) * scale[0:32]
            out[p + 32: p + 64] = d * (q2 - 32) * scale[32:64]
            out[p + 64: p + 96] = d * (q3 - 32) * scale[64:96]
            out[p + 96: p + 128] = d * (q4 - 32) * scale[96:128]
    return out


def quant_q8_K(x: np.ndarray) -> np.ndarray:
    """照抄 quantize_row_q8_K_ref：回傳反量化後的 activation（float32）。"""
    n = x.size
    out = np.empty(n, dtype=np.float32)
    for i in range(n // QK_K):
        blk = x[i * QK_K:(i + 1) * QK_K]
        ax = np.abs(blk)
        amax = ax.max()
        if amax == 0:
            out[i * QK_K:(i + 1) * QK_K] = 0
            continue
        # ggml 取「第一個達到 amax 的元素」作為 max（含正負號）
        idx = int(np.argmax(ax))
        maxv = float(blk[idx])
        iscale = -127.0 / maxv
        q = np.clip(np.rint(iscale * blk), -127, 127).astype(np.int8)
        out[i * QK_K:(i + 1) * QK_K] = q.astype(np.float32) / iscale
    return out


def read_expert(fh, off: int, n: int, typ: str) -> np.ndarray:
    fh.seek(off)
    raw = fh.read(n)
    if typ == "Q4_K":
        return dequant_q4_K(raw, n)
    if typ == "Q5_K":
        return dequant_q5_K(raw, n)
    if typ == "Q6_K":
        return dequant_q6_K(raw, n)
    raise SystemExit(f"selftest 暫不支援 {typ}")


def main() -> None:
    ap = argparse.ArgumentParser(description="SDQ MoE kernel 的 NumPy 對照")
    ap.add_argument("--model", required=True)
    ap.add_argument("--index", required=True, help="tools/gguf_index.py 產生的 JSON")
    ap.add_argument("--bin", required=True, help="epq-chat --selftest 的輸出")
    ap.add_argument("--layer", type=int, required=True)
    ap.add_argument("--expert", type=int, required=True)
    ap.add_argument("--n-tok", type=int, default=1, help="幾個 token 共用同一個 expert")
    ap.add_argument("--tol", type=float, default=2e-3, help="相對誤差上限")
    args = ap.parse_args()

    idx = json.loads(Path(args.index).read_text())
    layer = idx["per_layer"][args.layer]["experts"]
    n_embd = idx["n_embd"]
    n_ff = idx["n_ff_exp"]
    e = args.expert

    n_tok = args.n_tok
    raw = np.fromfile(args.bin, dtype=np.float32)
    assert raw.size == 2 * n_embd * n_tok, f"selftest 輸出大小 {raw.size} 應為 {2 * n_embd * n_tok}"
    cpp_out = raw[:n_embd * n_tok].reshape(n_tok, n_embd)
    x = raw[n_embd * n_tok:].reshape(n_tok, n_embd)

    with open(args.model, "rb") as fh:
        g = layer["ffn_gate_exps"]
        u = layer["ffn_up_exps"]
        d = layer["ffn_down_exps"]
        gate = read_expert(fh, g["file_offset"] + e * g["bytes_per_expert"],
                           n_ff * n_embd, g["type"]).reshape(n_ff, n_embd)
        up = read_expert(fh, u["file_offset"] + e * u["bytes_per_expert"],
                         n_ff * n_embd, u["type"]).reshape(n_ff, n_embd)
        down = read_expert(fh, d["file_offset"] + e * d["bytes_per_expert"],
                           n_embd * n_ff, d["type"]).reshape(n_embd, n_ff)

    # 逐 token：activation → Q8_K → gate/up → SiLU → Q8_K → down
    outs = []
    for t in range(n_tok):
        xq = quant_q8_K(x[t])
        gate_out = gate @ xq
        up_out = up @ xq
        h = (gate_out / (1.0 + np.exp(-gate_out))) * up_out
        hq = quant_q8_K(h)
        outs.append(down @ hq)
    out = np.stack(outs)

    diff = np.abs(out - cpp_out)
    scale = np.maximum(np.abs(out), 1e-3)
    rel = diff / scale
    print(f"layer={args.layer} expert={e} n_embd={n_embd} n_ff={n_ff}")
    print(f"  NumPy out[0:5] = {out[:5]}")
    print(f"  C++   out[0:5] = {cpp_out[:5]}")
    print(f"  max abs diff = {diff.max():.6g}   max rel = {rel.max():.6g}")
    print(f"  out norm = {np.linalg.norm(out):.4f}")
    ok = rel.max() <= args.tol
    print("  結果：", "PASS ✅" if ok else "FAIL ❌")
    raise SystemExit(0 if ok else 1)


if __name__ == "__main__":
    main()