package cpu

import (
	"math/rand"
	"testing"
)

func TestDotFloat32MatchesScalar(t *testing.T) {
	a := []float32{1, 2, 3, 4, 5, 6, 7, 8}
	b := []float32{8, 7, 6, 5, 4, 3, 2, 1}

	got := DotFloat32(a, b)
	want := dotFloat32Scalar(a, b)

	if diff := absDiff(got, want); diff > 1e-5 {
		t.Fatalf("dot mismatch: got %f want %f (diff %f)", got, want, diff)
	}
}

func TestDotFloat32ZeroLen(t *testing.T) {
	got := DotFloat32(nil, nil)
	if got != 0 {
		t.Fatalf("expected 0 for empty slices, got %f", got)
	}
}

func BenchmarkDotFloat32(b *testing.B) {
	size := 256
	a := make([]float32, size)
	bb := make([]float32, size)
	for i := 0; i < size; i++ {
		a[i] = rand.Float32()
		bb[i] = rand.Float32()
	}

	b.ReportAllocs()
	b.SetBytes(int64(size * 4 * 2))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		_ = DotFloat32(a, bb)
	}
}

func BenchmarkDotFloat32Scalar(b *testing.B) {
	size := 256
	a := make([]float32, size)
	bb := make([]float32, size)
	for i := 0; i < size; i++ {
		a[i] = rand.Float32()
		bb[i] = rand.Float32()
	}

	b.ReportAllocs()
	b.SetBytes(int64(size * 4 * 2))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		_ = dotFloat32Scalar(a, bb)
	}
}

func absDiff(a, b float32) float32 {
	if a > b {
		return a - b
	}
	return b - a
}

func TestAxpy(t *testing.T) {
	x := []float32{1, 2, 3}
	y := []float32{4, 5, 6}
	Axpy(2, x, y)
	want := []float32{6, 9, 12}
	for i := range y {
		if diff := absDiff(y[i], want[i]); diff > 1e-6 {
			t.Fatalf("axpy mismatch at %d: got %f want %f", i, y[i], want[i])
		}
	}
}

func BenchmarkDotFloat32Ptr(b *testing.B) {
	size := 256
	a := make([]float32, size)
	bb := make([]float32, size)
	for i := 0; i < size; i++ {
		a[i] = rand.Float32()
		bb[i] = rand.Float32()
	}

	b.ReportAllocs()
	b.SetBytes(int64(size * 4 * 2))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		_ = DotFloat32Ptr(&a[0], &bb[0], size)
	}
}

func BenchmarkAxpy(b *testing.B) {
	size := 256
	x := make([]float32, size)
	y := make([]float32, size)
	for i := 0; i < size; i++ {
		x[i] = rand.Float32()
		y[i] = rand.Float32()
	}
	alpha := float32(0.7)

	b.ReportAllocs()
	b.SetBytes(int64(size * 4 * 2))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		Axpy(alpha, x, y)
	}
}

func BenchmarkAxpyPtr(b *testing.B) {
	size := 256
	x := make([]float32, size)
	y := make([]float32, size)
	for i := 0; i < size; i++ {
		x[i] = rand.Float32()
		y[i] = rand.Float32()
	}
	alpha := float32(0.7)

	b.ReportAllocs()
	b.SetBytes(int64(size * 4 * 2))
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		AxpyPtr(alpha, &x[0], &y[0], size)
	}
}