#!/usr/bin/env python3
"""用 epq-chat 的 SDQ_MOE_DUMP 檔，驗證「整層」MoE 的計算（真實路由輸入）。

這是最接近模型實際執行的驗證：不是隨機 activation，而是模型自己算出來的
cur / selected_experts / weights，然後檢查我們的輸出是否等於 NumPy 重算的結果。
"""
import argparse, json, struct, sys
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).parent))
from moe_ref import quant_q8_K, read_expert

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--index", required=True)
    ap.add_argument("--dump", required=True)
    ap.add_argument("--tol", type=float, default=3e-3)
    a = ap.parse_args()

    b = open(a.dump, "rb").read()
    n_embd, n_tok, n_used, n_expert, _, _ = struct.unpack_from("<6i", b, 0)
    off = 24
    cur = np.frombuffer(b, dtype=np.float32, count=n_embd * n_tok, offset=off).reshape(n_tok, n_embd)
    off += 4 * n_embd * n_tok
    ids = np.frombuffer(b, dtype=np.int32, count=n_used * n_tok, offset=off).reshape(n_tok, n_used)
    off += 4 * n_used * n_tok
    wts = np.frombuffer(b, dtype=np.float32, count=n_used * n_tok, offset=off).reshape(n_tok, n_used)
    off += 4 * n_used * n_tok
    got = np.frombuffer(b, dtype=np.float32, count=n_embd * n_tok, offset=off).reshape(n_tok, n_embd)

    idx = json.loads(Path(a.index).read_text())
    il = int(Path(a.dump).name.rsplit(".", 1)[1])
    layer = idx["per_layer"][il]["experts"]
    n_ff = idx["n_ff_exp"]
    out = np.zeros((n_tok, n_embd), dtype=np.float32)

    # Worst-case effect on one output element of a single Q8_K LSB flip in h.
    # ggml accumulates the gate/up dots in float32 while this reference is exact,
    # so the two can disagree by ~1e-6 on h. When that crosses a Q8_K rounding
    # boundary one code moves by 1 and the layer output jumps by far more than
    # float noise. See docs/CORRECTNESS.md: a 1-LSB flip alone produced 7.9e-4.
    lsb = np.zeros(n_ff, dtype=np.float64)

    with open(a.model, "rb") as fh:
        cache = {}
        for t in range(n_tok):
            for k in range(n_used):
                e = int(ids[t, k])
                if e not in cache:
                    g, u, d = layer["ffn_gate_exps"], layer["ffn_up_exps"], layer["ffn_down_exps"]
                    Wg = read_expert(fh, g["file_offset"] + e * g["bytes_per_expert"], n_ff * n_embd, g["type"]).reshape(n_ff, n_embd)
                    Wu = read_expert(fh, u["file_offset"] + e * u["bytes_per_expert"], n_ff * n_embd, u["type"]).reshape(n_ff, n_embd)
                    Wd = read_expert(fh, d["file_offset"] + e * d["bytes_per_expert"], n_embd * n_ff, d["type"]).reshape(n_embd, n_ff)
                    cache[e] = (Wg, Wu, Wd)
                Wg, Wu, Wd = cache[e]
                xq = quant_q8_K(cur[t])
                # float64 for the gate/up dots: this is the reference, so it should
                # be the more accurate of the two.
                go = Wg.astype(np.float64) @ xq.astype(np.float64)
                uo = Wu.astype(np.float64) @ xq.astype(np.float64)
                h = ((go / (1.0 + np.exp(-go))) * uo).astype(np.float32)
                if t == 0 and k == 0:
                    ab = np.abs(h)
                    for bi in range(h.size // 256):
                        blk = ab[bi * 256:(bi + 1) * 256]
                        lsb[bi * 256:(bi + 1) * 256] = blk.max() / 127.0 if blk.max() > 0 else 0.0
                out[t] += (Wd @ quant_q8_K(h)) * wts[t, k]

    diff = np.abs(out - got)
    scale = np.maximum(np.abs(out), 1e-2)
    rel = (diff / scale).max()
    # bound[r] = sum_j |Wd[r,j]| * d_h(block(j)) = worst case for one flipped code
    bound = np.zeros((n_embd, n_ff), dtype=np.float64)
    for t in range(n_tok):
        for k in range(n_used):
            Wd = cache[int(ids[t, k])][2]
            bound += np.abs(Wd.astype(np.float64)) * lsb[None, :]
    print(f"layer {il}  n_tok={n_tok} n_used={n_used} ids[0]={ids[0].tolist()}")
    print(f"  NumPy out[0,:5]={out[0,:5]}")
    print(f"  SDQ   out[0,:5]={got[0,:5]}")
    print(f"  max abs diff={diff.max():.6g}  max rel={rel:.6g}  1-LSB-flip bound={bound.max():.6g}")
    ok = rel <= a.tol
    verdict = "PASS" if ok else ("PASS (quantization-boundary)" if diff.max() <= bound.max() else "FAIL")
    print("  結果：", verdict + (" ✅" if ok else (" ✅" if verdict != "FAIL" else " ❌")))
    sys.exit(0 if ok or verdict != "FAIL" else 1)

main()