package matmul

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

	"makarna/pkg/backend/cpu"
)

func TestGemvFloat32RangeMatchesScalar(t *testing.T) {
	rng := rand.New(rand.NewSource(1))

	cases := []struct {
		K      int
		N      int
		startN int
		endN   int
	}{
		{K: 3, N: 13, startN: 0, endN: 13},
		{K: 33, N: 17, startN: 0, endN: 17},
		{K: 33, N: 17, startN: 5, endN: 13},
		{K: 32, N: 16, startN: 0, endN: 16},
	}

	for _, tc := range cases {
		a := make([]float32, tc.K)
		for i := range a {
			a[i] = rng.Float32()*2 - 1
		}
		b := make([]float32, tc.N*tc.K)
		for i := range b {
			b[i] = rng.Float32()*2 - 1
		}

		want := make([]float32, tc.N)
		for n := tc.startN; n < tc.endN; n++ {
			var sum float32
			rowOff := n * tc.K
			for k := 0; k < tc.K; k++ {
				sum += a[k] * b[rowOff+k]
			}
			want[n] = sum
		}

		// Sanity: cpu.DotFloat32 should broadly match the scalar sum.
		for n := tc.startN; n < tc.endN; n++ {
			rowOff := n * tc.K
			gotCPU := cpu.DotFloat32(a, b[rowOff:rowOff+tc.K])
			diff := math.Abs(float64(gotCPU - want[n]))
			if diff > 1e-4 {
				t.Fatalf("DotFloat32 mismatch K=%d n=%d: got=%v want=%v", tc.K, n, gotCPU, want[n])
			}
		}

		got := make([]float32, tc.N)
		gemvFloat32Range(got, a, b, tc.K, tc.startN, tc.endN)

		const tol = 1e-4
		for n := tc.startN; n < tc.endN; n++ {
			diff := math.Abs(float64(got[n] - want[n]))
			if diff > tol {
				t.Fatalf("K=%d N=%d range=[%d,%d) n=%d: got=%v want=%v diff=%g",
					tc.K, tc.N, tc.startN, tc.endN, n, got[n], want[n], diff)
			}
		}
	}
}