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

package linalg

import (
	"math"
	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
	"testing"
)

// sparseFromDense builds a symmetric SparseCOO from a dense matrix,
// failing the test if the input is not symmetric enough to qualify.
func sparseFromDense(t *testing.T, vals []float64, n int) *core.SparseCOO {
	t.Helper()
	dense, err := core.FromFloats(vals, n, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	sp, err := core.SparseFrom(dense)
	if err != nil {
		t.Fatalf("SparseFrom: %v", err)
	}
	return sp
}

// residual norms ‖A·v − λ·v‖ for one eigenpair, the direct check that
// a returned pair is actually an eigenpair of the original matrix.
// vec holds eigenvectors as columns of an (n, k) array.
func residual(t *testing.T, sp *core.SparseCOO, vec *core.Array, lam float64, n, k, j int) float64 {
	t.Helper()
	x := make([]float64, n)
	for i := range n {
		x[i] = vec.FloatAt(i*k + j)
	}
	dense, err := sp.Dense()
	if err != nil {
		t.Fatalf("Dense: %v", err)
	}
	r := 0.0
	for i := range n {
		sum := 0.0
		for p := range n {
			sum += dense.FloatAt(i*n+p) * x[p]
		}
		r += (sum - lam*x[i]) * (sum - lam*x[i])
	}
	return math.Sqrt(r)
}

// TestSpEigenMatchesDense compares the sparse solver against the dense
// `Eigen` on the same matrix: the k largest-magnitude eigenvalues must
// agree, which is the property that matters for a partial solver.
func TestSpEigenMatchesDense(t *testing.T) {
	cases := []struct {
		name string
		vals []float64
		n    int
		k    int
	}{
		{
			name: "diagonal",
			vals: []float64{
				5, 0, 0, 0,
				0, -3, 0, 0,
				0, 0, 2, 0,
				0, 0, 0, -7,
			},
			n: 4, k: 2,
		},
		{
			name: "tridiagonal",
			vals: []float64{
				4, 1, 0, 0,
				1, 4, 1, 0,
				0, 1, 4, 1,
				0, 0, 1, 4,
			},
			n: 4, k: 2,
		},
		{
			name: "full_symmetric",
			vals: []float64{
				2, -1, 0.5, 0.3,
				-1, 3, 0.2, -0.4,
				0.5, 0.2, 1, 0.6,
				0.3, -0.4, 0.6, 2.5,
			},
			n: 4, k: 3,
		},
		{
			name: "repeated_spectrum",
			vals: []float64{
				1, 0, 0, 0,
				0, 1, 0, 0,
				0, 0, -2, 0,
				0, 0, 0, -2,
			},
			n: 4, k: 2,
		},
	}
	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			sp := sparseFromDense(t, tt.vals, tt.n)
			gotVals, gotVecs, err := SpEigen(sp, tt.k, nil)
			if err != nil {
				t.Fatalf("SpEigen: %v", err)
			}
			if gotVals.Len() != tt.k {
				t.Fatalf("values len %d, want %d", gotVals.Len(), tt.k)
			}
			if gotVecs.Shape()[0] != tt.n || gotVecs.Shape()[1] != tt.k {
				t.Fatalf("vectors shape %s, want [%d %d]", base.ShapeText(gotVecs.Shape()), tt.n, tt.k)
			}

			dense, err := core.FromFloats(tt.vals, tt.n, tt.n)
			if err != nil {
				t.Fatalf("FromFloats: %v", err)
			}
			refVals, _, err := Eigen(dense)
			if err != nil {
				t.Fatalf("Eigen: %v", err)
			}
			// Reference extremes by descending magnitude.
			ref := make([]float64, tt.n)
			for i := range tt.n {
				ref[i] = refVals.FloatAt(i)
			}
			for i := range tt.k {
				for j := i + 1; j < tt.n; j++ {
					if math.Abs(ref[j]) > math.Abs(ref[i]) {
						ref[i], ref[j] = ref[j], ref[i]
					}
				}
			}
			for i := range tt.k {
				g, w := gotVals.FloatAt(i), ref[i]
				if math.Abs(g-w) > 1e-8*(1+math.Abs(w)) {
					t.Fatalf("value[%d] = %.12g, want %.12g", i, g, w)
				}
			}

			// Each returned pair must be an eigenpair of A itself.
			for j := range tt.k {
				r := residual(t, sp, gotVecs, gotVals.FloatAt(j), tt.n, tt.k, j)
				if r > 1e-8 {
					t.Fatalf("eigenpair %d: residual ‖Av-λv‖ = %.3g, want <= 1e-8", j, r)
				}
			}
		})
	}
}

// TestSpEigenFullSpectrum checks that k = n recovers every eigenvalue,
// matching the dense solver's complete output.
func TestSpEigenFullSpectrum(t *testing.T) {
	const n = 5
	vals := []float64{
		4, 1, 0, 0, 0,
		1, 3, 2, 0, 0,
		0, 2, 1, -1, 0,
		0, 0, -1, 2, 0.5,
		0, 0, 0, 0.5, 5,
	}
	sp := sparseFromDense(t, vals, n)
	gotVals, _, err := SpEigen(sp, n, nil)
	if err != nil {
		t.Fatalf("SpEigen: %v", err)
	}
	dense, err := core.FromFloats(vals, n, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	refVals, _, err := Eigen(dense)
	if err != nil {
		t.Fatalf("Eigen: %v", err)
	}
	got := make([]float64, n)
	for i := range n {
		got[i] = gotVals.FloatAt(i)
	}
	for i := range n {
		for j := i + 1; j < n; j++ {
			if math.Abs(got[j]) > math.Abs(got[i]) {
				got[i], got[j] = got[j], got[i]
			}
		}
	}
	ref := make([]float64, n)
	for i := range n {
		ref[i] = refVals.FloatAt(i)
	}
	for i := range n {
		for j := i + 1; j < n; j++ {
			if math.Abs(ref[j]) > math.Abs(ref[i]) {
				ref[i], ref[j] = ref[j], ref[i]
			}
		}
	}
	for i := range n {
		if math.Abs(got[i]-ref[i]) > 1e-8*(1+math.Abs(ref[i])) {
			t.Fatalf("value[%d] = %.12g, want %.12g", i, got[i], ref[i])
		}
	}
}

// TestSpEigenLargerMatrix exercises the solver where it earns its
// keep: a sparse matrix far larger than the eigenpair count, so the
// answer comes from a partial decomposition rather than a full one.
// The spectrum is deliberately spread out with well-separated
// extremes, which is the regime Lanczos converges in; a clustered
// spectrum needs a far larger iteration budget than a partial solve
// normally gets.
func TestSpEigenLargerMatrix(t *testing.T) {
	const n = 24
	// A diagonally dominant tridiagonal matrix: diagonal spaced from
	// +n down to +1, off-diagonals 1. The diagonal is kept positive so
	// the spectrum is not symmetric about zero: a spectrum with ±pairs
	// of equal magnitude would leave the descending-magnitude ordering
	// with ties, and the test could not compare position by position.
	// The extremes sit well clear of the rest, so the top Ritz values
	// are accurate.
	vals := make([]float64, n*n)
	for i := range n {
		vals[i*n+i] = float64(n - i)
		if i+1 < n {
			vals[i*n+i+1] = 1
			vals[(i+1)*n+i] = 1
		}
	}
	sp := sparseFromDense(t, vals, n)
	const k = 3
	gotVals, gotVecs, err := SpEigen(sp, k, nil)
	if err != nil {
		t.Fatalf("SpEigen: %v", err)
	}
	dense, err := core.FromFloats(vals, n, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	refVals, _, err := Eigen(dense)
	if err != nil {
		t.Fatalf("Eigen: %v", err)
	}
	ref := make([]float64, n)
	for i := range n {
		ref[i] = refVals.FloatAt(i)
	}
	for i := range k {
		for j := i + 1; j < n; j++ {
			if math.Abs(ref[j]) > math.Abs(ref[i]) {
				ref[i], ref[j] = ref[j], ref[i]
			}
		}
	}
	for i := range k {
		g, w := gotVals.FloatAt(i), ref[i]
		if math.Abs(g-w) > 1e-8*(1+math.Abs(w)) {
			t.Fatalf("value[%d] = %.12g, want %.12g", i, g, w)
		}
	}
	for j := range k {
		// A partial solve delivers ‖Av-λv‖ ≈ β_m·|y_m|, the last
		// coupling times the tail of the tridiagonal eigenvector,
		// rather than machine epsilon. The eigenvalues above are
		// accurate to 1e-8 because the extremes are well separated;
		// the vectors carry the looser bound.
		r := residual(t, sp, gotVecs, gotVals.FloatAt(j), n, k, j)
		if r > 1e-5 {
			t.Fatalf("eigenpair %d: residual ‖Av-λv‖ = %.3g, want <= 1e-5", j, r)
		}
	}
}

// TestSpEigenPartialDecomposition exercises the path the sparse
// eigensolver exists for: n far larger than the eigenpair count, so
// the iteration budget stops well short of the dimension and the
// answer is a truncated decomposition rather than an exact one. The
// smaller tests all take n ≤ k+40, which decomposes fully and never
// reaches this branch.
func TestSpEigenPartialDecomposition(t *testing.T) {
	const n = 120
	// A diagonally dominant tridiagonal matrix whose diagonal falls off
	// like 1000/(i+1), so the top few eigenvalues are widely separated
	// relative to their magnitude. That separation is what lets a
	// truncated Lanczos converge on the wanted extremes inside the
	// budget; a tightly clustered spectrum (a near-constant diagonal,
	// say) converges far slower and would make the assertion about the
	// iteration budget rather than about the solver. The off-diagonal
	// of 1 is negligible against the leading diagonal entries, so the
	// spectrum stays close to the diagonal.
	vals := make([]float64, n*n)
	for i := range n {
		vals[i*n+i] = 1000 / float64(i+1)
		if i+1 < n {
			vals[i*n+i+1] = 1
			vals[(i+1)*n+i] = 1
		}
	}
	sp := sparseFromDense(t, vals, n)
	const k = 4
	gotVals, gotVecs, err := SpEigen(sp, k, nil)
	if err != nil {
		t.Fatalf("SpEigen: %v", err)
	}
	dense, err := core.FromFloats(vals, n, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	refVals, _, err := Eigen(dense)
	if err != nil {
		t.Fatalf("Eigen: %v", err)
	}
	ref := make([]float64, n)
	for i := range n {
		ref[i] = refVals.FloatAt(i)
	}
	for i := range k {
		for j := i + 1; j < n; j++ {
			if math.Abs(ref[j]) > math.Abs(ref[i]) {
				ref[i], ref[j] = ref[j], ref[i]
			}
		}
	}
	// The budget stops at k+40 = 44 of n = 120 steps, and the diagonal
	// is tightly clustered near the top, so the extremes are converged
	// to a relative few times 1e-7 rather than exactly. A tolerance
	// tighter than that would be asserting an exactness a truncated
	// decomposition does not claim.
	for i := range k {
		g, w := gotVals.FloatAt(i), ref[i]
		if math.Abs(g-w) > 1e-5*(1+math.Abs(w)) {
			t.Fatalf("value[%d] = %.12g, want %.12g", i, g, w)
		}
	}
	for j := range k {
		r := residual(t, sp, gotVecs, gotVals.FloatAt(j), n, k, j)
		if r > 1e-4 {
			t.Fatalf("eigenpair %d: residual ‖Av-λv‖ = %.3g, want <= 1e-4", j, r)
		}
	}
}

// TestSpEigenSumsDuplicates pins the COO contract: a coordinate listed
// twice contributes the sum of both values, not whichever comes last.
func TestSpEigenSumsDuplicates(t *testing.T) {
	// Build [[3,0],[0,1]] as (0,0)=1 plus (0,0)=2, plus the lower entry.
	idx, err := core.FromInts([]int64{0, 0, 0, 0, 1, 1}, 3, 2)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	vals, err := core.FromFloats([]float64{1, 2, 1}, 3)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	sp, err := core.NewSparseCOO(idx, vals, []int{2, 2})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	gotVals, _, err := SpEigen(sp, 2, nil)
	if err != nil {
		t.Fatalf("SpEigen: %v", err)
	}
	if math.Abs(gotVals.FloatAt(0)-3) > 1e-12 {
		t.Fatalf("largest value = %.12g, want 3 (duplicates must sum)", gotVals.FloatAt(0))
	}
	if math.Abs(gotVals.FloatAt(1)-1) > 1e-12 {
		t.Fatalf("second value = %.12g, want 1", gotVals.FloatAt(1))
	}
}

// TestSpEigenDeterminism checks that the unseeded call is reproducible
// and that an explicit seed is honoured, the contract callers rely on
// when they pin a start vector.
func TestSpEigenDeterminism(t *testing.T) {
	const n = 6
	vals := make([]float64, n*n)
	for i := range n {
		vals[i*n+i] = float64(i) - 3
		if i+1 < n {
			vals[i*n+i+1] = 0.5
			vals[(i+1)*n+i] = 0.5
		}
	}
	sp := sparseFromDense(t, vals, n)
	v1, u1, err := SpEigen(sp, 3, nil)
	if err != nil {
		t.Fatalf("SpEigen #1: %v", err)
	}
	v2, u2, err := SpEigen(sp, 3, nil)
	if err != nil {
		t.Fatalf("SpEigen #2: %v", err)
	}
	for i := range 3 {
		if v1.FloatAt(i) != v2.FloatAt(i) {
			t.Fatalf("unseeded run %d: value %.12g vs %.12g", i, v1.FloatAt(i), v2.FloatAt(i))
		}
	}
	for i := range u1.Len() {
		if u1.FloatAt(i) != u2.FloatAt(i) {
			t.Fatalf("unseeded run: vector element %d differs", i)
		}
	}
	// A different seed may pick a different start vector, but the
	// eigenvalues of a symmetric matrix do not depend on it.
	v3, _, err := SpEigen(sp, 3, core.NewGenerator(99))
	if err != nil {
		t.Fatalf("SpEigen seeded: %v", err)
	}
	for i := range 3 {
		if math.Abs(v1.FloatAt(i)-v3.FloatAt(i)) > 1e-8 {
			t.Fatalf("seeded run %d: value %.12g vs %.12g", i, v1.FloatAt(i), v3.FloatAt(i))
		}
	}
}

// TestSpEigenRejectsInvalid pins the error contract for every input
// the solver cannot honestly answer.
func TestSpEigenRejectsInvalid(t *testing.T) {
	t.Run("complex_sparse", func(t *testing.T) {
		idx, err := core.FromInts([]int64{0, 0}, 1, 2)
		if err != nil {
			t.Fatalf("FromInts: %v", err)
		}
		vals, err := core.FromComplexes([]complex128{1}, 1)
		if err != nil {
			t.Fatalf("FromComplexes: %v", err)
		}
		sp, err := core.NewSparseCOO(idx, vals, []int{1, 1})
		if err != nil {
			t.Fatalf("NewSparseCOO: %v", err)
		}
		if _, _, err := SpEigen(sp, 1, nil); err == nil {
			t.Fatal("expected an error for a complex sparse matrix")
		}
	})
	t.Run("not_square", func(t *testing.T) {
		idx, err := core.FromInts([]int64{0, 0}, 1, 2)
		if err != nil {
			t.Fatalf("FromInts: %v", err)
		}
		vals, err := core.FromFloats([]float64{1}, 1)
		if err != nil {
			t.Fatalf("FromFloats: %v", err)
		}
		sp, err := core.NewSparseCOO(idx, vals, []int{1, 2})
		if err != nil {
			t.Fatalf("NewSparseCOO: %v", err)
		}
		if _, _, err := SpEigen(sp, 1, nil); err == nil {
			t.Fatal("expected an error for a non-square matrix")
		}
	})
	t.Run("zero_sized", func(t *testing.T) {
		sp := &core.SparseCOO{Shape: []int{0, 0}}
		idx, err := core.FromInts(nil, 0, 2)
		if err != nil {
			t.Fatalf("FromInts: %v", err)
		}
		vals, err := core.FromFloats(nil, 0)
		if err != nil {
			t.Fatalf("FromFloats: %v", err)
		}
		sp.Indices, sp.Values = idx, vals
		if _, _, err := SpEigen(sp, 1, nil); err == nil {
			t.Fatal("expected an error for a zero-sized matrix")
		}
	})
	t.Run("k_out_of_range", func(t *testing.T) {
		sp := sparseFromDense(t, []float64{1, 0, 0, 1}, 2)
		if _, _, err := SpEigen(sp, 0, nil); err == nil {
			t.Fatal("expected an error for k = 0")
		}
		if _, _, err := SpEigen(sp, 3, nil); err == nil {
			t.Fatal("expected an error for k > n")
		}
	})
	t.Run("asymmetric", func(t *testing.T) {
		// A is not symmetric: A[0,1] = 1 but A[1,0] = 2.
		sp := sparseFromDense(t, []float64{1, 1, 2, 1}, 2)
		if _, _, err := SpEigen(sp, 1, nil); err == nil {
			t.Fatal("expected an error for an asymmetric matrix")
		}
	})
	t.Run("index_out_of_range", func(t *testing.T) {
		idx, err := core.FromInts([]int64{0, 5}, 1, 2)
		if err != nil {
			t.Fatalf("FromInts: %v", err)
		}
		vals, err := core.FromFloats([]float64{1}, 1)
		if err != nil {
			t.Fatalf("FromFloats: %v", err)
		}
		sp, err := core.NewSparseCOO(idx, vals, []int{2, 2})
		if err != nil {
			t.Fatalf("NewSparseCOO: %v", err)
		}
		if _, _, err := SpEigen(sp, 1, nil); err == nil {
			t.Fatal("expected an error for an out-of-range index")
		}
	})
}