"""Test Trill quantization: ternary, 4-bit, 8-bit, round-trip accuracy."""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import numpy as np
from trill_llm.model.quantization import TrillQuantizer, bits_per_weight, compression_ratio


def test_ternary_quantization():
    """Test ternary quantization round-trip."""
    quantizer = TrillQuantizer(format="ternary")
    weights = np.random.randn(64, 32).astype(np.float32) * 0.1

    packed = quantizer.quantize(weights)
    dequant = quantizer.dequantize(packed)

    # Ternary loses precision but preserves sign and approximate magnitude
    correlation = np.corrcoef(weights.flatten(), dequant.flatten())[0, 1]
    assert correlation > 0.5, f"Ternary correlation too low: {correlation}"
    print(f"  Ternary correlation: {correlation:.3f}")
    print(f"  BPW: {quantizer.bpw:.3f}, Compression: {compression_ratio('ternary'):.1f}x")


def test_4bit_quantization():
    """Test 4-bit quantization round-trip."""
    quantizer = TrillQuantizer(format="q4_k_m")
    weights = np.random.randn(128, 64).astype(np.float32) * 0.1

    packed = quantizer.quantize(weights)
    dequant = quantizer.dequantize(packed)

    # 4-bit should be closer to original
    max_error = np.max(np.abs(weights - dequant))
    rel_error = max_error / np.max(np.abs(weights))
    assert rel_error < 0.2, f"4-bit relative error too high: {rel_error}"
    print(f"  4-bit max relative error: {rel_error:.4f}")
    print(f"  BPW: {quantizer.bpw:.1f}, Compression: {compression_ratio('q4_k_m'):.1f}x")


def test_8bit_quantization():
    """Test 8-bit quantization round-trip."""
    quantizer = TrillQuantizer(format="q8_0")
    weights = np.random.randn(256, 128).astype(np.float32) * 0.1

    packed = quantizer.quantize(weights)
    dequant = quantizer.dequantize(packed)

    # 8-bit should be very close
    max_error = np.max(np.abs(weights - dequant))
    rel_error = max_error / np.max(np.abs(weights))
    assert rel_error < 0.02, f"8-bit relative error too high: {rel_error}"
    print(f"  8-bit max relative error: {rel_error:.5f}")


def test_fp16_passthrough():
    """Test fp16 passthrough (no quantization)."""
    quantizer = TrillQuantizer(format="fp16")
    weights = np.random.randn(64, 32).astype(np.float32)

    packed = quantizer.quantize(weights)
    dequant = quantizer.dequantize(packed)

    # fp16 should be nearly identical
    max_error = np.max(np.abs(weights - dequant))
    assert max_error < 0.01, f"fp16 error too high: {max_error}"
    print(f"  fp16 max error: {max_error:.6f}")


def test_bpw_table():
    """Test bits per weight table."""
    assert bits_per_weight("ternary") > 1.5
    assert bits_per_weight("q4_k_m") == 4.0
    assert bits_per_weight("q8_0") == 8.0
    assert bits_per_weight("fp16") == 16.0
    assert compression_ratio("ternary") > 10.0
    print(f"  Ternary BPW: {bits_per_weight('ternary'):.3f}")
    print(f"  Ternary compression: {compression_ratio('ternary'):.1f}x")


if __name__ == "__main__":
    print("Running quantization tests...")
    test_ternary_quantization()
    print("  ✓ test_ternary_quantization")
    test_4bit_quantization()
    print("  ✓ test_4bit_quantization")
    test_8bit_quantization()
    print("  ✓ test_8bit_quantization")
    test_fp16_passthrough()
    print("  ✓ test_fp16_passthrough")
    test_bpw_table()
    print("  ✓ test_bpw_table")
    print("\nAll quantization tests passed!")