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

package linalg

import (
	"math"
	"testing"

	core "sourcedock.dev/petrbalvin/tensor/internal/core"
)

// Dense solver benchmarks guard the factorisation kernels. Inputs are
// diagonally dominant so every path stays on the well-conditioned side
// and the timings measure the kernel, not pivoting churn.

func benchMatrix(b *testing.B, n int, spd bool) *core.Array {
	b.Helper()
	v := make([]float64, n*n)
	for i := range n {
		for j := range n {
			v[i*n+j] = float64((i*7+j*13)%11) - 5
		}
		v[i*n+i] += float64(n)
		if spd {
			// Symmetrise, then push the diagonal past the row sum so
			// positive definiteness is guaranteed by strict diagonal
			// dominance.
			for j := i + 1; j < n; j++ {
				avg := (v[i*n+j] + v[j*n+i]) / 2
				v[i*n+j], v[j*n+i] = avg, avg
			}
			row := float64(0)
			for j := range n {
				if j != i {
					row += math.Abs(v[i*n+j])
				}
			}
			v[i*n+i] += row
		}
	}
	a, err := core.FromFloats(v, n, n)
	if err != nil {
		b.Fatal(err)
	}
	return a
}

func benchRHS(b *testing.B, n int) *core.Array {
	b.Helper()
	v := make([]float64, n)
	for i := range v {
		v[i] = float64(i%9) - 4
	}
	x, err := core.FromFloats(v, n)
	if err != nil {
		b.Fatal(err)
	}
	return x
}

func BenchmarkSolve64(b *testing.B) {
	a := benchMatrix(b, 64, false)
	x := benchRHS(b, 64)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Solve(a, x); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkSolve256(b *testing.B) {
	a := benchMatrix(b, 256, false)
	x := benchRHS(b, 256)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Solve(a, x); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkInv256(b *testing.B) {
	a := benchMatrix(b, 256, false)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Inv(a); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkCholesky256(b *testing.B) {
	a := benchMatrix(b, 256, true)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Cholesky(a); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkQR256(b *testing.B) {
	a := benchMatrix(b, 256, false)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := QR(a); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkSVD128(b *testing.B) {
	a := benchMatrix(b, 128, true)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, _, err := SVD(a); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkDet256(b *testing.B) {
	a := benchMatrix(b, 256, false)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Det(a); err != nil {
			b.Fatal(err)
		}
	}
}