// 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/core"
)

// Benchmarks for the sparse surface: the matrix-vector and matrix-matrix
// products, the coordinate-to-compressed construction, and the Krylov
// solvers that drive them. Every input is a fixed literal pattern, so a
// run is reproducible and two builds compare like with like.

// benchSink keeps a kernel's result reachable so the compiler cannot
// elide the call it came from.
var (
	benchSinkCSR *SparseCSR
	benchSinkCSC *SparseCSC
	benchSinkArr *core.Array
)

// laplacianTriples returns the flat (row, col, value) triples of the
// n×n symmetric banded Laplacian of the given half-bandwidth: 2·width
// on the diagonal and -1 at every offset up to width. It is symmetric
// positive definite by strict diagonal dominance, which is what the
// Lanczos and conjugate-gradient benchmarks need, and its non-zeros per
// row are 2·width+1 at every row.
func laplacianTriples(n, width int) []float64 {
	tri := make([]float64, 0, n*(2*width+1)*3)
	for i := range n {
		tri = append(tri, float64(i), float64(i), float64(2*width))
		for d := 1; d <= width; d++ {
			j := i + d
			if j >= n {
				break
			}
			tri = append(tri, float64(i), float64(j), -1)
			tri = append(tri, float64(j), float64(i), -1)
		}
	}
	return tri
}

// duplicateTriples returns the triples of a matrix whose coordinates
// repeat: entry k carries row (k·13) mod n and column (k·37) mod n, so
// each coordinate appears about len/n times and the construction path
// has to merge duplicates rather than only sort distinct pairs. The
// values step along a literal modular sequence, so no merge cancels by
// accident.
func duplicateTriples(n, len_ int) []float64 {
	tri := make([]float64, 0, len_*3)
	for k := range len_ {
		row := (k * 13) % n
		col := (k * 37) % n
		v := float64((k*7)%23) - 11
		if v == 0 {
			v = 1.5
		}
		tri = append(tri, float64(row), float64(col), v)
	}
	return tri
}

// benchCOO builds a sparse matrix from flat (row, col, value) triples.
func benchCOO(b *testing.B, rows, cols int, tri []float64) *core.SparseCOO {
	b.Helper()
	nnz := len(tri) / 3
	idx := make([]int64, 0, nnz*2)
	vals := make([]float64, 0, nnz)
	for i := range nnz {
		idx = append(idx, int64(tri[i*3]), int64(tri[i*3+1]))
		vals = append(vals, tri[i*3+2])
	}
	indices, err := core.FromInts(idx, nnz, 2)
	if err != nil {
		b.Fatalf("FromInts: %v", err)
	}
	values, err := core.FromFloats(vals, nnz)
	if err != nil {
		b.Fatalf("FromFloats: %v", err)
	}
	coo, err := core.NewSparseCOO(indices, values, []int{rows, cols})
	if err != nil {
		b.Fatalf("NewSparseCOO: %v", err)
	}
	return coo
}

// benchVector returns a dense vector of length n whose entries come
// from a literal modular sequence, deliberately without zeros so a
// scaled product cannot skip work.
func benchVector(b *testing.B, n int) *core.Array {
	b.Helper()
	v := make([]float64, n)
	for i := range n {
		v[i] = float64((i*11)%17)/8 - 1
		if v[i] == 0 {
			v[i] = 0.25
		}
	}
	arr, err := core.FromFloats(v, n)
	if err != nil {
		b.Fatalf("FromFloats: %v", err)
	}
	return arr
}

func BenchmarkSparseCSRMatVec(b *testing.B) {
	for _, n := range []int{256, 131072} {
		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
			coo := benchCOO(b, n, n, laplacianTriples(n, 1))
			csr, err := CSRFromCOO(coo)
			if err != nil {
				b.Fatalf("CSRFromCOO: %v", err)
			}
			x := benchVector(b, n)
			for b.Loop() {
				benchSinkArr, _ = csr.MatVec(x)
			}
		})
	}
}

func BenchmarkSparseCSCMatVec(b *testing.B) {
	for _, n := range []int{256, 131072} {
		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
			coo := benchCOO(b, n, n, laplacianTriples(n, 1))
			csc, err := CSCFromCOO(coo)
			if err != nil {
				b.Fatalf("CSCFromCOO: %v", err)
			}
			x := benchVector(b, n)
			for b.Loop() {
				benchSinkArr, _ = csc.MatVec(x)
			}
		})
	}
}

func BenchmarkSparseCSRMatMulSparse(b *testing.B) {
	const n = 256
	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
		coo := benchCOO(b, n, n, laplacianTriples(n, 4))
		csr, err := CSRFromCOO(coo)
		if err != nil {
			b.Fatalf("CSRFromCOO: %v", err)
		}
		for b.Loop() {
			benchSinkCSR, _ = csr.MatMulSparse(csr)
		}
	})
}

func BenchmarkSparseCSRMatMulDense(b *testing.B) {
	const n = 512
	const cols = 4
	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
		coo := benchCOO(b, n, n, laplacianTriples(n, 2))
		csr, err := CSRFromCOO(coo)
		if err != nil {
			b.Fatalf("CSRFromCOO: %v", err)
		}
		v := make([]float64, n*cols)
		for i := range v {
			v[i] = float64((i*11)%17)/8 - 1
		}
		x, err := core.FromFloats(v, n, cols)
		if err != nil {
			b.Fatalf("FromFloats: %v", err)
		}
		for b.Loop() {
			benchSinkArr, _ = csr.MatMulDense(x)
		}
	})
}

func BenchmarkCSRFromCOODuplicates(b *testing.B) {
	const n = 1500
	tri := duplicateTriples(n, 6000)
	b.Run("nnz=6000", func(b *testing.B) {
		coo := benchCOO(b, n, n, tri)
		for b.Loop() {
			benchSinkCSR, _ = CSRFromCOO(coo)
		}
	})
}

func BenchmarkCSCFromCOODuplicates(b *testing.B) {
	const n = 1500
	tri := duplicateTriples(n, 6000)
	b.Run("nnz=6000", func(b *testing.B) {
		coo := benchCOO(b, n, n, tri)
		for b.Loop() {
			benchSinkCSC, _ = CSCFromCOO(coo)
		}
	})
}

// BenchmarkSpEigenBanded measures the real symmetric Krylov solve end
// to end: the coordinate construction, the banded matrix-vector
// products and the full reorthogonalisation of every Lanczos step.
func BenchmarkSpEigenBanded(b *testing.B) {
	const n = 800
	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
		coo := benchCOO(b, n, n, laplacianTriples(n, 1))
		for b.Loop() {
			benchSinkArr, _, _ = SpEigen(coo, 6, core.NewGenerator(7))
		}
	})
}

// BenchmarkSpEigenGeneralBanded measures the Arnoldi solve on the same
// matrix: one explicit orthogonalisation per basis column instead of
// the three-term recurrence.
func BenchmarkSpEigenGeneralBanded(b *testing.B) {
	const n = 400
	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
		coo := benchCOO(b, n, n, laplacianTriples(n, 1))
		for b.Loop() {
			benchSinkArr, _, _ = SpEigenGeneral(coo, 4, core.NewGenerator(7))
		}
	})
}

func BenchmarkSpEigenComplexBanded(b *testing.B) {
	const n = 400
	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
		idx := make([]int64, 0, n*6)
		vals := make([]complex128, 0, n*3)
		for i := range n {
			idx = append(idx, int64(i), int64(i))
			vals = append(vals, 2+0i)
			if i+1 < n {
				idx = append(idx, int64(i), int64(i+1))
				vals = append(vals, 1i)
				idx = append(idx, int64(i+1), int64(i))
				vals = append(vals, -1i)
			}
		}
		indices, err := core.FromInts(idx, len(vals), 2)
		if err != nil {
			b.Fatalf("FromInts: %v", err)
		}
		values, err := core.FromComplexes(vals, len(vals))
		if err != nil {
			b.Fatalf("FromComplexes: %v", err)
		}
		coo, err := core.NewSparseCOO(indices, values, []int{n, n})
		if err != nil {
			b.Fatalf("NewSparseCOO: %v", err)
		}
		for b.Loop() {
			benchSinkArr, _, _ = SpEigenComplex(coo, 4, core.NewGenerator(7))
		}
	})
}

// BenchmarkSpExpApplyBanded measures the Krylov projection of the
// matrix exponential action, which shares the Lanczos kernel with the
// symmetric eigensolver.
func BenchmarkSpExpApplyBanded(b *testing.B) {
	const n = 400
	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
		coo := benchCOO(b, n, n, laplacianTriples(n, 1))
		v := benchVector(b, n)
		for b.Loop() {
			benchSinkArr, _ = SpExpApply(coo, v, 0)
		}
	})
}

// BenchmarkSpSolveComplexCGBanded measures the conjugate-gradient solve
// whose every step is a complex sparse matrix-vector product with a
// Jacobi preconditioner.
func BenchmarkSpSolveComplexCGBanded(b *testing.B) {
	const n = 256
	b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
		idx := make([]int64, 0, n*6)
		vals := make([]complex128, 0, n*3)
		for i := range n {
			idx = append(idx, int64(i), int64(i))
			vals = append(vals, 4+0i)
			if i+1 < n {
				idx = append(idx, int64(i), int64(i+1))
				vals = append(vals, 1i)
				idx = append(idx, int64(i+1), int64(i))
				vals = append(vals, -1i)
			}
		}
		indices, err := core.FromInts(idx, len(vals), 2)
		if err != nil {
			b.Fatalf("FromInts: %v", err)
		}
		values, err := core.FromComplexes(vals, len(vals))
		if err != nil {
			b.Fatalf("FromComplexes: %v", err)
		}
		coo, err := core.NewSparseCOO(indices, values, []int{n, n})
		if err != nil {
			b.Fatalf("NewSparseCOO: %v", err)
		}
		ones := make([]complex128, n)
		for i := range ones {
			ones[i] = 1
		}
		rhs, err := core.FromComplexes(ones, n)
		if err != nil {
			b.Fatalf("FromComplexes: %v", err)
		}
		for b.Loop() {
			benchSinkArr, _ = SpSolveComplexCG(coo, rhs, 1e-12, 400)
		}
	})
}