"""GGUF export (T2.7): tensor types of a student export, with the dense matrices as trained (bf16) or Q8_0."""
import sys
from pathlib import Path

import numpy as np
import pytest
import torch

from lowbit.export import gguf as gx

GGUF_PY = Path("/path/to/llama.cpp/gguf-py")
pytestmark = pytest.mark.skipif(not GGUF_PY.exists(), reason="llama.cpp's gguf-py is not on this machine")


def test_dense_as_q8_0_quantises_matrices_only():
    sys.path.insert(0, str(GGUF_PY))
    import gguf
    T = gguf.GGMLQuantizationType
    torch.manual_seed(0)
    dense = {"model.layers.0.self_attn.q_proj.weight": torch.randn(64, 128, dtype=torch.bfloat16),
             "model.layers.0.input_layernorm.weight": torch.ones(128),
             "model.layers.0.mlp.gate.weight": torch.randn(8, 128),
             "model.embed_tokens.weight": torch.randn(32, 128, dtype=torch.bfloat16)}
    plain = gx.student_tensors(dense, {}, gguf)
    q8 = gx.student_tensors(dense, {}, gguf, dense_as="q8_0")
    assert plain["blk.0.attn_q.weight"][0] == T.BF16 and q8["blk.0.attn_q.weight"][0] == T.Q8_0
    assert q8["token_embd.weight"][0] == T.Q8_0
    assert q8["blk.0.attn_norm.weight"][0] == T.F32 and q8["blk.0.ffn_gate_inp.weight"][0] == T.F32  # norms, router
    back = gguf.quants.dequantize(q8["blk.0.attn_q.weight"][1], T.Q8_0)
    ref = dense["model.layers.0.self_attn.q_proj.weight"].float().numpy()
    assert back.shape == ref.shape and np.abs(back - ref).max() < 0.02 * np.abs(ref).max()