// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package core

import "testing"

// Dtype and split companions to the matmul benchmarks: bench_test.go,
// mat_tile_bench_test.go and mat_kernel_bench_test.go pin the float64
// shapes, so these pin the complex 2-D product, the float32, int and
// complex vector walks, and the vector split's per-worker floor. The
// operands are fixed literals, so a run is comparable to the next.

// benchMatMulComplexRect runs the n×k·k×m complex product under the
// default worker policy.
func benchMatMulComplexRect(b *testing.B, n, k, m int) {
	b.Helper()
	av := make([]complex128, n*k)
	bv := make([]complex128, k*m)
	for i := range av {
		av[i] = complex(float64(i%7)-3, float64(i%5)-2)
	}
	for i := range bv {
		bv[i] = complex(float64(i%5)-2, float64(i%11)-5)
	}
	a, err := FromComplexes(av, n, k)
	if err != nil {
		b.Fatal(err)
	}
	c, err := FromComplexes(bv, k, m)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := MatMul2D(a, c); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkMatMulComplex256 is the smallest complex square the split
// engages, one size above matMulComplexParallelMin.
func BenchmarkMatMulComplex256(b *testing.B) { benchMatMulComplexRect(b, 256, 256, 256) }

// BenchmarkMatMulComplex512 is the complex square product: its b
// payload is sixteen bytes per element, so sharing one b row across a
// panel carries the most weight here.
func BenchmarkMatMulComplex512(b *testing.B) { benchMatMulComplexRect(b, 512, 512, 512) }

// The complex rectangular extremes, mirroring the float64 shape set.
func BenchmarkMatMulComplexFat64x512x2048(b *testing.B) {
	benchMatMulComplexRect(b, 64, 512, 2048)
}

func BenchmarkMatMulComplexTall512x64x2048(b *testing.B) {
	benchMatMulComplexRect(b, 512, 64, 2048)
}

func BenchmarkMatMulComplexWide2048x64x512(b *testing.B) {
	benchMatMulComplexRect(b, 2048, 64, 512)
}

func BenchmarkMatMulComplexSkinny2048x512x64(b *testing.B) {
	benchMatMulComplexRect(b, 2048, 512, 64)
}

// benchMatVecTyped runs the 2-D×1-D (swap true) or 1-D×2-D (swap
// false) product of a typed pair built from fixed literals, under the
// default worker policy. The vector carries the inner dimension: the
// matrix's column count for a row-dot product, its row count for a
// column-dot one.
func benchMatVecTyped(b *testing.B, dt Dtype, rows, cols int, swap bool) {
	b.Helper()
	n := rows * cols
	flat := make([]float64, n)
	for i := range flat {
		flat[i] = float64(i%7) - 3
	}
	vlen := cols
	if !swap {
		vlen = rows
	}
	vecf := make([]float64, vlen)
	for i := range vecf {
		vecf[i] = float64(i%5) - 2
	}
	var mat, vec *Array
	var err error
	switch dt {
	case Float32:
		f32 := make([]float32, n)
		for i, v := range flat {
			f32[i] = float32(v)
		}
		v32 := make([]float32, vlen)
		for i, v := range vecf {
			v32[i] = float32(v)
		}
		mat, err = FromFloat32s(f32, rows, cols)
		if err == nil {
			vec, err = FromFloat32s(v32, vlen)
		}
	case Int:
		iv := make([]int64, n)
		for i, v := range flat {
			iv[i] = int64(v)
		}
		vv := make([]int64, vlen)
		for i, v := range vecf {
			vv[i] = int64(v)
		}
		mat, err = FromInts(iv, rows, cols)
		if err == nil {
			vec, err = FromInts(vv, vlen)
		}
	default:
		cv := make([]complex128, n)
		for i, v := range flat {
			cv[i] = complex(v, float64(i%5)-2)
		}
		cvv := make([]complex128, vlen)
		for i, v := range vecf {
			cvv[i] = complex(v, float64(i%3)-1)
		}
		mat, err = FromComplexes(cv, rows, cols)
		if err == nil {
			vec, err = FromComplexes(cvv, vlen)
		}
	}
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		var err error
		if swap {
			_, err = MatMul2D(mat, vec)
		} else {
			_, err = MatMul2D(vec, mat)
		}
		if err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkMatVecF32RowDots512x512 pins the float32 row dots, whose
// four-row walk is the float64 twin's structure.
func BenchmarkMatVecF32RowDots512x512(b *testing.B) {
	benchMatVecTyped(b, Float32, 512, 512, true)
}

// BenchmarkMatVecIntRowDots512x512 pins the int row dots.
func BenchmarkMatVecIntRowDots512x512(b *testing.B) { benchMatVecTyped(b, Int, 512, 512, true) }

// BenchmarkMatVecComplexRowDots256x256 pins the complex row dots, whose
// two-row walk is what a sixteen-byte sum chain can hold in registers.
func BenchmarkMatVecComplexRowDots256x256(b *testing.B) {
	benchMatVecTyped(b, Complex, 256, 256, true)
}

// BenchmarkVecMatF32ColDots128x2048 pins the float32 column dots.
func BenchmarkVecMatF32ColDots128x2048(b *testing.B) {
	benchMatVecTyped(b, Float32, 128, 2048, false)
}

// BenchmarkVecMatComplexColDots128x2048 pins the complex column dots.
func BenchmarkVecMatComplexColDots128x2048(b *testing.B) {
	benchMatVecTyped(b, Complex, 128, 2048, false)
}

// BenchmarkMatVecRowDots128x128 sits just above the vector split's
// per-worker floor: 16,384 terms fill at most one worker, so the shape
// measures the lone walk the floor keeps it on.
func BenchmarkMatVecRowDots128x128(b *testing.B) { benchMatVecTyped(b, Float, 128, 128, true) }

// BenchmarkVecMatColDots16x16384 pins the wide column dots: sixteen
// rows by 16,384 output columns, so each worker's band is far wider
// than a cache line.
func BenchmarkVecMatColDots16x16384(b *testing.B) { benchMatVecTyped(b, Float, 16, 16384, false) }

// BenchmarkVecMatColNarrowDots16384x16 is the thin-band shape: sixteen
// output columns of 16,384 terms, where the split hands a worker a band
// no wider than a cache line.
func BenchmarkVecMatColNarrowDots16384x16(b *testing.B) {
	benchMatVecTyped(b, Float, 16384, 16, false)
}