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

package linalg

import (
	"fmt"
	"math"
	"slices"
	"strings"
	"testing"

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

// Regression pins: a depth-first seed that ignored the live row
// permutation in the sparse LU, overflow states that slipped past the
// pivot and downdate guards, non-finite entries that sailed through
// the ILU intake and the symmetry screens, a minimum-degree
// absorption that handed a vertex its own index back, a zero
// hypotenuse in the complex SVD bulge chase, and a MatrixLog screen
// that refused positive definite matrices with small positive
// eigenvalues and a false "non-positive" verdict.

// TestSparseLUPermutedSeedFactorisation pins the seed fix: once an
// earlier column has swapped rows, the depth-first search over column
// k must enter through the rows' current positions in the factor, not
// through the labels they were stored under. The 4×4 matrix below
// swaps original row 2 to position 0 and row 3 to position 1 before
// column 1 is eliminated, so column 1's stored rows 2 and 3 now sit
// at positions 0 and 3.
func TestSparseLUPermutedSeedFactorisation(t *testing.T) {
	trips := [][3]float64{
		{0, 0, 1}, {0, 2, 1}, {0, 3, 1},
		{1, 0, 2}, {1, 2, 1}, {1, 3, 1},
		{2, 0, 5}, {2, 1, 7}, {2, 2, 8}, {2, 3, 1},
		{3, 1, 4}, {3, 2, 1}, {3, 3, 9},
	}
	idx := make([]int64, 0, 2*len(trips))
	vals := make([]float64, 0, len(trips))
	for _, e := range trips {
		idx = append(idx, int64(e[0]), int64(e[1]))
		vals = append(vals, e[2])
	}
	coo, err := core.NewSparseCOO(mustInts(t, idx, len(trips), 2),
		floatsToArray(vals, []int{len(trips)}), []int{4, 4})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	f, err := NewSparseLU(coo)
	if err != nil {
		t.Fatalf("NewSparseLU: %v", err)
	}
	if f.piv[0] != 2 || f.piv[1] != 3 {
		t.Fatalf("permutation %v, want it to start [2 3 ...]", f.piv)
	}
	// The swapped row 2 (now position 0) carries A[2,1] = 7 into
	// U row 0: the seed walk must have found its column.
	found := false
	for p, c := range f.rowCols[0] {
		if c == 1 {
			found = true
			if math.Abs(f.rowVals[0][p]-7) > 1e-12 {
				t.Fatalf("U[0,1] = %g, want 7", f.rowVals[0][p])
			}
		}
	}
	if !found {
		t.Fatal("U row 0 has no column-1 entry; the seed walk ignored the permutation")
	}
	// The factor must answer what the dense elimination answers.
	x, err := f.Solve(mustFloats(t, []float64{8, 9, 47, 47}))
	if err != nil {
		t.Fatalf("Solve: %v", err)
	}
	want := []float64{1, 2, 3, 4}
	for i := range 4 {
		if math.Abs(x.FloatAt(i)-want[i]) > 1e-12 {
			t.Fatalf("solve[%d] = %.15g, want %.15g", i, x.FloatAt(i), want[i])
		}
	}
}

// TestSparseLUOverflowPivotRefused pins the pivot guard: an
// elimination that squares the range (the update 1e308 − 1·1e308
// reads −2e308, which rounds to −Inf) used to slide past the NaN and
// zero tests and store an infinite pivot without a word.
func TestSparseLUOverflowPivotRefused(t *testing.T) {
	coo := cooFrom(t, []int64{0, 0, 0, 1, 1, 0, 1, 1},
		[]float64{1e308, 1e308, 1e308, -1e308}, []int{2, 2})
	f, err := NewSparseLU(coo)
	if err == nil {
		t.Fatalf("NewSparseLU stored an infinite pivot (diag = [%g, %g]) with no error", f.diag[0], f.diag[1])
	}
	if !strings.Contains(err.Error(), "overflow") {
		t.Fatalf("error = %v, want an overflow refusal", err)
	}
}

// TestSparseCholeskyOverflowPivotRefused pins the same guard on the
// Cholesky side: L[1,0] = 1e100/1e-100 = 1e200 and the pivot update
// squares it to 1e400, which overflows. The matrix is indefinite, but
// the rounded arithmetic cannot honestly deliver a
// positive-definiteness verdict, and the report is the overflow.
func TestSparseCholeskyOverflowPivotRefused(t *testing.T) {
	coo := cooFrom(t, []int64{0, 0, 0, 1, 1, 0, 1, 1},
		[]float64{1e-200, 1e100, 1e100, 1e200}, []int{2, 2})
	_, err := NewSparseCholesky(coo, SparseOrderingNatural)
	if err == nil {
		t.Fatal("NewSparseCholesky stored an infinite pivot with no error")
	}
	if !strings.Contains(err.Error(), "overflow") {
		t.Fatalf("error = %v, want an overflow refusal", err)
	}
}

// TestSparseILURefusesNonFinite pins the intake the complete sparse
// factorisations already have: a NaN or infinite entry is refused
// before the elimination turns it into a poisoned preconditioner.
func TestSparseILURefusesNonFinite(t *testing.T) {
	idx := []int64{0, 0, 0, 1, 1, 0, 1, 1, 1, 2, 2, 1, 2, 2}
	for name, bad := range map[string]float64{"NaN": math.NaN(), "Inf": math.Inf(1)} {
		t.Run(name, func(t *testing.T) {
			coo := cooFrom(t, idx, []float64{4, bad, -1, 4, -1, -1, 4}, []int{3, 3})
			if _, err := NewSparseILU(coo); err == nil {
				t.Fatal("NewSparseILU accepted a non-finite entry")
			} else if !strings.Contains(err.Error(), "not finite") {
				t.Fatalf("error = %v, want a non-finite refusal", err)
			}
		})
	}
	t.Run("finite", func(t *testing.T) {
		coo := cooFrom(t, idx, []float64{4, 1, -1, 4, -1, -1, 4}, []int{3, 3})
		if _, err := NewSparseILU(coo); err != nil {
			t.Fatalf("NewSparseILU refused a finite matrix: %v", err)
		}
	})
}

// TestMinimumDegreeAbsorptionDropsSelf pins the absorption rule:
// eliminating a vertex merges its adjacency into each surviving
// neighbour as adj(j) ∪ adj(p) \ {j}. The plain union hands j its own
// index back (it sits on p's list), which inflates the degree the
// selection scan reads and warps the order.
func TestMinimumDegreeAbsorptionDropsSelf(t *testing.T) {
	// Path 0 - 1 - 2: adj = [[1], [0, 2], [1]].
	coo := cooFrom(t, []int64{0, 1, 1, 0, 1, 2, 2, 1, 0, 0, 1, 1, 2, 2},
		[]float64{-1, -1, -1, -1, 4, 4, 4}, []int{3, 3})
	csc, err := CSCFromCOO(coo)
	if err != nil {
		t.Fatalf("CSCFromCOO: %v", err)
	}
	adj, err := symmetrisedAdjacency(csc)
	if err != nil {
		t.Fatalf("symmetrisedAdjacency: %v", err)
	}
	eliminated := []bool{true, false, false}
	degree := []int{1, 2, 1}
	frontier := &degreeFrontier{}
	absorbElement(adj, degree, eliminated, 0, frontier, &intArena{})
	for _, u := range adj[1] {
		if u == 1 {
			t.Fatalf("adj[1] = %v lists 1 itself after absorbing 0's element", adj[1])
		}
	}
	if degree[1] != 1 {
		t.Fatalf("vertex 1 keeps degree %d after the absorption, want 1", degree[1])
	}
	deg := 0
	for _, u := range adj[1] {
		if !eliminated[u] {
			deg++
		}
	}
	if deg != 1 {
		t.Fatalf("vertex 1 reads degree %d after the absorption, want 1", deg)
	}
	if !slices.IsSorted(adj[1]) {
		t.Fatalf("adj[1] = %v is not the sorted form the union expects", adj[1])
	}
	// With honest degrees, vertices 1 and 2 tie at degree 1 and the
	// tie breaks to the smaller index.
	order, err := minimumDegree(csc)
	if err != nil {
		t.Fatalf("minimumDegree: %v", err)
	}
	if !slices.Equal(order, []int{0, 1, 2}) {
		t.Fatalf("minimum degree order %v, want [0 1 2]", order)
	}
}

// TestSVDComplexRankOneRectangular pins the bulge chase's zero step:
// when a deflated diagonal meets a deflated bulge the rotation is the
// identity, where d[k]/0 raised NaNs that spilled through every
// factor. The rank-1 rectangular matrices below reach exactly that
// step.
func TestSVDComplexRankOneRectangular(t *testing.T) {
	t.Run("2x3explicit", func(t *testing.T) {
		raw := []complex128{1, 2, 3, 2, 4, 6}
		a := mustComplex(t, raw, 2, 3)
		u, sigma, vh, err := SVDComplex(a)
		if err != nil {
			t.Fatalf("SVDComplex: %v", err)
		}
		want := math.Sqrt(70.0)
		if rel := math.Abs(sigma.FloatAt(0)-want) / want; rel > 1e-12 {
			t.Fatalf("sigma[0] = %.15g, want %.15g (relative error %.3g)", sigma.FloatAt(0), want, rel)
		}
		if rel := sigma.FloatAt(1) / want; rel > 1e-12 {
			t.Fatalf("sigma[1] = %.15g, want 0 (relative %.3g)", sigma.FloatAt(1), rel)
		}
		checkUnitaryRows(t, u, sigma, vh, raw, 2, 3)
	})
	for _, shape := range [][2]int{{2, 3}, {2, 4}, {2, 5}, {3, 2}, {3, 4}, {4, 2}, {5, 2}} {
		t.Run(fmt.Sprintf("%dx%d", shape[0], shape[1]), func(t *testing.T) {
			m, n := shape[0], shape[1]
			a := make([]complex128, m*n)
			for i := range m {
				for j := range n {
					ui := complex(float64(i+1), float64(i))
					vj := complex(float64(-j), float64(j+1))
					a[i*n+j] = ui * vj
				}
			}
			aa, err := core.FromComplexes(a, m, n)
			if err != nil {
				t.Fatalf("FromComplexes: %v", err)
			}
			u, sigma, vh, err := SVDComplex(aa)
			if err != nil {
				t.Fatalf("SVDComplex: %v", err)
			}
			// A rank-1 matrix carries its whole Frobenius norm in the
			// leading singular value and nothing in the rest.
			fro := 0.0
			for _, z := range a {
				fro += real(z)*real(z) + imag(z)*imag(z)
			}
			fro = math.Sqrt(fro)
			if rel := math.Abs(sigma.FloatAt(0)-fro) / fro; rel > 1e-12 {
				t.Fatalf("sigma[0] = %.15g, want %.15g (relative error %.3g)", sigma.FloatAt(0), fro, rel)
			}
			if rel := sigma.FloatAt(1) / fro; rel > 1e-12 {
				t.Fatalf("sigma[1] = %.15g, want 0 (relative %.3g)", sigma.FloatAt(1), rel)
			}
			for k := 1; k < sigma.Len(); k++ {
				if math.IsNaN(sigma.FloatAt(k)) || math.IsInf(sigma.FloatAt(k), 0) || sigma.FloatAt(k) < 0 {
					t.Fatalf("sigma[%d] = %g is not a plain non-negative number", k, sigma.FloatAt(k))
				}
			}
			checkUnitaryRows(t, u, sigma, vh, a, m, n)
		})
	}
}

// checkUnitaryRows verifies the factors a rank-decomposed SVD must
// return: Uᴴ·U = I on the thin U, orthonormal rows of Vᴴ, and the
// reconstruction U·diag(σ)·Vᴴ back to a.
func checkUnitaryRows(t *testing.T, u, sigma, vh *core.Array, a []complex128, m, n int) {
	t.Helper()
	r := min(m, n)
	uc := u.RawComplexes()
	ucols := u.Shape()[1]
	urows := u.Shape()[0]
	for j := range r {
		for i := range r {
			s := complex(0, 0)
			for k := range urows {
				s += cmplxConjTest(uc[k*ucols+i]) * uc[k*ucols+j]
			}
			want := 0.0
			if i == j {
				want = 1
			}
			if d := cmplxAbsTest(s - complex(want, 0)); d > 1e-12 {
				t.Fatalf("(Uᴴ·U)[%d,%d] = %v, want %v", i, j, s, want)
			}
		}
	}
	vc := vh.RawComplexes()
	vcols := vh.Shape()[1]
	for j := range r {
		for i := range r {
			s := complex(0, 0)
			for k := range vcols {
				s += vc[i*vcols+k] * cmplxConjTest(vc[j*vcols+k])
			}
			want := 0.0
			if i == j {
				want = 1
			}
			if d := cmplxAbsTest(s - complex(want, 0)); d > 1e-12 {
				t.Fatalf("(Vᴴ·(Vᴴ)ᴴ)[%d,%d] = %v, want %v", i, j, s, want)
			}
		}
	}
	worst := 0.0
	for i := range m {
		for j := range n {
			s := complex(0, 0)
			for k := range r {
				s += uc[i*ucols+k] * complex(sigma.FloatAt(k), 0) * vc[k*vcols+j]
			}
			if d := cmplxAbsTest(s - a[i*n+j]); d > worst {
				worst = d
			}
		}
	}
	if worst > 1e-12*sigma.FloatAt(0) {
		t.Fatalf("reconstruction error %g exceeds %g", worst, 1e-12*sigma.FloatAt(0))
	}
}

func cmplxConjTest(z complex128) complex128 { return complex(real(z), -imag(z)) }

func cmplxAbsTest(z complex128) float64 { return math.Hypot(real(z), imag(z)) }

// TestMatrixLogSymmetricSmallSpectrum pins the symmetric route's
// negativity floor: a positive eigenvalue is logged whatever its size
// against the spectrum, where the old 1e-10 relative screen refused
// fine positive eigenvalues with a false "non-positive" verdict.
func TestMatrixLogSymmetricSmallSpectrum(t *testing.T) {
	tiny := mustFloats(t, []float64{1, 0, 0, 1e-12}, 2, 2)
	lg, err := MatrixLog(tiny)
	if err != nil {
		t.Fatalf("MatrixLog(diag(1, 1e-12)): %v", err)
	}
	for i, w := range []float64{0, math.Log(1e-12)} {
		if math.Abs(lg.FloatAt(i*2+i)-w) > 1e-9 {
			t.Fatalf("ln A[%d][%d] = %.16g, want %.16g", i, i, lg.FloatAt(i*2+i), w)
		}
	}
	// The smaller eigenvalue is 1e-16 of the larger, far below any
	// relative screen, and still logs honestly.
	big := mustFloats(t, []float64{1e20, 0, 0, 1e4}, 2, 2)
	lg, err = MatrixLog(big)
	if err != nil {
		t.Fatalf("MatrixLog(diag(1e20, 1e4)): %v", err)
	}
	for i, w := range []float64{math.Log(1e20), math.Log(1e4)} {
		if rel := math.Abs(lg.FloatAt(i*2+i)-w) / math.Abs(w); rel > 1e-12 {
			t.Fatalf("ln A[%d][%d] = %.16g, want %.16g", i, i, lg.FloatAt(i*2+i), w)
		}
	}
	for _, i := range []int{1, 2} {
		if math.Abs(lg.FloatAt(i)) > 1e-6 {
			t.Fatalf("ln A off-diagonal [%d] = %g, want 0", i, lg.FloatAt(i))
		}
	}
	// A genuinely negative eigenvalue is still refused.
	neg := mustFloats(t, []float64{0, 1, 1, 0}, 2, 2)
	if _, err := MatrixLog(neg); err == nil {
		t.Fatal("MatrixLog([[0,1],[1,0]]): want an error for the negative eigenvalue")
	}
	// MatrixSqrt keeps its own screen: positive spectra root, negative
	// spectra refuse, nothing moved.
	sq, err := MatrixSqrt(tiny)
	if err != nil {
		t.Fatalf("MatrixSqrt(diag(1, 1e-12)): %v", err)
	}
	for i, w := range []float64{1, 1e-6} {
		if math.Abs(sq.FloatAt(i*2+i)-w) > 1e-12 {
			t.Fatalf("√A[%d][%d] = %.16g, want %.16g", i, i, sq.FloatAt(i*2+i), w)
		}
	}
	sq, err = MatrixSqrt(big)
	if err != nil {
		t.Fatalf("MatrixSqrt(diag(1e20, 1e4)): %v", err)
	}
	for i, w := range []float64{1e10, 100} {
		if rel := math.Abs(sq.FloatAt(i*2+i)-w) / w; rel > 1e-12 {
			t.Fatalf("√A[%d][%d] = %.16g, want %.16g", i, i, sq.FloatAt(i*2+i), w)
		}
	}
	if _, err := MatrixSqrt(neg); err == nil {
		t.Fatal("MatrixSqrt of an indefinite matrix: want an error")
	}
}

// TestCholeskyDowndateOverflowRefused pins the downdate's range test:
// with a diagonal at 1e160 the squared terms overflow, and r2 reads
// NaN (Inf − Inf) or +Inf depending on the vector, both of which used
// to slide past the cone test into a silent NaN or infinite factor.
func TestCholeskyDowndateOverflowRefused(t *testing.T) {
	l := mustFloats(t, []float64{1e160}, 1, 1)
	for name, xv := range map[string]float64{"NaN": 1e160, "Inf": 1} {
		t.Run(name, func(t *testing.T) {
			out, err := CholeskyDowndate(l, mustFloats(t, []float64{xv}))
			if err == nil {
				t.Fatalf("downdate returned a factor with diagonal %g and no error", out.FloatAt(0))
			}
			if !strings.Contains(err.Error(), "overflow") {
				t.Fatalf("error = %v, want an overflow refusal", err)
			}
		})
	}
	// A rank-one update and a downdate inside the range are untouched.
	small := mustFloats(t, []float64{4}, 1, 1)
	up, err := CholeskyUpdate(small, mustFloats(t, []float64{1}))
	if err != nil {
		t.Fatalf("CholeskyUpdate: %v", err)
	}
	if math.Abs(up.FloatAt(0)-math.Sqrt(17)) > 1e-12 {
		t.Fatalf("update diagonal %g, want %.15g", up.FloatAt(0), math.Sqrt(17))
	}
	down, err := CholeskyDowndate(small, mustFloats(t, []float64{1}))
	if err != nil {
		t.Fatalf("CholeskyDowndate: %v", err)
	}
	if math.Abs(down.FloatAt(0)-math.Sqrt(15)) > 1e-12 {
		t.Fatalf("downdate diagonal %g, want %.15g", down.FloatAt(0), math.Sqrt(15))
	}
}

// TestSparseSymmetryChecksRefuseNonFinite pins the symmetry and
// Hermitian screens: a NaN compares unequal to everything, so the
// mirror test alone waved non-finite entries through as symmetric.
func TestSparseSymmetryChecksRefuseNonFinite(t *testing.T) {
	idx := []int64{0, 0, 0, 1, 1, 0, 1, 1}
	t.Run("symmetricNaN", func(t *testing.T) {
		coo := cooFrom(t, idx, []float64{4, math.NaN(), math.NaN(), 4}, []int{2, 2})
		c, err := cooToCSR(coo, "test")
		if err != nil {
			t.Fatalf("cooToCSR: %v", err)
		}
		if err := c.checkSymmetric("SpSolve"); err == nil {
			t.Fatal("checkSymmetric accepted a NaN entry")
		} else if !strings.Contains(err.Error(), "not finite") {
			t.Fatalf("error = %v, want a non-finite refusal", err)
		}
	})
	t.Run("symmetricInf", func(t *testing.T) {
		coo := cooFrom(t, idx, []float64{4, math.Inf(1), math.Inf(1), 4}, []int{2, 2})
		c, err := cooToCSR(coo, "test")
		if err != nil {
			t.Fatalf("cooToCSR: %v", err)
		}
		if err := c.checkSymmetric("SpSolve"); err == nil {
			t.Fatal("checkSymmetric accepted an infinite entry")
		} else if !strings.Contains(err.Error(), "not finite") {
			t.Fatalf("error = %v, want a non-finite refusal", err)
		}
	})
	t.Run("hermitianNaN", func(t *testing.T) {
		coo := cooComplexFrom(t, idx, []complex128{4, cmplxNaN(), cmplxNaN(), 4}, []int{2, 2})
		c, err := cooToComplexCSR(coo, "test")
		if err != nil {
			t.Fatalf("cooToComplexCSR: %v", err)
		}
		if err := c.checkHermitian("SpSolveComplexCG"); err == nil {
			t.Fatal("checkHermitian accepted a NaN entry")
		} else if !strings.Contains(err.Error(), "not finite") {
			t.Fatalf("error = %v, want a non-finite refusal", err)
		}
	})
	t.Run("hermitianInf", func(t *testing.T) {
		inf := complex(math.Inf(1), 0)
		coo := cooComplexFrom(t, idx, []complex128{4, inf, inf, 4}, []int{2, 2})
		c, err := cooToComplexCSR(coo, "test")
		if err != nil {
			t.Fatalf("cooToComplexCSR: %v", err)
		}
		if err := c.checkHermitian("SpSolveComplexCG"); err == nil {
			t.Fatal("checkHermitian accepted an infinite entry")
		} else if !strings.Contains(err.Error(), "not finite") {
			t.Fatalf("error = %v, want a non-finite refusal", err)
		}
	})
	t.Run("legalInputs", func(t *testing.T) {
		coo := cooFrom(t, idx, []float64{4, 1, 1, 4}, []int{2, 2})
		c, err := cooToCSR(coo, "test")
		if err != nil {
			t.Fatalf("cooToCSR: %v", err)
		}
		if err := c.checkSymmetric("SpSolve"); err != nil {
			t.Fatalf("checkSymmetric refused a finite symmetric matrix: %v", err)
		}
		hcoo := cooComplexFrom(t, idx, []complex128{4, 1 + 2i, 1 - 2i, 4}, []int{2, 2})
		hc, err := cooToComplexCSR(hcoo, "test")
		if err != nil {
			t.Fatalf("cooToComplexCSR: %v", err)
		}
		if err := hc.checkHermitian("SpSolveComplexCG"); err != nil {
			t.Fatalf("checkHermitian refused a finite Hermitian matrix: %v", err)
		}
	})
	t.Run("publicEntries", func(t *testing.T) {
		sp := cooFrom(t, idx, []float64{4, math.NaN(), math.NaN(), 4}, []int{2, 2})
		if _, err := SpSolve(sp, mustFloats(t, []float64{1, 1}), 0, 0); err == nil {
			t.Fatal("SpSolve accepted a NaN entry")
		}
		hsp := cooComplexFrom(t, idx, []complex128{4, cmplxNaN(), cmplxNaN(), 4}, []int{2, 2})
		b := mustComplex(t, []complex128{1, 1}, 2)
		if _, err := SpSolveComplexCG(hsp, b, 0, 0); err == nil {
			t.Fatal("SpSolveComplexCG accepted a NaN entry")
		}
	})
}

// cooComplexFrom builds a complex-valued sparse COO, failing the test
// on a bad shape.
func cooComplexFrom(t *testing.T, idx []int64, vals []complex128, shape []int) *core.SparseCOO {
	t.Helper()
	i, err := core.FromInts(idx, len(idx)/len(shape), len(shape))
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	v, err := core.FromComplexes(vals, len(vals))
	if err != nil {
		t.Fatalf("FromComplexes: %v", err)
	}
	sp, err := core.NewSparseCOO(i, v, shape)
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	return sp
}

func cmplxNaN() complex128 { return complex(math.NaN(), math.NaN()) }