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

package linalg

import (
	"fmt"
	"testing"

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

// Benchmarks for the dense LU kernel, the sparse factorisations and the
// sparse rank-one sweep. Every input is built from fixed literals, so
// the patterns, the pivots and the orderings are the same on every run.

// denseLUInput builds an n×n row-major matrix whose rows view one flat
// buffer, with a pristine copy the benchmark restores from: the
// factorisation consumes its argument, and rebuilding the matrix inside
// the timed region would measure the rebuild.
func denseLUInput(n int) (rows [][]float64, pristine []float64) {
	flat := make([]float64, n*n)
	for i := range n {
		for j := range n {
			flat[i*n+j] = float64((i*5+j*11)%13) - 6
		}
		// Diagonal dominance keeps the elimination on the
		// well-conditioned side, so the timing measures the kernel.
		flat[i*n+i] += float64(2 * n)
	}
	rows = make([][]float64, n)
	for i := range n {
		rows[i] = flat[i*n : (i+1)*n]
	}
	pristine = make([]float64, len(flat))
	copy(pristine, flat)
	return rows, pristine
}

func restoreDenseRows(rows [][]float64, pristine []float64) {
	off := 0
	for _, row := range rows {
		copy(row, pristine[off:off+len(row)])
		off += len(row)
	}
}

// BenchmarkDenseLUFactor measures the LU elimination alone at the sizes
// the solvers reach.
func BenchmarkDenseLUFactor(b *testing.B) {
	for _, n := range []int{128, 256, 512} {
		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
			rows, pristine := denseLUInput(n)
			b.ReportAllocs()
			for b.Loop() {
				restoreDenseRows(rows, pristine)
				base.Factor(rows)
			}
		})
	}
}

// denseLUSolveInput builds the same matrix as a core array, for the
// end-to-end solve path.
func denseLUSolveInput(b *testing.B, n int) (*core.Array, *core.Array) {
	b.Helper()
	flat := make([]float64, n*n)
	for i := range n {
		for j := range n {
			flat[i*n+j] = float64((i*5+j*11)%13) - 6
		}
		flat[i*n+i] += float64(2 * n)
	}
	a, err := core.FromFloats(flat, n, n)
	if err != nil {
		b.Fatal(err)
	}
	rhs := make([]float64, n)
	for i := range rhs {
		rhs[i] = float64(i%7) - 3
	}
	x, err := core.FromFloats(rhs, n)
	if err != nil {
		b.Fatal(err)
	}
	return a, x
}

// BenchmarkDenseLUSolve measures Solve end to end: the copy of the
// matrix, the factorisation, the permutation and the substitution.
func BenchmarkDenseLUSolve(b *testing.B) {
	for _, n := range []int{128, 256, 512} {
		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
			a, x := denseLUSolveInput(b, n)
			b.ReportAllocs()
			for b.Loop() {
				if _, err := Solve(a, x); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}

// sparseCOOFrom assembles a COO matrix from the triplets a builder
// appended, refusing nothing: the builders below emit canonical,
// duplicate-free triplets.
func sparseCOOFrom(b *testing.B, n int, idx []int64, vals []float64) *core.SparseCOO {
	b.Helper()
	indices, err := core.FromInts(idx, len(vals), 2)
	if err != nil {
		b.Fatal(err)
	}
	coo, err := core.NewSparseCOO(indices, floatsToArray(vals, []int{len(vals)}), []int{n, n})
	if err != nil {
		b.Fatal(err)
	}
	return coo
}

// gridLaplacian builds the 5-point Laplacian on a w×h grid in row-major
// order: symmetric positive definite, banded, and the pattern every
// sparse direct solver is measured on.
func gridLaplacian(b *testing.B, w, h int) *core.SparseCOO {
	b.Helper()
	n := w * h
	idx := make([]int64, 0, 5*n)
	vals := make([]float64, 0, 5*n)
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	at := func(x, y int) int { return y*w + x }
	for y := range h {
		for x := range w {
			add(at(x, y), at(x, y), 4)
			if x+1 < w {
				add(at(x, y), at(x+1, y), -1)
				add(at(x+1, y), at(x, y), -1)
			}
			if y+1 < h {
				add(at(x, y), at(x, y+1), -1)
				add(at(x, y+1), at(x, y), -1)
			}
		}
	}
	return sparseCOOFrom(b, n, idx, vals)
}

// arrowHead builds the arrowhead matrix of order n: a diagonal plus a
// dense first row and column. The diagonal dominates the first arrow
// (n+1 against n−1), so the matrix is positive definite, and the
// ordering has a genuine choice to make on the arrow.
func arrowHead(b *testing.B, n int) *core.SparseCOO {
	b.Helper()
	idx := make([]int64, 0, 3*n)
	vals := make([]float64, 0, 3*n)
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	add(0, 0, float64(n+1))
	for i := 1; i < n; i++ {
		add(i, i, float64(i+2))
		add(0, i, 1)
		add(i, 0, 1)
	}
	return sparseCOOFrom(b, n, idx, vals)
}

// bandedNonsymmetric builds a banded, diagonally dominant matrix with a
// periodic spike below the diagonal: every tenth column has a
// subdiagonal entry larger than its diagonal, so the partial pivoting
// of the LU factorisation swaps rows and its label bookkeeping is
// exercised rather than measured cold.
func bandedNonsymmetric(b *testing.B, n int) *core.SparseCOO {
	b.Helper()
	idx := make([]int64, 0, 6*n)
	vals := make([]float64, 0, 6*n)
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	for i := range n {
		d := 6.0
		if i%10 == 3 {
			d = 0.5 // the spike below takes this column's pivot
		}
		add(i, i, d)
		if i+1 < n {
			add(i, i+1, 1+0.25*float64(i%3))
			add(i+1, i, 1.5+0.5*float64(i%5))
		}
		if i+3 < n {
			add(i, i+3, 0.5)
		}
	}
	return sparseCOOFrom(b, n, idx, vals)
}

// swapHeavyBanded builds a tridiagonal matrix that pivots on every
// column: a small diagonal against a large subdiagonal, so the largest
// entry at or below the diagonal is always the row below. The factor
// stays banded, so the measurement is the pivot bookkeeping rather
// than the elimination arithmetic.
func swapHeavyBanded(b *testing.B, n int) *core.SparseCOO {
	b.Helper()
	idx := make([]int64, 0, 3*n)
	vals := make([]float64, 0, 3*n)
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	for i := range n {
		add(i, i, 0.125)
		if i+1 < n {
			add(i, i+1, 1)
			add(i+1, i, 16)
		}
	}
	return sparseCOOFrom(b, n, idx, vals)
}

// BenchmarkSparseLUFactorSwapped measures the elimination of a system
// that pivots at every column, so the cost of relabelling L's stored
// rows is part of the number.
func BenchmarkSparseLUFactorSwapped(b *testing.B) {
	for _, n := range []int{100, 200, 400} {
		coo := swapHeavyBanded(b, n)
		b.Run(fmt.Sprintf("banded-%d", n), func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				if _, err := NewSparseLU(coo); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}

// denseRHS builds a deterministic right hand side of length n.
func denseRHS(b *testing.B, n int) *core.Array {
	b.Helper()
	v := make([]float64, n)
	for i := range v {
		v[i] = float64(i%11) - 5 + 0.5*float64(i%3)
	}
	x, err := core.FromFloats(v, n)
	if err != nil {
		b.Fatal(err)
	}
	return x
}

// BenchmarkSparseCholeskyFactor measures the symbolic and numeric
// elimination of a few hundred unknowns: the Laplacian grid with the
// natural order, then the same grid through the reverse Cuthill-McKee
// order, which meets a much smaller factor.
func BenchmarkSparseCholeskyFactor(b *testing.B) {
	coo := gridLaplacian(b, 19, 19)
	b.Run("grid-natural", func(b *testing.B) {
		b.ReportAllocs()
		for b.Loop() {
			if _, err := NewSparseCholesky(coo, SparseOrderingNatural); err != nil {
				b.Fatal(err)
			}
		}
	})
	b.Run("grid-rcm", func(b *testing.B) {
		b.ReportAllocs()
		for b.Loop() {
			if _, err := NewSparseCholesky(coo, SparseOrderingReverseCuthillMcKee); err != nil {
				b.Fatal(err)
			}
		}
	})
	coo = arrowHead(b, 300)
	b.Run("arrow-natural", func(b *testing.B) {
		b.ReportAllocs()
		for b.Loop() {
			if _, err := NewSparseCholesky(coo, SparseOrderingNatural); err != nil {
				b.Fatal(err)
			}
		}
	})
}

// BenchmarkSparseCholeskySolve measures the solve on a factor built
// once outside the loop: the gather, the two substitutions and the
// scatter.
func BenchmarkSparseCholeskySolve(b *testing.B) {
	for _, size := range []struct {
		name     string
		w, h     int
		ordering SparseOrdering
	}{
		{"grid-natural", 19, 19, SparseOrderingNatural},
		{"grid-rcm", 19, 19, SparseOrderingReverseCuthillMcKee},
	} {
		coo := gridLaplacian(b, size.w, size.h)
		f, err := NewSparseCholesky(coo, size.ordering)
		if err != nil {
			b.Fatal(err)
		}
		rhs := denseRHS(b, size.w*size.h)
		b.Run(size.name, func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				if _, err := f.Solve(rhs); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}

// BenchmarkSparseLUFactor measures the left-looking elimination of a
// banded nonsymmetric system with pivoting.
func BenchmarkSparseLUFactor(b *testing.B) {
	for _, n := range []int{200, 400} {
		coo := bandedNonsymmetric(b, n)
		b.Run(fmt.Sprintf("banded-%d", n), func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				if _, err := NewSparseLU(coo); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}

// BenchmarkSparseLUSolve measures the forward and backward substitution
// on a factor built once outside the loop.
func BenchmarkSparseLUSolve(b *testing.B) {
	const n = 400
	coo := bandedNonsymmetric(b, n)
	f, err := NewSparseLU(coo)
	if err != nil {
		b.Fatal(err)
	}
	rhs := denseRHS(b, n)
	b.Run("banded-400", func(b *testing.B) {
		b.ReportAllocs()
		for b.Loop() {
			if _, err := f.Solve(rhs); err != nil {
				b.Fatal(err)
			}
		}
	})
}

// sparseCholRankOneInput factors the tridiagonal Laplacian and returns
// the factor with a two-entry vector whose support sits on the stored
// pattern, so both the update and the downdate are accepted.
func sparseCholRankOneInput(b *testing.B, n int) (*SparseCholesky, *core.Array) {
	b.Helper()
	idx := make([]int64, 0, 3*n)
	vals := make([]float64, 0, 3*n)
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	for i := range n {
		add(i, i, 4)
		if i+1 < n {
			add(i, i+1, -1)
			add(i+1, i, -1)
		}
	}
	coo := sparseCOOFrom(b, n, idx, vals)
	f, err := NewSparseCholesky(coo, SparseOrderingNatural)
	if err != nil {
		b.Fatal(err)
	}
	x := core.New(core.Float, []int{n}...)
	x.RawFloats()[0] = 0.5
	x.RawFloats()[1] = 0.25
	return f, x
}

// BenchmarkSparseCholRankOnePair measures one rank-one update followed
// by the matching downdate: each iteration leaves the factor where it
// found it, so the sweep is timed rather than the refactorisation the
// modification replaces.
func BenchmarkSparseCholRankOnePair(b *testing.B) {
	for _, n := range []int{150, 300} {
		f, x := sparseCholRankOneInput(b, n)
		b.Run(fmt.Sprintf("tridiag-%d", n), func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				if err := f.Update(x); err != nil {
					b.Fatal(err)
				}
				if err := f.Downdate(x); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}

// BenchmarkSparseCholUpdateRefusal measures the refusal path: a vector
// whose support reaches outside the stored pattern, so every call
// returns the pattern error after the fill check has walked the support.
// It is the cost the check pays before it can refuse, and the error
// construction is deliberately part of it.
func BenchmarkSparseCholUpdateRefusal(b *testing.B) {
	const n = 300
	f, _ := sparseCholRankOneInput(b, n)
	x := core.New(core.Float, []int{n}...)
	x.RawFloats()[0] = 0.5
	x.RawFloats()[n-1] = 0.25
	b.ReportAllocs()
	for b.Loop() {
		if err := f.Update(x); err == nil {
			b.Fatal("expected the update to need fill the pattern does not hold")
		}
	}
}