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

// spdTridiagonal builds a symmetric diagonally dominant tridiagonal
// matrix, which is positive-definite by the Gershgorin bound, so the
// conjugate gradient is guaranteed to apply. The diagonal is kept
// above the off-diagonal sum in absolute value.
func spdTridiagonal(n int, diag, off float64) []float64 {
	vals := make([]float64, n*n)
	for i := range n {
		vals[i*n+i] = diag
		if i+1 < n {
			vals[i*n+i+1] = off
			vals[(i+1)*n+i] = off
		}
	}
	return vals
}

// denseSolve solves A·x = b through the dense LU path, the reference
// the sparse solver is compared against.
func denseSolve(t *testing.T, vals []float64, n int, rhs []float64) []float64 {
	t.Helper()
	a, err := core.FromFloats(vals, n, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	b, err := core.FromFloats(rhs, n)
	if err != nil {
		t.Fatalf("FromFloats rhs: %v", err)
	}
	x, err := Solve(a, b)
	if err != nil {
		t.Fatalf("Solve: %v", err)
	}
	out := make([]float64, n)
	for i := range n {
		out[i] = x.FloatAt(i)
	}
	return out
}

// residualNorm norms ‖b − A·x‖₂, the direct check that a returned x
// solves the system rather than merely being returned.
func residualNorm(t *testing.T, vals []float64, x, rhs []float64, n int) float64 {
	t.Helper()
	sum := 0.0
	for i := range n {
		ax := 0.0
		for j := range n {
			ax += vals[i*n+j] * x[j]
		}
		d := rhs[i] - ax
		sum += d * d
	}
	return math.Sqrt(sum)
}

// TestSpSolveMatchesDense compares the sparse solver against the dense
// `Solve` on symmetric positive-definite systems, which is the
// contract the two share.
func TestSpSolveMatchesDense(t *testing.T) {
	cases := []struct {
		name string
		n    int
		diag float64
		off  float64
	}{
		{name: "identity_like", n: 3, diag: 4, off: 0},
		{name: "weakly_coupled", n: 5, diag: 4, off: -1},
		{name: "strongly_coupled", n: 8, diag: 10, off: 3},
		{name: "larger_system", n: 40, diag: 5, off: 1.5},
	}
	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			vals := spdTridiagonal(tt.n, tt.diag, tt.off)
			sp := sparseFromDense(t, vals, tt.n)
			rhs := make([]float64, tt.n)
			for i := range tt.n {
				rhs[i] = float64(i+1) / float64(tt.n)
			}
			b, err := core.FromFloats(rhs, tt.n)
			if err != nil {
				t.Fatalf("FromFloats: %v", err)
			}
			got, err := SpSolve(sp, b, 0, 0)
			if err != nil {
				t.Fatalf("SpSolve: %v", err)
			}
			want := denseSolve(t, vals, tt.n, rhs)
			x := make([]float64, tt.n)
			for i := range tt.n {
				x[i] = got.FloatAt(i)
			}
			for i := range tt.n {
				if math.Abs(x[i]-want[i]) > 1e-8*(1+math.Abs(want[i])) {
					t.Fatalf("x[%d] = %.12g, want %.12g", i, x[i], want[i])
				}
			}
			if r := residualNorm(t, vals, x, rhs, tt.n); r > 1e-8 {
				t.Fatalf("residual ‖b-Ax‖ = %.3g, want <= 1e-8", r)
			}
		})
	}
}

// TestSpSolveNonDiagonalMatrix exercises a symmetric
// positive-definite matrix that is not tridiagonal, so the sparse
// structure carries a genuinely two-dimensional sparsity pattern.
func TestSpSolveNonDiagonalMatrix(t *testing.T) {
	const n = 6
	// A = L·Lᵀ for a lower-triangular L with a positive diagonal,
	// which is symmetric positive-definite by construction.
	l := []float64{
		2, 0, 0, 0, 0, 0,
		1, 3, 0, 0, 0, 0,
		0, 1, 2, 0, 0, 0,
		1, 0, 1, 4, 0, 0,
		0, 0, 0, 1, 2, 0,
		0, 1, 0, 0, 1, 3,
	}
	vals := make([]float64, n*n)
	for i := range n {
		for j := range n {
			s := 0.0
			for k := range n {
				if k <= i && k <= j {
					s += l[i*n+k] * l[j*n+k]
				}
			}
			vals[i*n+j] = s
		}
	}
	sp := sparseFromDense(t, vals, n)
	rhs := []float64{1, 2, 3, 4, 5, 6}
	b, err := core.FromFloats(rhs, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	got, err := SpSolve(sp, b, 0, 0)
	if err != nil {
		t.Fatalf("SpSolve: %v", err)
	}
	x := make([]float64, n)
	for i := range n {
		x[i] = got.FloatAt(i)
	}
	want := denseSolve(t, vals, n, rhs)
	for i := range n {
		if math.Abs(x[i]-want[i]) > 1e-8*(1+math.Abs(want[i])) {
			t.Fatalf("x[%d] = %.12g, want %.12g", i, x[i], want[i])
		}
	}
}

// TestSpSolveZeroRightHandSide pins the degenerate contract: the
// solution of A·x = 0 is the zero vector, returned without dividing
// by a zero residual norm.
func TestSpSolveZeroRightHandSide(t *testing.T) {
	const n = 4
	vals := spdTridiagonal(n, 4, -1)
	sp := sparseFromDense(t, vals, n)
	b, err := core.Zeros(core.Float, n)
	if err != nil {
		t.Fatalf("Zeros: %v", err)
	}
	got, err := SpSolve(sp, b, 0, 0)
	if err != nil {
		t.Fatalf("SpSolve: %v", err)
	}
	for i := range n {
		if got.FloatAt(i) != 0 {
			t.Fatalf("x[%d] = %.12g, want 0 for a zero right-hand side", i, got.FloatAt(i))
		}
	}
}

// TestSpSolveDeterminism checks that the solve is reproducible: the
// iteration starts from a fixed zero guess and draws nothing random.
func TestSpSolveDeterminism(t *testing.T) {
	const n = 12
	vals := spdTridiagonal(n, 6, -1)
	sp := sparseFromDense(t, vals, n)
	rhs := make([]float64, n)
	for i := range n {
		rhs[i] = math.Sin(float64(i + 1))
	}
	b, err := core.FromFloats(rhs, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	x1, err := SpSolve(sp, b, 0, 0)
	if err != nil {
		t.Fatalf("SpSolve #1: %v", err)
	}
	x2, err := SpSolve(sp, b, 0, 0)
	if err != nil {
		t.Fatalf("SpSolve #2: %v", err)
	}
	for i := range n {
		if x1.FloatAt(i) != x2.FloatAt(i) {
			t.Fatalf("x[%d] = %.12g vs %.12g across runs", i, x1.FloatAt(i), x2.FloatAt(i))
		}
	}
}

// TestSpSolveTolerance drives the stopping rule explicitly: a loose
// tolerance stops early with a larger residual, a tight one converges
// further, and both are reported honestly.
func TestSpSolveTolerance(t *testing.T) {
	const n = 60
	vals := spdTridiagonal(n, 4, -1)
	sp := sparseFromDense(t, vals, n)
	rhs := make([]float64, n)
	for i := range n {
		rhs[i] = 1
	}
	b, err := core.FromFloats(rhs, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	loose, err := SpSolve(sp, b, 1e-2, 0)
	if err != nil {
		t.Fatalf("SpSolve loose: %v", err)
	}
	tight, err := SpSolve(sp, b, 1e-12, 0)
	if err != nil {
		t.Fatalf("SpSolve tight: %v", err)
	}
	xl := make([]float64, n)
	xt := make([]float64, n)
	for i := range n {
		xl[i] = loose.FloatAt(i)
		xt[i] = tight.FloatAt(i)
	}
	rl := residualNorm(t, vals, xl, rhs, n)
	rt := residualNorm(t, vals, xt, rhs, n)
	if rl < rt {
		t.Fatalf("loose tolerance gave residual %.3g, tighter than the tight run's %.3g", rl, rt)
	}
	if rt > 1e-9 {
		t.Fatalf("tight tolerance gave residual %.3g, want <= 1e-9", rt)
	}
}

// TestSpSolveNoConvergence pins the unconverged contract: the budget
// runs out and the solve reports the residual instead of returning a
// silent approximation.
func TestSpSolveNoConvergence(t *testing.T) {
	const n = 200
	vals := spdTridiagonal(n, 2, -0.999)
	sp := sparseFromDense(t, vals, n)
	rhs := make([]float64, n)
	for i := range n {
		rhs[i] = 1
	}
	b, err := core.FromFloats(rhs, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	// A nearly singular system needs far more than two steps.
	if _, err := SpSolve(sp, b, 1e-14, 2); err == nil {
		t.Fatal("expected an error when the iteration budget is exhausted")
	}
}

// TestSpSolveRejectsInvalid pins the error contract for every input
// the solver cannot honestly answer.
func TestSpSolveRejectsInvalid(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)
		}
		b, _ := core.FromFloats([]float64{1}, 1)
		if _, err := SpSolve(sp, b, 0, 0); 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)
		}
		b, _ := core.FromFloats([]float64{1}, 1)
		if _, err := SpSolve(sp, b, 0, 0); err == nil {
			t.Fatal("expected an error for a non-square matrix")
		}
	})
	t.Run("zero_sized", func(t *testing.T) {
		idx, _ := core.FromInts(nil, 0, 2)
		vals, _ := core.FromFloats(nil, 0)
		sp := &core.SparseCOO{Indices: idx, Values: vals, Shape: []int{0, 0}}
		b, _ := core.FromFloats(nil, 0)
		if _, err := SpSolve(sp, b, 0, 0); err == nil {
			t.Fatal("expected an error for a zero-sized matrix")
		}
	})
	t.Run("rhs_wrong_length", func(t *testing.T) {
		vals := spdTridiagonal(3, 4, -1)
		sp := sparseFromDense(t, vals, 3)
		b, _ := core.FromFloats([]float64{1, 2}, 2)
		if _, err := SpSolve(sp, b, 0, 0); err == nil {
			t.Fatal("expected an error for a right-hand side of the wrong length")
		}
	})
	t.Run("rhs_not_vector", func(t *testing.T) {
		vals := spdTridiagonal(2, 4, -1)
		sp := sparseFromDense(t, vals, 2)
		b, _ := core.FromFloats([]float64{1, 2, 3, 4}, 2, 2)
		if _, err := SpSolve(sp, b, 0, 0); err == nil {
			t.Fatal("expected an error for a rank-2 right-hand side")
		}
	})
	t.Run("complex_rhs", func(t *testing.T) {
		vals := spdTridiagonal(2, 4, -1)
		sp := sparseFromDense(t, vals, 2)
		b, _ := core.FromComplexes([]complex128{1, 1}, 2)
		if _, err := SpSolve(sp, b, 0, 0); err == nil {
			t.Fatal("expected an error for a complex right-hand side")
		}
	})
	t.Run("zero_diagonal", func(t *testing.T) {
		// Symmetric but singular, and the Jacobi preconditioner
		// cannot divide by a zero diagonal.
		sp := sparseFromDense(t, []float64{0, 1, 1, 0}, 2)
		b, _ := core.FromFloats([]float64{1, 1}, 2)
		if _, err := SpSolve(sp, b, 0, 0); err == nil {
			t.Fatal("expected an error for a zero diagonal entry")
		}
	})
	t.Run("asymmetric", func(t *testing.T) {
		sp := sparseFromDense(t, []float64{4, 1, 2, 4}, 2)
		b, _ := core.FromFloats([]float64{1, 1}, 2)
		if _, err := SpSolve(sp, b, 0, 0); err == nil {
			t.Fatal("expected an error for an asymmetric matrix")
		}
	})
	t.Run("indefinite", func(t *testing.T) {
		// Symmetric with eigenvalues 3 and -1, so not
		// positive-definite. The right-hand side must excite the
		// negative eigendirection [1,-1]: with b = [1,1] the system
		// still has the exact solution [1/3,1/3] and the curvature
		// stays positive, so the matrix would not be caught.
		sp := sparseFromDense(t, []float64{1, 2, 2, 1}, 2)
		b, _ := core.FromFloats([]float64{1, -1}, 2)
		if _, err := SpSolve(sp, b, 0, 0); err == nil {
			t.Fatal("expected an error for an indefinite matrix")
		}
	})
}