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

// TestSpSolveBiCGSTABNonsymmetric checks the solver against the dense
// Solve on a 4×4 system with genuinely asymmetric entries.
func TestSpSolveBiCGSTABNonsymmetric(t *testing.T) {
	dense := [][]float64{
		{10, 1, 0, 2},
		{-1, 12, 3, 0},
		{0, -2, 15, 1},
		{2, 0, -1, 8},
	}
	entries := make([][3]float64, 0, 16)
	for i := range 4 {
		for j := range 4 {
			if dense[i][j] != 0 {
				entries = append(entries, [3]float64{float64(i), float64(j), dense[i][j]})
			}
		}
	}
	idx := make([]int64, 0, len(entries))
	vals := make([]float64, 0, len(entries))
	for _, e := range entries {
		idx = append(idx, int64(e[0]), int64(e[1]))
		vals = append(vals, e[2])
	}
	indices, err := core.FromInts(idx, len(entries), 2)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	coo, err := core.NewSparseCOO(indices, floatsToArray(vals, []int{len(entries)}), []int{4, 4})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	b := mustFloats(t, []float64{1, 2, 3, 4})
	x, err := SpSolveBiCGSTAB(coo, b, 1e-12, 0)
	if err != nil {
		t.Fatalf("SpSolveBiCGSTAB: %v", err)
	}
	ref, err := Solve(mustFloats(t, []float64{
		10, 1, 0, 2,
		-1, 12, 3, 0,
		0, -2, 15, 1,
		2, 0, -1, 8,
	}, 4, 4), b)
	if err != nil {
		t.Fatalf("Solve: %v", err)
	}
	for i := range 4 {
		if math.Abs(x.FloatAt(i)-ref.FloatAt(i)) > 1e-9 {
			t.Fatalf("x[%d] = %.14g, want %.14g", i, x.FloatAt(i), ref.FloatAt(i))
		}
	}
}

// TestSpSolveBiCGSTABConvective builds a 50×50 shifted Laplacian with
// an asymmetric convection term, the shape transport discretisations
// produce, and checks the residual the solver promises.
func TestSpSolveBiCGSTABConvective(t *testing.T) {
	const n = 50
	idx := make([]int64, 0, 4*n)
	vals := make([]float64, 0, 4*n)
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	for i := range n {
		add(i, i, 4)
		if i+1 < n {
			add(i, i+1, -1)
			add(i+1, i, -1+0.5) // asymmetric neighbour coupling
		}
	}
	indices, err := core.FromInts(idx, len(vals), 2)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	coo, err := core.NewSparseCOO(indices, floatsToArray(vals, []int{len(vals)}), []int{n, n})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	bv := make([]float64, n)
	for i := range n {
		bv[i] = float64(i+1) / float64(n+1)
	}
	b := mustFloats(t, bv)
	x, err := SpSolveBiCGSTAB(coo, b, 1e-10, 0)
	if err != nil {
		t.Fatalf("SpSolveBiCGSTAB: %v", err)
	}
	// Residual check against the sparse product path.
	ax, err := cooToCSRForTest(t, coo).MatVec(x)
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	res := 0.0
	for i := range n {
		d := b.FloatAt(i) - ax.FloatAt(i)
		res += d * d
	}
	if math.Sqrt(res) > 1e-9*math.Sqrt(dotF64(bv, bv)) {
		t.Fatalf("relative residual %g exceeds 1e-9", math.Sqrt(res)/math.Sqrt(dotF64(bv, bv)))
	}
}

// cooToCSRForTest converts a COO to the public CSR for residual checks.
func cooToCSRForTest(t *testing.T, coo *core.SparseCOO) *SparseCSR {
	t.Helper()
	csr, err := CSRFromCOO(coo)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	return csr
}

// TestSpSolveBiCGSTABConvergenceFailure reports an honest failure when
// the budget cannot reach the tolerance.
func TestSpSolveBiCGSTABConvergenceFailure(t *testing.T) {
	const n = 30
	idx := make([]int64, 0, 3*n)
	vals := make([]float64, 0, 3*n)
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	for i := range n {
		add(i, i, 4)
		if i+1 < n {
			add(i, i+1, -1)
			add(i+1, i, -0.5)
		}
	}
	indices, _ := core.FromInts(idx, len(vals), 2)
	coo, err := core.NewSparseCOO(indices, floatsToArray(vals, []int{len(vals)}), []int{n, n})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	b := mustFloats(t, make([]float64, n))
	for i := range n {
		b.RawFloats()[i] = 1
	}
	if _, err := SpSolveBiCGSTAB(coo, b, 1e-14, 2); err == nil {
		t.Fatal("tight tolerance with a two-step budget: want an error")
	}
}

// TestSpSolveBiCGSTABErrors pins the validation surface.
func TestSpSolveBiCGSTABErrors(t *testing.T) {
	// Build the COO inputs directly for the error paths.
	mkCOO := func(t *testing.T, rows, cols int, entries [][3]float64) *core.SparseCOO {
		t.Helper()
		idx := make([]int64, 0, len(entries)*2)
		vals := make([]float64, 0, len(entries))
		for _, e := range entries {
			idx = append(idx, int64(e[0]), int64(e[1]))
			vals = append(vals, e[2])
		}
		indices, err := core.FromInts(idx, len(entries), 2)
		if err != nil {
			t.Fatalf("FromInts: %v", err)
		}
		coo, err := core.NewSparseCOO(indices, floatsToArray(vals, []int{len(entries)}), []int{rows, cols})
		if err != nil {
			t.Fatalf("NewSparseCOO: %v", err)
		}
		return coo
	}
	b2 := mustFloats(t, []float64{1, 1})
	if _, err := SpSolveBiCGSTAB(mkCOO(t, 2, 3, [][3]float64{{0, 0, 2}}), b2, 0, 0); err == nil {
		t.Fatal("non-square matrix: want an error")
	}
	// A missing diagonal entry defeats the Jacobi preconditioner.
	gap := mkCOO(t, 2, 2, [][3]float64{{0, 1, 1}, {1, 1, 3}})
	if _, err := SpSolveBiCGSTAB(gap, b2, 0, 0); err == nil {
		t.Fatal("missing diagonal: want an error")
	}
	bad := mkCOO(t, 2, 2, [][3]float64{{0, 0, 2}, {1, 1, 3}})
	if _, err := SpSolveBiCGSTAB(bad, mustFloats(t, []float64{1, 2, 3}), 0, 0); err == nil {
		t.Fatal("wrong right-hand side length: want an error")
	}
	complexVals := core.New(core.Complex, 2)
	cx := mkCOO(t, 2, 2, [][3]float64{{0, 0, 1}, {1, 1, 1}})
	cx.Values = complexVals
	if _, err := SpSolveBiCGSTAB(cx, b2, 0, 0); err == nil {
		t.Fatal("complex values: want an error")
	}
}