#!/usr/bin/env python3
"""用 NumPy 重算一整層分頁 MoE，並與 sdq-chat 的輸出比對。

`sdq-chat` 在 SDQ_MOE_DUMP=<prefix> SDQ_MOE_DUMP_LAYER=<il> 時，會從自訂 MoE op
內寫出：
    int32 meta[6] = { n_embd, n_tok, n_expert_used, n_expert, 1, 2 }
    float  cur   [n_tok][n_embd]        (FFN 輸入)
    int32  ids   [n_tok][n_expert_used] (被路由到的 expert)
    float  wts   [n_tok][n_expert_used] (路由權重)
    float  out   [n_tok][n_embd]        (SDQ MoE 輸出)

本腳本從 GGUF 讀出同一批 expert 權重，以 fp32 重算 SiLU-gated expert 前向並比對，
一次檢查 pager + 自訂 kernel + 權重縮放整條路徑。
"""
import importlib.util
import os
import struct
import sys

import numpy as np

_HERE = os.path.dirname(os.path.abspath(__file__))
spec = importlib.util.spec_from_file_location('vl', os.path.join(_HERE, 'verify_epq.py'))
vl = importlib.util.module_from_spec(spec)
spec.loader.exec_module(vl)

_ARCH = 'openpangu-v2'


def dequant_mat(raw, n_rows, ne0, tt):
    nrow = vl.row_bytes(ne0, tt)
    return np.stack([vl.dequant(raw[r * nrow:(r + 1) * nrow], ne0, tt) for r in range(n_rows)])


def main():
    dump = sys.argv[1]
    il = int(sys.argv[2])
    kv, ts = vl.parse_headers()
    n_embd = int(kv[f'{_ARCH}.embedding_length'])
    n_ff = int(kv[f'{_ARCH}.expert_feed_forward_length'])
    n_expert = int(kv[f'{_ARCH}.expert_count'])

    raw = open(dump, 'rb').read()
    meta = struct.unpack('<6i', raw[:24])
    n_tok, n_used = meta[1], meta[2]
    off = 24
    cur = np.frombuffer(raw, dtype=np.float32, count=n_tok * n_embd, offset=off).reshape(n_tok, n_embd)
    off += 4 * n_tok * n_embd
    ids = np.frombuffer(raw, dtype=np.int32, count=n_tok * n_used, offset=off).reshape(n_tok, n_used)
    off += 4 * n_tok * n_used
    wts = np.frombuffer(raw, dtype=np.float32, count=n_tok * n_used, offset=off).reshape(n_tok, n_used)
    off += 4 * n_tok * n_used
    got = np.frombuffer(raw, dtype=np.float32, count=n_tok * n_embd, offset=off).reshape(n_tok, n_embd)
    print(f'layer {il}: n_tok={n_tok} n_expert_used={n_used} n_expert={n_expert} '
          f'n_embd={n_embd} n_ff={n_ff}')

    def seg(nm, ie):
        si, off_t, dims, tt = ts[f'blk.{il}.{nm}']
        nrow = vl.row_bytes(dims[0], tt)
        nb = nrow * dims[1]
        with open(vl.SHARDS[si], 'rb') as f:
            f.seek(off_t + ie * nb)
            return f.read(nb), nrow, tt, dims[0]

    ref = np.zeros((n_tok, n_embd), dtype=np.float32)
    experts = sorted(set(int(e) for e in ids.reshape(-1)))
    print(f'distinct experts: {len(experts)}')
    for ie in experts:
        g_raw, g_nrow, g_t, g_ne0 = seg('ffn_gate_exps.weight', ie)
        u_raw, _, u_t, u_ne0 = seg('ffn_up_exps.weight', ie)
        d_raw, _, d_t, d_ne0 = seg('ffn_down_exps.weight', ie)
        Wg = dequant_mat(g_raw, n_ff, g_ne0, g_t)
        Wu = dequant_mat(u_raw, n_ff, u_ne0, u_t)
        Wd = dequant_mat(d_raw, n_embd, d_ne0, d_t)
        rows = [t for t in range(n_tok) if ie in ids[t]]
        x = cur[rows].T                       # [n_embd, T]
        w = np.array([wts[t][list(ids[t]).index(ie)] for t in rows], dtype=np.float32)
        g = Wg @ x
        h = (g / (1.0 + np.exp(-g))) * (Wu @ x)
        y = Wd @ h                             # [n_embd, T]
        ref[rows] += (y * w).T
        del Wg, Wu, Wd

    denom = max(float(np.abs(ref).max()), 1e-9)
    diff = np.abs(got - ref)
    corr = float(np.corrcoef(got.reshape(-1), ref.reshape(-1))[0, 1])
    print(f'max|got-ref| = {diff.max():.6g}   max|ref| = {denom:.6g}   '
          f'max_rel = {diff.max() / denom:.6g}   corr = {corr:.8f}')
    ok = corr > 0.99 and diff.max() / denom < 0.05
    print('PASS' if ok else 'FAIL')
    return 0 if ok else 1


if __name__ == '__main__':
    sys.exit(main())