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

// complexCOO builds a SparseCOO from (row, col, val) triples.
func complexCOO(t *testing.T, n int, entries []complex128) *core.SparseCOO {
	t.Helper()
	idx := make([]int64, 0, 2*len(entries)/3)
	vals := make([]complex128, 0, len(entries)/3)
	for i := 0; i+2 < len(entries); i += 3 {
		idx = append(idx, int64(real(entries[i])), int64(real(entries[i+1])))
		vals = append(vals, entries[i+2])
	}
	idxArr, err := core.FromInts(idx, len(vals), 2)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	valArr, err := core.FromComplexes(vals, len(vals))
	if err != nil {
		t.Fatalf("FromComplexes: %v", err)
	}
	coo, err := core.NewSparseCOO(idxArr, valArr, []int{n, n})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	return coo
}

// TestSpSolveComplexCG pins the Hermitian solver: the diagonal
// complex matrix with Hermitian couplings must reproduce the known
// solution to tolerance.
func TestSpSolveComplexCG(t *testing.T) {
	// A = diag(2, 3) with coupling i on both sides: Hermitian.
	sp := complexCOO(t, 2, []complex128{
		0, 0, 2,
		1, 1, 3,
		0, 1, 1i,
		1, 0, -1i,
	})
	b, err := core.FromComplexes([]complex128{1 + 2i, -1}, 2)
	if err != nil {
		t.Fatalf("FromComplexes: %v", err)
	}
	x, err := SpSolveComplexCG(sp, b, 1e-12, 50)
	if err != nil {
		t.Fatalf("SpSolveComplexCG: %v", err)
	}
	// Verify A·x = b directly.
	cm, err := cooToComplexCSR(sp, "test")
	if err != nil {
		t.Fatalf("cooToComplexCSR: %v", err)
	}
	xv := make([]complex128, 2)
	for i := range 2 {
		xv[i] = x.ComplexAt(i)
	}
	ax := make([]complex128, 2)
	cm.matVec(xv, ax)
	for i := range 2 {
		if cmplxAbs(ax[i]-b.ComplexAt(i)) > 1e-10 {
			t.Fatalf("residual[%d] = %v, want 0", i, ax[i]-b.ComplexAt(i))
		}
	}
}

// TestSpSolveComplexCGRejectsNonHermitian pins the symmetry gate.
func TestSpSolveComplexCGRejectsNonHermitian(t *testing.T) {
	sp := complexCOO(t, 2, []complex128{
		0, 0, 2,
		1, 1, 3,
		0, 1, 1i,
		1, 0, 1i, // same sign: not conjugate mirrors
	})
	b, _ := core.FromComplexes([]complex128{1, 1}, 2)
	if _, err := SpSolveComplexCG(sp, b, 1e-12, 50); err == nil {
		t.Fatal("SpSolveComplexCG accepted a non-Hermitian matrix")
	}
}

// TestSpSolveComplexBiCGSTAB pins the general solver on a genuinely
// non-Hermitian complex system.
func TestSpSolveComplexBiCGSTAB(t *testing.T) {
	sp := complexCOO(t, 3, []complex128{
		0, 0, 2 + 1i,
		1, 1, 3,
		2, 2, 1 - 1i,
		0, 1, 1i,
		1, 0, 2, // A[1][0] ≠ conj(A[0][1]): non-Hermitian
		1, 2, 0.5,
	})
	b, err := core.FromComplexes([]complex128{1, 2 - 1i, 3}, 3)
	if err != nil {
		t.Fatalf("FromComplexes: %v", err)
	}
	x, err := SpSolveComplexBiCGSTAB(sp, b, 1e-12, 100)
	if err != nil {
		t.Fatalf("SpSolveComplexBiCGSTAB: %v", err)
	}
	cm, err := cooToComplexCSR(sp, "test")
	if err != nil {
		t.Fatalf("cooToComplexCSR: %v", err)
	}
	xv := make([]complex128, 3)
	for i := range 3 {
		xv[i] = x.ComplexAt(i)
	}
	ax := make([]complex128, 3)
	cm.matVec(xv, ax)
	for i := range 3 {
		if cmplxAbs(ax[i]-b.ComplexAt(i)) > 1e-9 {
			t.Fatalf("residual[%d] = %v, want 0", i, ax[i]-b.ComplexAt(i))
		}
	}
}

// TestSpSolveComplexBiCGSTABNaN pins the NaN guards: a poisoned entry
// must be a breakdown error, never a converged garbage answer.
func TestSpSolveComplexBiCGSTABNaN(t *testing.T) {
	sp := complexCOO(t, 2, []complex128{
		0, 0, 1,
		1, 1, complex(math.NaN(), 0),
	})
	b, _ := core.FromComplexes([]complex128{1, 1}, 2)
	if _, err := SpSolveComplexBiCGSTAB(sp, b, 1e-10, 50); err == nil {
		t.Fatal("BiCGSTAB reported convergence on a NaN matrix")
	}
}

// TestSpEigenComplex pins the Hermitian Lanczos against a matrix with
// known spectrum: eigenvalues and the residual ‖A·v − λ·v‖.
func TestSpEigenComplex(t *testing.T) {
	// [[2, i], [−i, 2]] has eigenvalues 3 and 1 with eigenvectors
	// (1, −i)/√2 (for 3) and (1, i)/√2 (for 1).
	sp := complexCOO(t, 2, []complex128{
		0, 0, 2,
		1, 1, 2,
		0, 1, 1i,
		1, 0, -1i,
	})
	vals, vecs, err := SpEigenComplex(sp, 2, core.NewGenerator(7))
	if err != nil {
		t.Fatalf("SpEigenComplex: %v", err)
	}
	hi, lo := vals.FloatAt(0), vals.FloatAt(1)
	if math.Abs(hi-3) > 1e-10 || math.Abs(lo-1) > 1e-10 {
		t.Fatalf("eigenvalues (%g, %g), want (3, 1)", hi, lo)
	}
	cm, _ := cooToComplexCSR(sp, "test")
	for j := range 2 {
		v := make([]complex128, 2)
		for i := range 2 {
			v[i] = vecs.ComplexAt(i*2 + j)
		}
		if n := normC(v); math.Abs(n-1) > 1e-10 {
			t.Fatalf("eigenvector %d norm %g, want 1", j, n)
		}
		av := make([]complex128, 2)
		cm.matVec(v, av)
		lam := complex(vals.FloatAt(j), 0)
		for i := range 2 {
			if cmplxAbs(av[i]-lam*v[i]) > 1e-9 {
				t.Fatalf("eigenpair %d residual[%d] = %v", j, i, av[i]-lam*v[i])
			}
		}
	}
}

// TestSpEigenComplexRejectsNonHermitian pins the symmetry gate.
func TestSpEigenComplexRejectsNonHermitian(t *testing.T) {
	sp := complexCOO(t, 2, []complex128{
		0, 0, 2,
		1, 1, 2,
		0, 1, 1i,
		1, 0, 1i,
	})
	if _, _, err := SpEigenComplex(sp, 1, nil); err == nil {
		t.Fatal("SpEigenComplex accepted a non-Hermitian matrix")
	}
}

// TestComplexCSRFromCOOMergesDuplicates pins the COO semantics: a
// coordinate repeated twice sums.
func TestComplexCSRFromCOOMergesDuplicates(t *testing.T) {
	sp := complexCOO(t, 2, []complex128{
		0, 0, 1 + 1i,
		0, 0, 2 - 1i,
		1, 1, 4,
	})
	cm, err := cooToComplexCSR(sp, "test")
	if err != nil {
		t.Fatalf("cooToComplexCSR: %v", err)
	}
	v, ok := cm.at(0, 0)
	if !ok || v != 3 {
		t.Fatalf("merged (0,0) = %v ok=%v, want 3+0i", v, ok)
	}
	if got := len(cm.vals); got != 2 {
		t.Fatalf("stored %d entries, want 2", got)
	}
}