// 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"
	"sourcedock.dev/petrbalvin/tensor/internal/engine"
)

// Decomposition benchmarks for the shapes the solver benchmarks in
// perf_bench_test.go do not cover: the tall least-squares system, the
// pivoted QR sweep, the complex and matrix-function paths, and the SVD
// shapes whose cost sits in the reconstruction kernels rather than in
// the bidiagonalisation. Every input is a fixed literal formula, so the
// timings compare like with like across runs.

// benchDecompFloats fills an m×n row-major matrix with a deterministic
// literal pattern: a diagonally weighted band plus a bounded pseudo-random
// ripple.
func benchDecompFloats(m, n int, diagonal float64) []float64 {
	v := make([]float64, m*n)
	for i := range m {
		for j := range n {
			v[i*n+j] = 0.5*float64((i*11+j*7)%17) - 4 + 0.25*float64((i*j*5)%13)
		}
		v[i*n+i%n] += diagonal
	}
	return v
}

// benchDecompSPD builds a symmetric strictly diagonally dominant
// matrix, which is positive definite however the ripple lands.
func benchDecompSPD(n int) []float64 {
	v := make([]float64, n*n)
	for i := range n {
		for j := range n {
			v[i*n+j] = float64((i*5+j*3)%9) - 4
		}
	}
	for i := range n {
		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 := 0.0
		for j := range n {
			if j != i {
				row += v[i*n+j]
				if v[i*n+j] < 0 {
					row += 2 * -v[i*n+j]
				}
			}
		}
		v[i*n+i] += row + float64(n)
	}
	return v
}

// benchDecompHermitian builds a Hermitian matrix with a real diagonal
// whose entries are well separated, so the Jacobi sweep deflates
// quickly.
func benchDecompHermitian(n int) []complex128 {
	v := make([]complex128, n*n)
	for i := range n {
		v[i*n+i] = complex(4*float64(i)+8, 0)
		for j := i + 1; j < n; j++ {
			re := 0.5*float64((i*3+j*5)%7) - 1.5
			im := 0.25*float64((i+j)%5) - 0.5
			v[i*n+j] = complex(re, im)
			v[j*n+i] = complex(re, -im)
		}
	}
	return v
}

func benchDecompArray(b *testing.B, v []float64, m, n int) *core.Array {
	b.Helper()
	a, err := core.FromFloats(v, m, n)
	if err != nil {
		b.Fatal(err)
	}
	return a
}

func benchDecompComplexArray(b *testing.B, v []complex128, m, n int) *core.Array {
	b.Helper()
	a, err := core.FromComplexes(v, m, n)
	if err != nil {
		b.Fatal(err)
	}
	return a
}

// BenchmarkCholesky512 measures the blocked Cholesky sweep above the
// size where the solver benchmark's 256 stops.
func BenchmarkCholesky512(b *testing.B) {
	a := benchDecompArray(b, benchDecompSPD(512), 512, 512)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := Cholesky(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkLeastSquares512x32 measures the tall least-squares route,
// where the QR factor's orthogonal accumulation dominates.
func BenchmarkLeastSquares512x32(b *testing.B) {
	const m, n = 512, 32
	a := benchDecompArray(b, benchDecompFloats(m, n, 3), m, n)
	rhs, err := core.FromFloats(benchDecompFloats(m, 1, 1), m)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := LeastSquares(a, rhs); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkRRQR192x48 measures the pivoted QR sweep, whose cost is the
// per-step column-norm scan and the reflector application.
func BenchmarkRRQR192x48(b *testing.B) {
	const m, n = 192, 48
	a := benchDecompArray(b, benchDecompFloats(m, n, 3), m, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, _, _, err := RRQR(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkRRQR128x32 measures the pivoted QR sweep at the size where a
// single reflector's work is small enough that crew sizing decides it.
func BenchmarkRRQR128x32(b *testing.B) {
	const m, n = 128, 32
	a := benchDecompArray(b, benchDecompFloats(m, n, 3), m, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, _, _, err := RRQR(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSVDReconstruct1024x16 measures the SVD shapes whose cost sits
// in the reconstruction kernels: the full QR factor is far larger than
// the bidiagonalisation.
func BenchmarkSVDReconstruct1024x16(b *testing.B) {
	const m, n = 1024, 16
	a := benchDecompArray(b, benchDecompFloats(m, n, 2), m, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, _, err := SVD(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkEigenComplex8 measures the complex Hermitian Jacobi sweep at
// the largest size its convergence floor admits: the sweep's off-mass
// threshold sits below the rounding floor the rotations leave behind
// above 8, so a bigger input errors out before it measures anything.
func BenchmarkEigenComplex8(b *testing.B) {
	const n = 8
	a := benchDecompComplexArray(b, benchDecompHermitian(n), n, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := EigenComplex(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSVDComplex96x48 measures the direct complex bidiagonalisation
// and its Golub-Reinsch sweep.
func BenchmarkSVDComplex96x48(b *testing.B) {
	const m, n = 96, 48
	flat := benchDecompFloats(m, n, 2)
	v := make([]complex128, m*n)
	for i := range m * n {
		v[i] = complex(flat[i], 0.5*float64((i*3)%7)-1.5)
	}
	a := benchDecompComplexArray(b, v, m, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, _, err := SVDComplex(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkEigenGeneral32 measures the Hessenberg reduction, the shifted
// QR sweep and the eigenvector back-substitution on a general matrix.
func BenchmarkEigenGeneral32(b *testing.B) {
	const n = 32
	v := benchDecompFloats(n, n, 4)
	a := benchDecompArray(b, v, n, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := EigenGeneral(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSchurComplex48 measures the complex Schur decomposition the
// matrix functions are built on.
func BenchmarkSchurComplex48(b *testing.B) {
	const n = 48
	flat := benchDecompFloats(n, n, 3)
	v := make([]complex128, n*n)
	for i := range n * n {
		v[i] = complex(flat[i], 0.25*float64((i*5)%9)-1)
	}
	a := benchDecompComplexArray(b, v, n, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := SchurComplex(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkMatrixExp32 measures the Padé scaling-and-squaring kernel.
func BenchmarkMatrixExp32(b *testing.B) {
	const n = 32
	v := benchDecompFloats(n, n, 0)
	v[0] = -0.5 // keep the 1-norm inside the unscaled degrees
	a := benchDecompArray(b, v, n, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := MatrixExp(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkMatrixSqrt64 measures the symmetric eigen route of a matrix
// function, whose cost is the eigendecomposition plus two triple
// products.
func BenchmarkMatrixSqrt64(b *testing.B) {
	const n = 64
	a := benchDecompArray(b, benchDecompSPD(n), n, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := MatrixSqrt(a); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkMatrixLogSchur32 measures the Schur-Parlett logarithm, the
// square-root walk and the Mercator series included.
func BenchmarkMatrixLogSchur32(b *testing.B) {
	const n = 32
	flat := benchDecompSPD(n)
	v := benchDecompComplexSPD(flat, n)
	a := benchDecompComplexArray(b, v, n, n)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := MatrixLog(a); err != nil {
			b.Fatal(err)
		}
	}
}

// benchDecompComplexSPD lifts a symmetric positive definite real matrix
// to complex with a small positive imaginary part on the strict upper
// triangle and its conjugate below, keeping the spectrum off the
// non-positive real axis.
func benchDecompComplexSPD(flat []float64, n int) []complex128 {
	v := make([]complex128, n*n)
	for i := range n {
		for j := range n {
			v[i*n+j] = complex(flat[i*n+j], 0)
		}
	}
	for i := range n {
		for j := i + 1; j < n; j++ {
			im := 0.05 * float64((i+j)%4+1)
			v[i*n+j] = complex(flat[i*n+j], im)
			v[j*n+i] = complex(flat[j*n+i], -im)
		}
	}
	return v
}

// BenchmarkTikhonov256x64 measures the SVD solve route, whose cost is a
// full SVD of the system matrix plus the Uᵀb and V·d products.
func BenchmarkTikhonov256x64(b *testing.B) {
	const m, n = 256, 64
	a := benchDecompArray(b, benchDecompFloats(m, n, 2), m, n)
	rhs, err := core.FromFloats(benchDecompFloats(m, 1, 1), m)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := SolveTikhonov(a, rhs, 0.5); err != nil {
			b.Fatal(err)
		}
	}
}

// TestDecompWorkerCountBitIdentity pins the parallel splits: every
// kernel this file benchmarks must answer bit for bit the same whether
// the crew is one goroutine or the machine's full width. A split that
// moved an addend or a boundary shows up here before the oracle's
// smaller pinned shapes see it.
func TestDecompWorkerCountBitIdentity(t *testing.T) {
	prev := engine.SetNumWorkers(1)
	defer engine.SetNumWorkers(prev)

	const (
		cholN = 320
		lsM   = 384
		lsN   = 24
		rrM   = 160
		rrN   = 40
		svdM  = 384
		svdN  = 24
		eigN  = 192
		cplxN = 8
	)
	spd := benchDecompSPD(cholN)
	tall := benchDecompFloats(lsM, lsN, 3)
	rr := benchDecompFloats(rrM, rrN, 3)
	wide := benchDecompFloats(svdM, svdN, 2)
	herm := benchDecompHermitian(cplxN)
	square := benchDecompSPD(eigN)
	complexSquare := make([]complex128, cplxN*cplxN)
	flatComplex := benchDecompFloats(cplxN, cplxN, 3)
	for i := range cplxN * cplxN {
		complexSquare[i] = complex(flatComplex[i], 0.25*float64((i*5)%9)-1)
	}

	spdA := mustFromFloats(t, spd, cholN, cholN)
	tallA := mustFromFloats(t, tall, lsM, lsN)
	rrA := mustFromFloats(t, rr, rrM, rrN)
	wideA := mustFromFloats(t, wide, svdM, svdN)
	hermA := mustFromComplexes(t, herm, cplxN, cplxN)
	squareA := mustFromFloats(t, square, eigN, eigN)
	complexA := mustFromComplexes(t, complexSquare, cplxN, cplxN)
	rhs := mustFromFloats(t, benchDecompFloats(lsM, 1, 1), lsM)
	rrRHS := mustFromFloats(t, benchDecompFloats(rrM, 1, 1), rrM)
	// A tall system wide enough that the blocked Qᵀb dispatches a crew,
	// so the split is compared against the serial order too.
	const lsWideN = 48
	wideLS := mustFromFloats(t, benchDecompFloats(lsM, lsWideN, 3), lsM, lsWideN)

	type snap struct {
		name string
		run  func() []*core.Array
	}
	cases := []snap{
		{"Cholesky", func() []*core.Array {
			l, err := Cholesky(spdA)
			if err != nil {
				t.Fatal(err)
			}
			return []*core.Array{l}
		}},
		{"LeastSquares", func() []*core.Array {
			x, err := LeastSquares(tallA, rhs)
			if err != nil {
				t.Fatal(err)
			}
			return []*core.Array{x}
		}},
		{"LeastSquaresBlockedQtb", func() []*core.Array {
			x, err := LeastSquares(wideLS, rhs)
			if err != nil {
				t.Fatal(err)
			}
			return []*core.Array{x}
		}},
		{"RRQR", func() []*core.Array {
			q, r, perm, rank, err := RRQR(rrA)
			if err != nil {
				t.Fatal(err)
			}
			pf := make([]float64, len(perm)+1)
			for i, p := range perm {
				pf[i] = float64(p)
			}
			pf[len(perm)] = float64(rank)
			return []*core.Array{q, r, mustFromFloats(t, pf, len(pf), 1)}
		}},
		{"SolveRRQR", func() []*core.Array {
			x, err := SolveRRQR(rrA, rrRHS)
			if err != nil {
				t.Fatal(err)
			}
			return []*core.Array{x}
		}},
		{"SVD", func() []*core.Array {
			u, s, vt, err := SVD(wideA)
			if err != nil {
				t.Fatal(err)
			}
			return []*core.Array{u, s, vt}
		}},
		{"Eigen", func() []*core.Array {
			v, q, err := Eigen(squareA)
			if err != nil {
				t.Fatal(err)
			}
			return []*core.Array{v, q}
		}},
		{"EigenComplex", func() []*core.Array {
			v, q, err := EigenComplex(hermA)
			if err != nil {
				t.Fatal(err)
			}
			return []*core.Array{v, q}
		}},
		{"SVDComplex", func() []*core.Array {
			u, s, vh, err := SVDComplex(complexA)
			if err != nil {
				t.Fatal(err)
			}
			return []*core.Array{u, s, vh}
		}},
	}

	serial := make([][]*core.Array, len(cases))
	for i, c := range cases {
		serial[i] = c.run()
	}
	engine.SetNumWorkers(0) // the machine's full width
	for i, c := range cases {
		got := c.run()
		for k := range got {
			if got[k].Len() != serial[i][k].Len() {
				t.Fatalf("%s: result %d length %d under the full crew, %d serial",
					c.name, k, got[k].Len(), serial[i][k].Len())
			}
			if !rawBitsEqual(got[k], serial[i][k]) {
				t.Fatalf("%s: result %d differs bitwise between the serial and parallel crew", c.name, k)
			}
		}
	}
}

// rawBitsEqual compares two arrays' payloads bit for bit, complex
// payloads included.
func rawBitsEqual(a, b *core.Array) bool {
	if a.Dtype() == core.Complex || b.Dtype() == core.Complex {
		ac, bc := a.RawComplexes(), b.RawComplexes()
		if len(ac) != len(bc) {
			return false
		}
		for i := range ac {
			if ac[i] != bc[i] {
				return false
			}
		}
		return true
	}
	af, bf := a.RawFloats(), b.RawFloats()
	if len(af) != len(bf) {
		return false
	}
	for i := range af {
		if af[i] != bf[i] {
			return false
		}
	}
	return true
}

// BenchmarkLeastSquaresTall measures the solve at the shapes the
// reflector route exists for: many more rows than columns, where forming
// Q would dominate everything else.
func BenchmarkLeastSquaresTall(b *testing.B) {
	for _, c := range []struct{ m, n int }{{512, 32}, {2048, 64}, {8192, 16}} {
		a := make([]float64, c.m*c.n)
		s := uint64(20260920)
		for i := range a {
			s = s*6364136223846793005 + 1442695040888963407
			a[i] = float64((s>>40)%9+1) * 0.5
		}
		bm := make([]float64, c.m)
		for i := range bm {
			bm[i] = float64(i%11) - 5
		}
		am, err := core.FromFloats(a, c.m, c.n)
		if err != nil {
			b.Fatal(err)
		}
		bv, err := core.FromFloats(bm, c.m, 1)
		if err != nil {
			b.Fatal(err)
		}
		b.Run(fmt.Sprintf("%dx%d", c.m, c.n), func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				if _, err := LeastSquares(am, bv); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}