#!/usr/bin/env python3
"""Recompute a whole paged MoE layer in NumPy and compare with sdq-chat's output.

`sdq-chat` with SDQ_MOE_DUMP=<prefix> SDQ_MOE_DUMP_LAYER=<il> writes, from inside
the custom MoE op:
    int32 meta[6] = { n_embd, n_tok, n_expert_used, n_expert, 1, 2 }
    float  cur   [n_tok][n_embd]        (FFN input)
    int32  ids   [n_tok][n_expert_used] (routed experts)
    float  wts   [n_tok][n_expert_used] (routed weights)
    float  out   [n_tok][n_embd]        (SDQ MoE output)

This script reads the same expert weights straight out of the GGUF, runs the
SiLU-gated expert forward in fp32 and diffs the result, so the whole
pager + custom kernel + weight-scaling path is checked at once.
"""
import importlib.util
import struct
import sys

import numpy as np

spec = importlib.util.spec_from_file_location('vl', str(__import__('pathlib').Path(__file__).with_name('verify_ling.py')))
vl = importlib.util.module_from_spec(spec)
spec.loader.exec_module(vl)

QK_K = 256


def dequant_q4_K_vec(raw, n_elem):
    """raw: uint8 [..., 144] -> float32 [..., 256]"""
    raw = raw.reshape(-1, 144)
    nb = raw.shape[0]
    out = np.empty((nb, QK_K), dtype=np.float32)
    d = raw[:, 0:2].copy().view(np.float16).astype(np.float32)[:, 0]
    dmin = raw[:, 2:4].copy().view(np.float16).astype(np.float32)[:, 0]
    b = np.ascontiguousarray(raw[:, 4:16])           # 12 packed 6-bit scale bytes
    sc = np.empty((nb, 8), dtype=np.int32)
    mn = np.empty((nb, 8), dtype=np.int32)
    for j in range(4):
        sc[:, j] = b[:, j] & 63
        mn[:, j] = b[:, j + 4] & 63
    for j in range(4, 8):
        sc[:, j] = (b[:, j + 4] & 0xF) | ((b[:, j - 4] >> 6) << 4)
        mn[:, j] = (b[:, j + 4] >> 4) | ((b[:, j] >> 6) << 4)
    qs = raw[:, 16:144].reshape(nb, 4, 32).astype(np.int32)
    q = np.empty((nb, QK_K), dtype=np.int32)
    for j in range(4):
        q[:, j * 64:j * 64 + 32] = qs[:, j, :] & 0xF
        q[:, j * 64 + 32:j * 64 + 64] = qs[:, j, :] >> 4
    gidx = np.repeat(np.arange(8), 32)               # 32-value group -> (sc, mn) index
    out[:] = d[:, None] * sc[:, gidx] * q - dmin[:, None] * mn[:, gidx]
    return out


def dequant_q6_K_vec(raw, n_elem):
    raw = raw.reshape(-1, 210)
    nb = raw.shape[0]
    out = np.empty((nb, QK_K), dtype=np.float32)
    d = raw[:, 208:210].copy().view(np.float16).astype(np.float32)[:, 0]
    ql = raw[:, 0:128].reshape(nb, 4, 32).astype(np.int32)      # 4 x 32 low/high nibbles
    qh = raw[:, 128:192].reshape(nb, 2, 32).astype(np.int32)
    sc = raw[:, 192:208].copy().view(np.int8).astype(np.int32)
    q = np.empty((nb, QK_K), dtype=np.int32)
    sidx = np.empty(QK_K, dtype=np.int32)
    for g in range(2):
        for l in range(32):
            is_ = l // 16
            o = g * 128
            a = ql[:, 2 * g, l]        # ql[g*64 + l]
            b_ = ql[:, 2 * g + 1, l]      # ql[g*64 + l + 32]
            h = qh[:, g, l]
            q[:, o + l] = ((a & 0xF) | ((h & 3) << 4)) - 32
            q[:, o + l + 32] = ((b_ & 0xF) | (((h >> 2) & 3) << 4)) - 32
            q[:, o + l + 64] = ((a >> 4) | (((h >> 4) & 3) << 4)) - 32
            q[:, o + l + 96] = ((b_ >> 4) | (((h >> 6) & 3) << 4)) - 32
            sidx[o + l] = g * 8 + is_
            sidx[o + l + 32] = g * 8 + is_ + 2
            sidx[o + l + 64] = g * 8 + is_ + 4
            sidx[o + l + 96] = g * 8 + is_ + 6
    out[:] = d[:, None] * sc[:, sidx] * q
    return out


def dequant_mat(raw, n_rows, ne0, tt):
    nrow = vl.row_bytes(ne0, tt)
    nblk = ne0 // QK_K
    arr = np.frombuffer(raw[:n_rows * nrow], dtype=np.uint8).reshape(n_rows, nblk, nrow // nblk)
    if tt == 12:
        return dequant_q4_K_vec(arr, ne0).reshape(n_rows, ne0)
    if tt == 14:
        return dequant_q6_K_vec(arr, ne0).reshape(n_rows, ne0)
    raise ValueError(tt)


def main():
    dump = sys.argv[1]
    il = int(sys.argv[2])
    kv, ts = vl.parse_headers()
    n_embd = int(kv['bailingmoe3.embedding_length'])
    n_ff = int(kv['bailingmoe3.expert_feed_forward_length'])
    n_expert = int(kv['bailingmoe3.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)
        h = ((Wg @ x) / (1.0 + np.exp(-(Wg @ x)))) * (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())