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

package linalg

import (
	"math"
	"testing"

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

// TestSpEigenGeneralRotationBlocks pins the real general solver: block
// rotations give complex conjugate eigenvalue pairs a symmetric-only
// method could never reach, and the answer must agree with the dense
// EigenGeneral on the same matrix.
func TestSpEigenGeneralRotationBlocks(t *testing.T) {
	// Block diagonal: rotations by 5 and 2 plus one 7: spectrum
	// {7, ±5i, ±2i}.
	dense := []float64{
		0, -5, 0, 0, 0,
		5, 0, 0, 0, 0,
		0, 0, 0, -2, 0,
		0, 0, 2, 0, 0,
		0, 0, 0, 0, 7,
	}
	d, err := core.FromFloats(dense, 5, 5)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	sp, err := core.SparseFrom(d)
	if err != nil {
		t.Fatalf("SparseFrom: %v", err)
	}
	vals, vecs, err := SpEigenGeneral(sp, 3, core.NewGenerator(31))
	if err != nil {
		t.Fatalf("SpEigenGeneral: %v", err)
	}
	want, _, err := EigenGeneral(d)
	if err != nil {
		t.Fatalf("EigenGeneral: %v", err)
	}
	// The spectrum is compared as a multiset against the dense one: a
	// conjugate pair shares one magnitude, so which of the two a Krylov
	// run reports at which index is an arbitrary tie-break of its own
	// rounding, not a property of the matrix, and pinning it made the
	// test depend on the last bit of a magnitude. Every computed value
	// must still match some dense value, and the residual loop below
	// pins each value to its own vector.
	used := make([]bool, 3)
	for j := range 3 {
		got := vals.ComplexAt(j)
		best, bestDist := -1, math.Inf(1)
		for i := range 3 {
			if used[i] {
				continue
			}
			w := want.ComplexAt(i)
			if dist := math.Hypot(real(got)-real(w), imag(got)-imag(w)); dist < bestDist {
				best, bestDist = i, dist
			}
		}
		if best < 0 || bestDist > 1e-8 {
			t.Fatalf("value[%d] = %v matches no dense value (closest %v at distance %.3g)",
				j, got, want.ComplexAt(best), bestDist)
		}
		used[best] = true
	}
	// Residuals ‖A·v − λ·v‖ against the original sparse operator. The
	// eigenvectors may carry any complex phase, so the full complex
	// vector enters the check.
	for j := range 3 {
		v := make([]complex128, 5)
		for i := range 5 {
			v[i] = vecs.ComplexAt(i*3 + j)
		}
		av := make([]complex128, 5)
		for i := range 5 {
			for p := range 5 {
				av[i] += complex(dense[i*5+p], 0) * v[p]
			}
		}
		lam := vals.ComplexAt(j)
		for i := range 5 {
			res := av[i] - lam*v[i]
			if math.Hypot(real(res), imag(res)) > 1e-7 {
				t.Fatalf("residual[%d][%d] = %v", j, i, res)
			}
		}
	}
}

// TestSpEigenGeneralComplexTriangular pins the complex general solver:
// a non-Hermitian triangular operator whose spectrum is its diagonal.
func TestSpEigenGeneralComplexTriangular(t *testing.T) {
	// Upper triangular with distinct diagonal; the off-diagonal
	// couplings make it genuinely non-normal.
	entries := []complex128{
		0, 0, 2 + 3i,
		1, 1, -1 + 1i,
		2, 2, 0.5 - 2i,
		0, 1, 0.7 + 0.3i,
		1, 2, -0.4,
	}
	idx := make([]int64, 0, 10)
	valsIn := make([]complex128, 0, 5)
	for i := 0; i+2 < len(entries); i += 3 {
		idx = append(idx, int64(real(entries[i])), int64(real(entries[i+1])))
		valsIn = append(valsIn, entries[i+2])
	}
	idxArr, err := core.FromInts(idx, len(valsIn), 2)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	valArr, err := core.FromComplexes(valsIn, len(valsIn))
	if err != nil {
		t.Fatalf("FromComplexes: %v", err)
	}
	sp, err := core.NewSparseCOO(idxArr, valArr, []int{3, 3})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	vals, vecs, err := SpEigenGeneralComplex(sp, 2, core.NewGenerator(5))
	if err != nil {
		t.Fatalf("SpEigenGeneralComplex: %v", err)
	}
	// |2+3i| ≈ 3.606 > |0.5−2i| ≈ 2.062 > |−1+i| ≈ 1.414.
	wantTop := complex(2, 3)
	wantSecond := complex(0.5, -2)
	for j, want := range []complex128{wantTop, wantSecond} {
		got := vals.ComplexAt(j)
		if math.Abs(real(got)-real(want)) > 1e-9 || math.Abs(imag(got)-imag(want)) > 1e-9 {
			t.Fatalf("value[%d] = %v, want %v", j, got, want)
		}
	}
	// Residual against the dense operator.
	dense := make([]complex128, 9)
	for i := 0; i+2 < len(entries); i += 3 {
		dense[int(real(entries[i]))*3+int(real(entries[i+1]))] = entries[i+2]
	}
	for j := range 2 {
		v := make([]complex128, 3)
		for i := range 3 {
			v[i] = vecs.ComplexAt(i*2 + j)
		}
		var n2 float64
		for _, z := range v {
			n2 += real(z)*real(z) + imag(z)*imag(z)
		}
		if math.Abs(n2-1) > 1e-8 {
			t.Fatalf("vector %d norm² = %g, want 1", j, n2)
		}
		av := make([]complex128, 3)
		for i := range 3 {
			for p := range 3 {
				av[i] += dense[i*3+p] * v[p]
			}
		}
		lam := vals.ComplexAt(j)
		for i := range 3 {
			res := av[i] - lam*v[i]
			if math.Hypot(real(res), imag(res)) > 1e-8 {
				t.Fatalf("residual[%d][%d] = %v", j, i, res)
			}
		}
	}
}

// TestSpEigenGeneralLargerMatrix pins convergence on a bigger
// nonsymmetric operator: the top eigenvalues must match the dense
// reference.
func TestSpEigenGeneralLargerMatrix(t *testing.T) {
	const n = 40
	g := core.NewGenerator(77)
	dense := make([]float64, n*n)
	for i := range n * n {
		dense[i] = g.NormalUnit()
	}
	// A sprinkle of larger entries decides the spectrum's top end.
	for i := range n {
		dense[i*n+i] += 6 * float64(n-i) / float64(n)
	}
	d, err := core.FromFloats(dense, n, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	sp, err := core.SparseFrom(d)
	if err != nil {
		t.Fatalf("SparseFrom: %v", err)
	}
	vals, _, err := SpEigenGeneral(sp, 2, core.NewGenerator(9))
	if err != nil {
		t.Fatalf("SpEigenGeneral: %v", err)
	}
	want, _, err := EigenGeneral(d)
	if err != nil {
		t.Fatalf("EigenGeneral: %v", err)
	}
	// A real matrix carries conjugate twins of equal magnitude, so
	// the k-th slot may hold either member: match against the set.
	for j := range 2 {
		got := vals.ComplexAt(j)
		ok := cmplxAbs(got-want.ComplexAt(j)) <= 1e-6*cmplxAbs(got) ||
			cmplxAbs(got-complexConj(want.ComplexAt(j))) <= 1e-6*cmplxAbs(got)
		if !ok {
			t.Fatalf("value[%d] = %v, dense says %v (or its conjugate)", j, got, want.ComplexAt(j))
		}
	}
}

// TestSpEigenGeneralErrors pins the routing and validation.
func TestSpEigenGeneralErrors(t *testing.T) {
	d, _ := core.FromFloats([]float64{0, -1, 1, 0}, 2, 2)
	sp, _ := core.SparseFrom(d)
	if _, _, err := SpEigenGeneral(sp, 0, nil); err == nil {
		t.Error("k = 0 accepted")
	}
	if _, _, err := SpEigenGeneral(sp, 3, nil); err == nil {
		t.Error("k > n accepted")
	}
	// Complex input routes to the complex entry point.
	idx, _ := core.FromInts([]int64{0, 0, 1, 1}, 2, 2)
	cv, _ := core.FromComplexes([]complex128{1, 2}, 2)
	cc, _ := core.NewSparseCOO(idx, cv, []int{2, 2})
	if _, _, err := SpEigenGeneral(cc, 1, nil); err == nil {
		t.Error("SpEigenGeneral accepted complex values")
	}
	rv, _ := core.FromFloats([]float64{1, 2}, 2)
	rc, _ := core.NewSparseCOO(idx, rv, []int{2, 2})
	if _, _, err := SpEigenGeneralComplex(rc, 1, nil); err == nil {
		t.Error("SpEigenGeneralComplex accepted real values")
	}
}