package tests

import (
	"math"
	"testing"
	
	"makarna/pkg/tensor"
)

func float32NearlyEqual(a, b, epsilon float32) bool {
	return float32(math.Abs(float64(a-b))) <= epsilon
}

func TestDequantizeQ4_K(t *testing.T) {
	// Create a BlockQ4_K with known parameters
	// D = 1.0 (FP16: 0x3C00)
	// DMin = 0.0 (FP16: 0x0000)
	// Scales: all 1s (ls=1, lm=0)
	// QS: all 8s (0x88) -> q=8
	// Expected: w = 8 * (1.0 * 1.0) - 0 = 8.0
	
	// Create scales array
	// j < 4: scales[j] = ls, scales[j+4] = lm
	// ls = 1, lm = 0
	// scales[0..3] = 1, scales[4..7] = 0
	// j >= 4 (packed):
	// scales[j+4] (8..11) = (ls & 0xF) | ((lm & 0xF) << 4) = 1 | 0 = 1
	// scales[j-4] (0..3) |= ((ls >> 4) << 6) = 0
	// scales[j]   (4..7) |= ((lm >> 4) << 6) = 0
	
	scales := [12]uint8{
		1, 1, 1, 1, // ls (0..3)
		0, 0, 0, 0, // lm (0..3) initially
		1, 1, 1, 1, // ls/lm (4..7) packed: (1 | 0<<4) = 1
	}
	
	qs := [128]uint8{}
	for i := range qs {
		qs[i] = 0x88 // q=8 for both low and high nibbles
	}
	
	block := &tensor.BlockQ4_K{
		D:      0x3C00, // 1.0
		DMin:   0x0000, // 0.0
		Scales: scales,
		QS:     qs,
	}
	
	out := make([]float32, 256)
	tensor.DequantizeQ4_K(block, out)
	
	for i, v := range out {
		if !float32NearlyEqual(v, 8.0, 1e-4) {
			t.Errorf("DequantizeQ4_K[%d] = %f, expected 8.0", i, v)
		}
	}
}

func TestFP16ToFP32(t *testing.T) {
	tests := []struct {
		in  uint16
		out float32
	}{
		{0x3C00, 1.0},
		{0x0000, 0.0},
		{0xC000, -2.0},
		{0x7C00, float32(math.Inf(1))},
	}
	
	for _, tc := range tests {
		if res := tensor.FP16ToFP32(tc.in); res != tc.out {
			t.Errorf("FP16ToFP32(0x%X) = %f, expected %f", tc.in, res, tc.out)
		}
	}
}

func TestDequantizeQ8_K(t *testing.T) {
	// Create a BlockQ8_K with known parameters
	// D = 2.0, QS = [1, 2, 3, ..., 127, -127, -126, ..., -1, 0, ...]
	// Expected output: D * qs[i]
	
	d := float32(0.5)
	
	var qs [256]int8
	for i := 0; i < 256; i++ {
		qs[i] = int8((i % 255) - 127) // Range: -127 to 127
	}
	
	var bsums [16]int16 // Not used in dequantization
	
	block := &tensor.BlockQ8_K{
		D:     d,
		QS:    qs,
		BSums: bsums,
	}
	
	out := make([]float32, 256)
	tensor.DequantizeQ8_K(block, out)
	
	for i := 0; i < 256; i++ {
		expected := d * float32(qs[i])
		if !float32NearlyEqual(out[i], expected, 1e-6) {
			t.Errorf("DequantizeQ8_K[%d] = %f, expected %f", i, out[i], expected)
		}
	}
}