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

package linalg

import (
	"testing"

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

// Iterative solver benchmarks: the Krylov solves and the least-squares
// recursions, with and without the ILU(0) preconditioner. Every system
// is built from fixed literals so a run is reproducible, and each one
// is sized to finish well inside a second, which keeps the timings on
// the solvers' own per-iteration work rather than on input
// construction.

// benchGridCOO builds the five-point stencil of an nx by ny grid as an
// n x n COO matrix (n = nx·ny) in natural row ordering. The diagonal is
// 4 and the vertical coupling −1 on both sides. A nonzero skew makes
// the horizontal coupling −1∓skew, which breaks the symmetry while
// keeping the matrix diagonally dominant; skew 0 leaves it symmetric
// positive-definite.
func benchGridCOO(b *testing.B, nx, ny int, skew float64) *core.SparseCOO {
	b.Helper()
	n := nx * ny
	idx := make([]int64, 0, 5*n*2)
	vals := make([]float64, 0, 5*n)
	add := func(i, j int, v float64) {
		idx = append(idx, int64(i), int64(j))
		vals = append(vals, v)
	}
	for r := range ny {
		for c := range nx {
			i := r*nx + c
			add(i, i, 4)
			if c+1 < nx {
				add(i, i+1, -1-skew)
			}
			if c > 0 {
				add(i, i-1, -1+skew)
			}
			if r+1 < ny {
				add(i, i+nx, -1)
			}
			if r > 0 {
				add(i, i-nx, -1)
			}
		}
	}
	indices, err := core.FromInts(idx, len(vals), 2)
	if err != nil {
		b.Fatal(err)
	}
	values, err := core.FromFloats(vals, len(vals))
	if err != nil {
		b.Fatal(err)
	}
	coo, err := core.NewSparseCOO(indices, values, []int{n, n})
	if err != nil {
		b.Fatal(err)
	}
	return coo
}

// benchTallCOO builds the overdetermined (2n − 2) x n matrix whose
// first n rows are the identity and whose remaining rows are the
// second-difference stencil over three adjacent columns. It is full
// rank and banded, so LSQR and LSMR converge in a bounded number of
// steps.
func benchTallCOO(b *testing.B, n int) *core.SparseCOO {
	b.Helper()
	m := 2*n - 2
	idx := make([]int64, 0, (4*n)*2)
	vals := make([]float64, 0, 4*n)
	add := func(i, j int, v float64) {
		idx = append(idx, int64(i), int64(j))
		vals = append(vals, v)
	}
	for j := range n {
		add(j, j, 1)
	}
	for r := range n - 2 {
		i := n + r
		add(i, r, -1)
		add(i, r+1, 2)
		add(i, r+2, -1)
	}
	indices, err := core.FromInts(idx, len(vals), 2)
	if err != nil {
		b.Fatal(err)
	}
	values, err := core.FromFloats(vals, len(vals))
	if err != nil {
		b.Fatal(err)
	}
	coo, err := core.NewSparseCOO(indices, values, []int{m, n})
	if err != nil {
		b.Fatal(err)
	}
	return coo
}

// benchPatternRHS builds the deterministic right-hand side of length n
// from a fixed integer pattern.
func benchPatternRHS(b *testing.B, n int) *core.Array {
	b.Helper()
	v := make([]float64, n)
	for i := range v {
		v[i] = float64(i%7) - 3
	}
	a, err := core.FromFloats(v, n)
	if err != nil {
		b.Fatal(err)
	}
	return a
}

// BenchmarkGMRESSparse measures a restarted GMRES solve of a
// nonsymmetric five-point system.
func BenchmarkGMRESSparse(b *testing.B) {
	const n = 400
	coo := benchGridCOO(b, 20, 20, 0.5)
	c, err := cooToCSR(coo, "BenchmarkGMRESSparse")
	if err != nil {
		b.Fatal(err)
	}
	rowStart, colIdx, vals := c.rowStart, c.colIdx, c.vals
	op := func(v *core.Array) (*core.Array, error) {
		out := core.New(core.Float, n)
		dst := out.RawFloats()
		for i := range n {
			sum := 0.0
			for p := rowStart[i]; p < rowStart[i+1]; p++ {
				sum += vals[p] * v.FloatAt(colIdx[p])
			}
			dst[i] = sum
		}
		return out, nil
	}
	bv := benchPatternRHS(b, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := GMRES(op, bv, 30, 200, 1e-10); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSpSolveCG measures a conjugate-gradient solve of a
// symmetric five-point system under the Jacobi preconditioner.
func BenchmarkSpSolveCG(b *testing.B) {
	coo := benchGridCOO(b, 20, 20, 0)
	bv := benchPatternRHS(b, 400)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := SpSolve(coo, bv, 1e-10, 0); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSpSolveCGILU measures the same system under the ILU(0)
// preconditioner.
func BenchmarkSpSolveCGILU(b *testing.B) {
	coo := benchGridCOO(b, 20, 20, 0)
	bv := benchPatternRHS(b, 400)
	ilu, err := NewSparseILU(coo)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := SpSolve(coo, bv, 1e-10, 0, ilu); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSpSolveBiCGSTAB measures a BiCGSTAB solve of a nonsymmetric
// five-point system under the Jacobi preconditioner.
func BenchmarkSpSolveBiCGSTAB(b *testing.B) {
	coo := benchGridCOO(b, 20, 20, 0.5)
	bv := benchPatternRHS(b, 400)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := SpSolveBiCGSTAB(coo, bv, 1e-10, 0); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSpSolveBiCGSTABILU measures the same system under the
// ILU(0) preconditioner.
func BenchmarkSpSolveBiCGSTABILU(b *testing.B) {
	coo := benchGridCOO(b, 20, 20, 0.5)
	bv := benchPatternRHS(b, 400)
	ilu, err := NewSparseILU(coo)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := SpSolveBiCGSTAB(coo, bv, 1e-10, 0, ilu); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSpLSQR measures the Golub-Kahan least-squares recursion of
// SpLSQR on an overdetermined banded system.
func BenchmarkSpLSQR(b *testing.B) {
	const n = 400
	coo := benchTallCOO(b, n)
	bv := benchPatternRHS(b, 2*n-2)
	_, info, err := SpLSQR(coo, bv, 1e-10, 0, 0)
	if err != nil {
		b.Fatal(err)
	}
	steps := float64(info.Iterations)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := SpLSQR(coo, bv, 1e-10, 0, 0); err != nil {
			b.Fatal(err)
		}
	}
	// The system and the tolerance are fixed, so the step count is too:
	// reporting it shows the size of the work the timing covers. It goes
	// after the loop, which clears any metric reported before it.
	b.ReportMetric(steps, "steps")
}

// BenchmarkSpLSMR measures the LSMR recursion on the same system.
func BenchmarkSpLSMR(b *testing.B) {
	const n = 400
	coo := benchTallCOO(b, n)
	bv := benchPatternRHS(b, 2*n-2)
	_, info, err := SpLSMR(coo, bv, 1e-10, 0, 0)
	if err != nil {
		b.Fatal(err)
	}
	steps := float64(info.Iterations)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := SpLSMR(coo, bv, 1e-10, 0, 0); err != nil {
			b.Fatal(err)
		}
	}
	b.ReportMetric(steps, "steps")
}