package tests

import (
	"math"
	"math/rand"
	"testing"
	"unsafe"

	"makarna/pkg/backend/cpu"
	"makarna/pkg/backend/cpu/nn"
	"makarna/pkg/quant"
	"makarna/pkg/tensor"
)

func maxAbsDiff(a, b []float32) float32 {
	max := float32(0)
	for i := range a {
		d := a[i] - b[i]
		if d < 0 {
			d = -d
		}
		if d > max {
			max = d
		}
	}
	return max
}

func mse(a, b []float32) float32 {
	var s float32
	for i := range a {
		d := a[i] - b[i]
		s += d * d
	}
	return s / float32(len(a))
}

func makeTestBlock(seed int64, scale float32) []float32 {
	r := rand.New(rand.NewSource(seed))
	out := make([]float32, 256)
	for i := range out {
		out[i] = (r.Float32()*2 - 1) * scale
	}
	return out
}

func TestHarness_CPU_QuantDequantRoundTrip(t *testing.T) {
	cases := []struct {
		name      string
		seed      int64
		scale     float32
		mseMax    float32
		maxAbsMax float32
	}{
		{name: "small", seed: 1, scale: 0.01, mseMax: 1e-4, maxAbsMax: 0.1},
		{name: "medium", seed: 2, scale: 1.0, mseMax: 5.0, maxAbsMax: 2.0},
		{name: "large", seed: 3, scale: 100.0, mseMax: 5e4, maxAbsMax: 200.0},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			inp := makeTestBlock(tc.seed, tc.scale)

			q2 := quant.QuantizeQ2K(inp)
			q3 := quant.QuantizeQ3K(inp)
			q4 := quant.QuantizeQ4K(inp)
			q6 := quant.QuantizeQ6K(inp)
			q8 := quant.QuantizeQ8K(inp)

			if len(q2) != 84 {
				t.Fatalf("q2k size=%d", len(q2))
			}
			if len(q3) != 110 {
				t.Fatalf("q3k size=%d", len(q3))
			}
			if len(q4) != 144 {
				t.Fatalf("q4k size=%d", len(q4))
			}
			if len(q6) != 210 {
				t.Fatalf("q6k size=%d", len(q6))
			}
			if len(q8) != 292 {
				t.Fatalf("q8k size=%d", len(q8))
			}

			out2 := make([]float32, 256)
			out3 := make([]float32, 256)
			out4 := make([]float32, 256)
			out6 := make([]float32, 256)
			out8 := make([]float32, 256)

			tensor.DequantizeQ2_K((*tensor.BlockQ2_K)(unsafe.Pointer(&q2[0])), out2)
			tensor.DequantizeQ3_K((*tensor.BlockQ3_K)(unsafe.Pointer(&q3[0])), out3)
			tensor.DequantizeQ4_K((*tensor.BlockQ4_K)(unsafe.Pointer(&q4[0])), out4)
			tensor.DequantizeQ6_K((*tensor.BlockQ6_K)(unsafe.Pointer(&q6[0])), out6)
			tensor.DequantizeQ8_K((*tensor.BlockQ8_K)(unsafe.Pointer(&q8[0])), out8)

			outs := []struct {
				name string
				out  []float32
			}{
				{name: "q2k", out: out2},
				{name: "q3k", out: out3},
				{name: "q4k", out: out4},
				{name: "q6k", out: out6},
				{name: "q8k", out: out8},
			}

			for _, o := range outs {
				for i, v := range o.out {
					if math.IsNaN(float64(v)) || math.IsInf(float64(v), 0) {
						t.Fatalf("%s invalid at %d: %f", o.name, i, v)
					}
				}
				m := mse(inp, o.out)
				mx := maxAbsDiff(inp, o.out)
				if m > tc.mseMax {
					t.Fatalf("%s mse=%f > %f", o.name, m, tc.mseMax)
				}
				if mx > tc.maxAbsMax {
					t.Fatalf("%s maxAbs=%f > %f", o.name, mx, tc.maxAbsMax)
				}
			}
		})
	}
}

func TestHarness_CPU_NNOpsSanity(t *testing.T) {
	dim := 32
	seq := 4
	r := rand.New(rand.NewSource(123))

	xData := make([]float32, seq*dim)
	wData := make([]float32, dim)
	for i := range xData {
		xData[i] = r.Float32()*2 - 1
	}
	for i := range wData {
		wData[i] = r.Float32()*2 - 1
	}

	x := cpu.NewTensor(tensor.Shape{seq, dim}, append([]float32(nil), xData...))
	w := cpu.NewTensor(tensor.Shape{dim}, append([]float32(nil), wData...))
	if err := nn.RMSNorm(x, w, 1e-5); err != nil {
		t.Fatalf("rmsnorm: %v", err)
	}
	for i, v := range x.DataFloat32() {
		if math.IsNaN(float64(v)) || math.IsInf(float64(v), 0) {
			t.Fatalf("rmsnorm invalid at %d: %f", i, v)
		}
	}

	headDim := 16
	numHeads := dim / headDim
	pos := make([]int, seq)
	for i := range pos {
		pos[i] = i
	}
	if err := nn.RoPE(x, pos, headDim, 10000); err != nil {
		t.Fatalf("rope: %v", err)
	}
	for i, v := range x.DataFloat32() {
		if math.IsNaN(float64(v)) || math.IsInf(float64(v), 0) {
			t.Fatalf("rope invalid at %d: %f", i, v)
		}
	}

	_ = numHeads

	row := cpu.NewTensor(tensor.Shape{dim}, append([]float32(nil), xData[:dim]...))
	if err := nn.Softmax(row); err != nil {
		t.Fatalf("softmax: %v", err)
	}
	var sum float32
	for i, v := range row.DataFloat32() {
		if v <= 0 || math.IsNaN(float64(v)) || math.IsInf(float64(v), 0) {
			t.Fatalf("softmax invalid at %d: %f", i, v)
		}
		sum += v
	}
	if d := float32(math.Abs(float64(sum - 1))); d > 1e-4 {
		t.Fatalf("softmax sum=%f", sum)
	}
}