#!/usr/bin/env python3
# Parity check between ternaria's Rust DT3 packer and llama.cpp's C
# dequantization: blocks packed by pack_dt3 must dequantize (through
# dequantize_row_dt3, exposed by `test-dt3 --dequant`) to exactly
# d1*t1 + d2*t2 computed independently in numpy from the same trits and
# fp16-rounded scales.
#
# Must run inside the ternaria environment, e.g.:
#   cd /path/to/ternaria && uv run python /path/to/llama.cpp/tests/test-dt3-rust-parity.py \
#       /path/to/llama.cpp/build/bin/test-dt3

import subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np

from ternaria._core import pack_dt3

QK_DT3 = 128
BLOCK_BYTES = 56


def main() -> int:
    if len(sys.argv) != 2:
        print(f"usage: {sys.argv[0]} /path/to/test-dt3", file=sys.stderr)
        return 1
    test_bin = Path(sys.argv[1])

    rng = np.random.default_rng(20260810)

    rows, cols = 8, 1024
    t1 = rng.integers(-1, 2, size=(rows, cols)).astype(np.int8)
    t2 = rng.integers(-1, 2, size=(rows, cols)).astype(np.int8)
    s1 = rng.normal(size=(rows, cols // QK_DT3)).astype(np.float32)
    s2 = (0.25 * rng.normal(size=(rows, cols // QK_DT3))).astype(np.float32)

    # edge cases: an all-zero block, a block with negative scales, and a
    # block that is non-trivial only in the qh region (elements 120..127)
    t1[0, :QK_DT3] = 0
    t2[0, :QK_DT3] = 0
    s1[0, 0] = 0.0
    s2[0, 0] = 0.0
    s1[0, 1] = -abs(s1[0, 1])
    s2[0, 1] = -abs(s2[0, 1])
    t1[1, :QK_DT3] = 0
    t2[1, :QK_DT3] = 0
    t1[1, 120:128] = [-1, 1, 0, -1, 1, -1, 0, 1]
    t2[1, 120:128] = [1, -1, 1, 0, 0, 1, -1, -1]

    raw = np.asarray(pack_dt3(t1, s1, t2, s2), dtype=np.uint8)
    assert raw.size == rows * (cols // QK_DT3) * BLOCK_BYTES, raw.size

    with tempfile.TemporaryDirectory() as tmp:
        raw_path = Path(tmp) / "dt3.bin"
        out_path = Path(tmp) / "out.f32"
        raw_path.write_bytes(raw.tobytes())
        subprocess.run([str(test_bin), "--dequant", str(raw_path), str(out_path)], check=True)
        got = np.fromfile(out_path, dtype=np.float32).reshape(rows, cols)

    # what the packed bytes mean: fp16-rounded scales times the trits
    d1 = s1.astype(np.float16).astype(np.float32).repeat(QK_DT3, axis=1)
    d2 = s2.astype(np.float16).astype(np.float32).repeat(QK_DT3, axis=1)
    expected = d1 * t1.astype(np.float32) + d2 * t2.astype(np.float32)

    if not np.array_equal(got, expected):
        bad = np.nonzero(got != expected)
        print(f"FAILED: {len(bad[0])} of {got.size} elements differ", file=sys.stderr)
        r, c = bad[0][0], bad[1][0]
        print(f"first mismatch at ({r}, {c}): got {got[r, c]}, expected {expected[r, c]}", file=sys.stderr)
        return 1

    print(f"ok: {got.size} weights bit-exact between Rust pack_dt3 and C dequantize_row_dt3")
    return 0


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