#!/usr/bin/env python3
"""Byte-level cross-check of the NumPy dequantizers against ggml's own.

`ggml_deq <type> <gguf> <offset> <n>` (built from sdq-ling/ggml_deq.cpp) dumps one
dequantized row using ggml's dequantize_row_*; this script compares it with the
NumPy implementation used by verify_ling.py / verify_moe_layer.py.
"""
import importlib.util
import subprocess
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)

GGML_DEQ = str(__import__('pathlib').Path(__file__).resolve().parents[1] / 'llama.cpp/build/bin/ggml_deq')


def main():
    kv, ts = vl.parse_headers()
    cases = [(2, 'ffn_gate_exps.weight', 12), (2, 'ffn_up_exps.weight', 12),
             (2, 'ffn_down_exps.weight', 14), (5, 'ffn_gate_exps.weight', 12),
             (5, 'ffn_down_exps.weight', 14), (20, 'ffn_gate_exps.weight', 12),
             (23, 'ffn_down_exps.weight', 14)]
    all_ok = True
    for il, name, tt in cases:
        si, off, dims, t = ts[f'blk.{il}.{name}']
        assert t == tt, (name, t, tt)
        nrow = vl.row_bytes(dims[0], tt)
        ne0 = dims[0]
        for ie in (0, 17, 127):
            eo = off + ie * nrow * dims[1]
            out = subprocess.run([GGML_DEQ, str(tt), vl.SHARDS[si], str(eo), str(ne0)],
                                 capture_output=True, check=True)
            ref = np.frombuffer(out.stdout, dtype=np.float32)
            with open(vl.SHARDS[si], 'rb') as f:
                f.seek(eo)
                mine = vl.dequant(f.read(nrow), ne0, tt)
            d = float(np.abs(ref - mine).max())
            ok = d == 0.0
            all_ok &= ok
            print(f'blk.{il}.{name} E{ie} type={vl.NAME.get(tt)} n={ne0}: '
                  f'max|ggml - numpy| = {d:g}  {"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())