#!/usr/bin/env python3
"""Validate the SDQ paged MoE kernel for Ling-3.0-flash (bailingmoe3) with pure NumPy.

The weights are read straight out of the sharded GGUF, dequantized with ggml's
definitions, and the SiLU-gated expert forward pass is recomputed in fp32. The
result is compared against `sdq-chat --selftest IL:IE:SEED:OUT.bin`.

Ling-3.0-flash expert tensors (Q4_K_M):
    ffn_gate_exps [n_embd=2560, n_ff_exp=768, n_expert=512]  q4_K
    ffn_up_exps   [2560, 768, 512]                            q4_K
    ffn_down_exps [768, 2560, 512]                            q4_K or q6_K
"""
import os
import struct
import sys

import numpy as np

QK_K = 256
MODEL = os.environ.get('MODEL', '/root/app/ling/tiny/models/Ling-3.0-tiny-Q4_K_M.gguf')
SHARDS = [MODEL.replace('00001-of', f'{i:05d}-of') for i in (1, 2)]

TS = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8}
BLK = {12: (256, 144), 13: (256, 176), 14: (256, 210), 7: (32, 24), 8: (32, 34),
       2: (32, 18), 3: (32, 20), 6: (32, 22), 0: (1, 4), 1: (1, 2)}
FMT = {0: '<B', 1: '<b', 2: '<H', 3: '<h', 4: '<I', 5: '<i', 6: '<f', 7: '<B',
       10: '<Q', 11: '<q', 12: '<d'}
NAME = {12: 'q4_K', 13: 'q5_K', 14: 'q6_K', 7: 'q5_1', 8: 'q8_0', 6: 'q5_0', 2: 'q4_0'}


def rstr(f):
    n = struct.unpack('<Q', f.read(8))[0]
    return f.read(n)


def read_val(f, t):
    if t == 8:
        return rstr(f).decode('utf-8', 'replace')
    if t == 9:
        et = struct.unpack('<I', f.read(4))[0]
        n = struct.unpack('<Q', f.read(8))[0]
        if et == 8:
            return [rstr(f).decode('utf-8', 'replace') for _ in range(n)]
        buf = f.read(n * TS[et])
        return np.frombuffer(buf, dtype=np.dtype(FMT[et])).tolist()
    buf = f.read(TS[t])
    return np.frombuffer(buf, dtype=np.dtype(FMT[t]))[0].item()


def parse_headers():
    kv = {}
    tensors = {}
    for si, path in enumerate(SHARDS):
        with open(path, 'rb') as f:
            assert f.read(4) == b'GGUF'
            f.read(4)
            n_tensors = struct.unpack('<Q', f.read(8))[0]
            n_kv = struct.unpack('<Q', f.read(8))[0]
            for _ in range(n_kv):
                k = rstr(f).decode()
                kv[k] = read_val(f, struct.unpack('<I', f.read(4))[0])
            align = kv.get('general.alignment', 32)
            pending = []
            for _ in range(n_tensors):
                name = rstr(f).decode()
                nd = struct.unpack('<I', f.read(4))[0]
                dims = struct.unpack(f'<{nd}Q', f.read(8 * nd))
                tt = struct.unpack('<I', f.read(4))[0]
                off = struct.unpack('<Q', f.read(8))[0]
                pending.append((name, dims, tt, off))
            # data_start is only known after the whole tensor table has been read
            data_start = (f.tell() + align - 1) // align * align
            for name, dims, tt, off in pending:
                tensors[name] = (si, data_start + off, dims, tt)
    return kv, tensors


def fp16(x):
    return np.frombuffer(x, dtype=np.float16)[0].astype(np.float32)


def dequant_q4_K(raw, n):
    nb = n // QK_K
    out = np.empty(n, dtype=np.float32)
    for i in range(nb):
        b = raw[i * 144:(i + 1) * 144]
        d = fp16(b[0:2])
        dmin = fp16(b[2:4])
        sc = np.frombuffer(b, dtype=np.uint8, count=12, offset=4)
        qs = np.frombuffer(b, dtype=np.uint8, count=128, offset=16)
        base = i * QK_K
        is_ = 0
        for j in range(0, QK_K, 64):
            if is_ < 4:
                sc1, m1 = int(sc[is_] & 63), int(sc[is_ + 4] & 63)
                sc2, m2 = int(sc[is_ + 1] & 63), int(sc[is_ + 5] & 63)
            else:
                sc1 = int((sc[is_ + 4] & 0xF) | ((sc[is_ - 4] >> 6) << 4))
                m1 = int((sc[is_ + 4] >> 4) | ((sc[is_] >> 6) << 4))
                sc2 = int((sc[is_ + 5] & 0xF) | ((sc[is_ - 3] >> 6) << 4))
                m2 = int((sc[is_ + 5] >> 4) | ((sc[is_ + 1] >> 6) << 4))
            q = qs[j // 2: j // 2 + 32]
            out[base + j: base + j + 32] = d * sc1 * (q & 0xF) - dmin * m1
            out[base + j + 32: base + j + 64] = d * sc2 * (q >> 4) - dmin * m2
            is_ += 2
    return out


def dequant_q6_K(raw, n):
    nb = n // QK_K
    out = np.empty(n, dtype=np.float32)
    for i in range(nb):
        b = raw[i * 210:(i + 1) * 210]
        d = fp16(b[208:210])
        ql = np.frombuffer(b, dtype=np.uint8, count=128, offset=0)
        qh = np.frombuffer(b, dtype=np.uint8, count=64, offset=128)
        sc = np.frombuffer(b, dtype=np.int8, count=16, offset=192)
        base = i * QK_K
        for g in range(2):   # two 128-element groups per block
            for l in range(32):
                is_ = l // 16
                q1 = ((ql[g * 64 + l] & 0xF) | (((qh[g * 32 + l] >> 0) & 3) << 4)).astype(np.int16) - 32
                q2 = ((ql[g * 64 + l + 32] & 0xF) | (((qh[g * 32 + l] >> 2) & 3) << 4)).astype(np.int16) - 32
                q3 = ((ql[g * 64 + l] >> 4) | (((qh[g * 32 + l] >> 4) & 3) << 4)).astype(np.int16) - 32
                q4 = ((ql[g * 64 + l + 32] >> 4) | (((qh[g * 32 + l] >> 6) & 3) << 4)).astype(np.int16) - 32
                off0 = base + g * 128
                out[off0 + l] = d * sc[g * 8 + is_] * q1
                out[off0 + l + 32] = d * sc[g * 8 + is_ + 2] * q2
                out[off0 + l + 64] = d * sc[g * 8 + is_ + 4] * q3
                out[off0 + l + 96] = d * sc[g * 8 + is_ + 6] * q4
    return out


def dequant(raw, n, tt):
    if tt == 12:
        return dequant_q4_K(raw, n)
    if tt == 14:
        return dequant_q6_K(raw, n)
    raise ValueError(f'unsupported type {tt} ({NAME.get(tt, tt)})')


def row_bytes(n, tt):
    bs, ts = BLK[tt]
    return n // bs * ts


def read_expert(tensors, il, ie):
    def seg(name):
        si, off, dims, tt = tensors[f'blk.{il}.{name}']
        nrow = row_bytes(dims[0], tt)
        n = dims[1] * nrow
        with open(SHARDS[si], 'rb') as f:
            f.seek(off + ie * n)
            return f.read(n), nrow, tt, dims[0]
    g = seg('ffn_gate_exps.weight')
    u = seg('ffn_up_exps.weight')
    d = seg('ffn_down_exps.weight')
    return g, u, d


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


def main():
    cases = []
    for path in sys.argv[1:]:
        # st_l<IL>e<IE>.bin
        base = os.path.basename(path)
        body = base[len('st_l'):-len('.bin')]
        il, ie = body.split('e')
        cases.append((int(il), int(ie), path))
    if not cases:
        print('usage: verify_ling.py st_l<IL>e<IE>.bin ...')
        return 1

    kv, tensors = parse_headers()
    n_embd = int(kv['bailingmoe3.embedding_length'])
    n_ff = int(kv['bailingmoe3.expert_feed_forward_length'])
    print(f"arch=bailingmoe3 n_expert={int(kv['bailingmoe3.expert_count'])} "
          f"n_expert_used={int(kv['bailingmoe3.expert_used_count'])} "
          f"n_embd={n_embd} n_ff_exp={n_ff}")

    all_ok = True
    for il, ie, path in cases:
        raw = open(path, 'rb').read()
        n_tok = len(raw) // 4 // (2 * n_embd)
        got = np.frombuffer(raw[:n_tok * n_embd * 4], dtype=np.float32).reshape(n_tok, n_embd)
        x = np.frombuffer(raw[n_tok * n_embd * 4:], dtype=np.float32).reshape(n_tok, n_embd)

        g_seg, u_seg, d_seg = read_expert(tensors, il, ie)
        Wg = load_matrix(g_seg, n_ff)
        Wu = load_matrix(u_seg, n_ff)
        Wd = load_matrix(d_seg, n_embd)

        rels, corrs = [], []
        for t in range(n_tok):
            xt = x[t]
            g = Wg @ xt
            u = Wu @ xt
            h = (g / (1.0 + np.exp(-g))) * u
            ref = Wd @ h
            denom = max(float(np.abs(ref).max()), 1e-9)
            rels.append(float(np.abs(got[t] - ref).max()) / denom)
            corrs.append(float(np.corrcoef(got[t], ref)[0, 1]))
        rel = max(rels)
        corr = min(corrs)
        ok = rel < 0.05 and corr > 0.99
        all_ok &= ok
        print(f'L{il} E{ie}: gate={NAME[g_seg[2]]} up={NAME[u_seg[2]]} down={NAME[d_seg[2]]}'
              f' | max_rel_err={rel:.4f} corr={corr:.6f} | {"PASS" if ok else "FAIL"}')

    print('\nALL PASS' if all_ok else '\nSOME FAILED')
    return 0 if all_ok else 1


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