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

package linalg

import (
	"math"
	"slices"
	"strings"
	"testing"

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

// gridLaplacianCOO builds the 5-point Laplacian on a w by h grid: the
// standard sparse positive definite test matrix, symmetric with a
// dominant diagonal. When shuffle is true the grid vertices are
// relabelled by a fixed seeded permutation first, so the stored order
// carries no bandwidth for the natural ordering to lean on.
func gridLaplacianCOOShuffled(t *testing.T, w, h int, shuffle bool) *core.SparseCOO {
	t.Helper()
	idx := make([]int64, 0, 5*w*h)
	vals := make([]float64, 0, 5*w*h)
	label := func(x, y int) int { return y*w + x }
	if shuffle {
		g := core.NewGenerator(5)
		perm := make([]int, w*h)
		for i := range perm {
			perm[i] = i
		}
		// Fisher-Yates with the seeded generator, a fixed relabelling.
		for i := len(perm) - 1; i > 0; i-- {
			j := int(g.Next() % uint64(i+1))
			perm[i], perm[j] = perm[j], perm[i]
		}
		label = func(x, y int) int { return perm[y*w+x] }
	}
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	at := func(x, y int) int { return label(x, y) }
	for y := range h {
		for x := range w {
			add(at(x, y), at(x, y), 4)
			if x+1 < w {
				add(at(x, y), at(x+1, y), -1)
				add(at(x+1, y), at(x, y), -1)
			}
			if y+1 < h {
				add(at(x, y), at(x, y+1), -1)
				add(at(x, y+1), at(x, y), -1)
			}
		}
	}
	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{w * h, w * h})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	return coo
}

// gridLaplacianCOO builds the plain row-major grid.
func gridLaplacianCOO(t *testing.T, w, h int) *core.SparseCOO {
	t.Helper()
	return gridLaplacianCOOShuffled(t, w, h, false)
}

func TestSparseCholeskySolvesLaplacian(t *testing.T) {
	const w, h = 12, 10
	coo := gridLaplacianCOO(t, w, h)
	n := w * h
	csr, err := CSRFromCOO(coo)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	// The truth is constructive: xTrue is fixed, b comes from the
	// independently tested CSR MatVec, and the factor must invert it.
	xTrue := core.New(core.Float, n)
	for i := range n {
		xTrue.RawFloats()[i] = math.Sin(float64(i)) + float64(i%7)*0.1
	}
	b, err := csr.MatVec(xTrue)
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	for _, ordering := range []SparseOrdering{SparseOrderingNatural, SparseOrderingReverseCuthillMcKee} {
		f, err := NewSparseCholesky(coo, ordering)
		if err != nil {
			t.Fatalf("NewSparseCholesky(%d): %v", ordering, err)
		}
		x, err := f.Solve(b)
		if err != nil {
			t.Fatalf("Solve(%d): %v", ordering, err)
		}
		worst := 0.0
		for i := range n {
			if d := math.Abs(x.FloatAt(i) - xTrue.FloatAt(i)); d > worst {
				worst = d
			}
		}
		if worst > 1e-9 {
			t.Fatalf("ordering %d: worst solution error %.3g, want under 1e-9", ordering, worst)
		}
		// The residual through the original matrix closes the loop:
		// A·x must give b back.
		ax, err := csr.MatVec(x)
		if err != nil {
			t.Fatalf("MatVec: %v", err)
		}
		res := 0.0
		for i := range n {
			if d := math.Abs(ax.FloatAt(i) - b.FloatAt(i)); d > res {
				res = d
			}
		}
		if res > 1e-8 {
			t.Fatalf("ordering %d: residual %.3g, want under 1e-8", ordering, res)
		}
	}
}

// irregularSPDCOO builds a diagonally dominant symmetric matrix with
// long-range random couplings: the pattern no bandwidth ordering can
// tame, the case the minimum degree order exists for.
func irregularSPDCOO(t *testing.T, seed int64, n int) *core.SparseCOO {
	t.Helper()
	g := core.NewGenerator(seed)
	deg := make([]float64, n)
	type entry struct {
		r, c int
	}
	var entries []entry
	for range 700 {
		i := int(g.Next() % uint64(n))
		j := int(g.Next() % uint64(n))
		if i == j {
			continue
		}
		entries = append(entries, entry{i, j})
		deg[i]++
		deg[j]++
	}
	idx := make([]int64, 0, 2*len(entries)+2*n)
	vals := make([]float64, 0, 2*len(entries)+n)
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	for _, e := range entries {
		add(e.r, e.c, -1)
		add(e.c, e.r, -1)
	}
	for i := range n {
		add(i, i, float64(deg[i])+2)
	}
	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)
	}
	return coo
}

// TestSparseCholeskyMinimumDegreeOrdering pins the minimum degree
// order's property: on both a shuffled mesh and a long-range
// irregular pattern its factor stores fewer non-zeros than the
// reverse Cuthill-McKee factor, and the solve stays exact. The
// measured numbers sit far inside the bounds and land in the log.
func TestSparseCholeskyMinimumDegreeOrdering(t *testing.T) {
	coo := gridLaplacianCOOShuffled(t, 15, 15, true)
	nat, err := NewSparseCholesky(coo, SparseOrderingNatural)
	if err != nil {
		t.Fatalf("NewSparseCholesky(natural): %v", err)
	}
	rcm, err := NewSparseCholesky(coo, SparseOrderingReverseCuthillMcKee)
	if err != nil {
		t.Fatalf("NewSparseCholesky(rcm): %v", err)
	}
	md, err := NewSparseCholesky(coo, SparseOrderingMinimumDegree)
	if err != nil {
		t.Fatalf("NewSparseCholesky(md): %v", err)
	}
	t.Logf("grid n=225: natural %d, rcm %d, md %d", nat.NNZ(), rcm.NNZ(), md.NNZ())
	if md.NNZ() >= rcm.NNZ() {
		t.Fatalf("md fill %d is not below rcm fill %d on the grid", md.NNZ(), rcm.NNZ())
	}
	// Long-range irregular pattern: both orders must still solve, and
	// the minimum degree order must keep its edge.
	irr := irregularSPDCOO(t, 11, 300)
	csr, err := CSRFromCOO(irr)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	xTrue := core.New(core.Float, 300)
	for i := range 300 {
		xTrue.RawFloats()[i] = math.Cos(0.3*float64(i)) + float64(i%7)
	}
	b, err := csr.MatVec(xTrue)
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	rcmI, err := NewSparseCholesky(irr, SparseOrderingReverseCuthillMcKee)
	if err != nil {
		t.Fatalf("NewSparseCholesky(rcm): %v", err)
	}
	mdI, err := NewSparseCholesky(irr, SparseOrderingMinimumDegree)
	if err != nil {
		t.Fatalf("NewSparseCholesky(md): %v", err)
	}
	t.Logf("irregular n=300: rcm %d, md %d", rcmI.NNZ(), mdI.NNZ())
	if mdI.NNZ() >= rcmI.NNZ() {
		t.Fatalf("md fill %d is not below rcm fill %d on the irregular pattern", mdI.NNZ(), rcmI.NNZ())
	}
	for name, f := range map[string]*SparseCholesky{"rcm": rcmI, "md": mdI} {
		x, err := f.Solve(b)
		if err != nil {
			t.Fatalf("%s solve: %v", name, err)
		}
		for i := range 300 {
			if math.Abs(x.FloatAt(i)-xTrue.FloatAt(i)) > 1e-9 {
				t.Fatalf("%s: entry %d error %.3g", name, i, math.Abs(x.FloatAt(i)-xTrue.FloatAt(i)))
			}
		}
	}
	perm := md.Permutation()
	sorted := slices.Clone(perm)
	slices.Sort(sorted)
	for i := range sorted {
		if sorted[i] != i {
			t.Fatalf("permutation entry %d holds %d; not a permutation", i, sorted[i])
		}
	}
}

// TestSparseCholeskyMatchesDense factors the same grid Laplacian as a
// dense matrix and requires the two solvers to agree: the sparse
// factorisation is a different algorithm for the same A⁻¹.
func TestSparseCholeskyMatchesDense(t *testing.T) {
	const w, h = 8, 8
	coo := gridLaplacianCOO(t, w, h)
	n := w * h
	csr, err := CSRFromCOO(coo)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	xTrue := core.New(core.Float, n)
	for i := range n {
		xTrue.RawFloats()[i] = math.Cos(0.3*float64(i)) + float64(i%5)
	}
	b, err := csr.MatVec(xTrue)
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	f, err := NewSparseCholesky(coo, SparseOrderingReverseCuthillMcKee)
	if err != nil {
		t.Fatalf("NewSparseCholesky: %v", err)
	}
	sparse, err := f.Solve(b)
	if err != nil {
		t.Fatalf("Solve: %v", err)
	}
	denseVals := make([]float64, n*n)
	nnz := coo.Indices.Shape()[0]
	for i := range nnz {
		r := int(coo.Indices.RawInts()[i*2])
		c := int(coo.Indices.RawInts()[i*2+1])
		denseVals[r*n+c] = coo.Values.FloatAt(i)
	}
	dense, err := Solve(floatsToArray(denseVals, []int{n, n}), b)
	if err != nil {
		t.Fatalf("dense Solve: %v", err)
	}
	scale := 0.0
	for i := range n {
		if d := math.Abs(dense.FloatAt(i)); d > scale {
			scale = d
		}
	}
	for i := range n {
		if math.Abs(sparse.FloatAt(i)-dense.FloatAt(i)) > 1e-9*scale {
			t.Fatalf("entry %d: sparse %.12g vs dense %.12g", i, sparse.FloatAt(i), dense.FloatAt(i))
		}
	}
}

// TestSparseCholeskyFillMeasures pins the property the orderings
// exist for: on a shuffled mesh pattern the reverse Cuthill-McKee
// factor must store markedly fewer non-zeros than the natural order
// factor. The bounds carry slack on purpose; the measured numbers
// land far inside them and the log line records them, so a regression
// in the ordering shows up as a test failure, not as a slow solver.
func TestSparseCholeskyFillMeasures(t *testing.T) {
	coo := gridLaplacianCOOShuffled(t, 15, 15, true)
	natural, err := NewSparseCholesky(coo, SparseOrderingNatural)
	if err != nil {
		t.Fatalf("NewSparseCholesky(natural): %v", err)
	}
	rcm, err := NewSparseCholesky(coo, SparseOrderingReverseCuthillMcKee)
	if err != nil {
		t.Fatalf("NewSparseCholesky(rcm): %v", err)
	}
	t.Logf("n=225 shuffled: natural L nnz %d, rcm L nnz %d", natural.NNZ(), rcm.NNZ())
	if rcm.NNZ() >= natural.NNZ() {
		t.Fatalf("rcm fill %d is not below natural fill %d", rcm.NNZ(), natural.NNZ())
	}
	if rcm.NNZ() > natural.NNZ()/2 {
		t.Fatalf("rcm fill %d did not at least halve natural fill %d", rcm.NNZ(), natural.NNZ())
	}
	// Both factors must still solve: the ordering changes the fill,
	// never the answer.
	csr, err := CSRFromCOO(coo)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	xTrue := core.New(core.Float, 225)
	for i := range xTrue.Len() {
		xTrue.RawFloats()[i] = math.Cos(0.7*float64(i)) + float64(i%11)*0.2
	}
	b, err := csr.MatVec(xTrue)
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	for name, f := range map[string]*SparseCholesky{"natural": natural, "rcm": rcm} {
		x, err := f.Solve(b)
		if err != nil {
			t.Fatalf("%s solve: %v", name, err)
		}
		for i := range xTrue.Len() {
			if math.Abs(x.FloatAt(i)-xTrue.FloatAt(i)) > 1e-9 {
				t.Fatalf("%s: entry %d error %.3g", name, i, math.Abs(x.FloatAt(i)-xTrue.FloatAt(i)))
			}
		}
	}
	perm := rcm.Permutation()
	sorted := slices.Clone(perm)
	slices.Sort(sorted)
	for i := range sorted {
		if sorted[i] != i {
			t.Fatalf("permutation entry %d holds %d; not a permutation", i, sorted[i])
		}
	}
	// The permutation is a copy: moving the caller's slice must not
	// move the factor's.
	perm[0] = -1
	if rcm.Permutation()[0] == -1 {
		t.Fatal("Permutation exposed the factor's internal slice")
	}
}

func TestSparseCholeskyRefusals(t *testing.T) {
	good := gridLaplacianCOO(t, 4, 4)
	if _, err := NewSparseCholesky(good, SparseOrdering(7)); err == nil {
		t.Fatal("an unknown ordering was accepted")
	}
	// A stored upper entry without its lower counterpart is a silent
	// asymmetry the factor refuses to inherit. Entries: (0,0)=4,
	// (0,1)=1, (1,1)=4, (1,2)=1, (2,2)=4, (2,3)=1, (3,3)=4: the upper
	// (0,1) has no lower (1,0).
	oneWay, err := core.NewSparseCOO(
		mustInts(t, []int64{0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3}, 7, 2),
		floatsToArray([]float64{4, 1, 4, 1, 4, 1, 4}, []int{7}),
		[]int{4, 4})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	if _, err := NewSparseCholesky(oneWay, SparseOrderingNatural); err == nil || !strings.Contains(err.Error(), "counterpart") {
		t.Fatalf("one-sided upper entry: %v", err)
	}
	// A conflicting pair refuses too. Entries: (0,0)=4, (0,1)=1,
	// (1,1)=4, (1,2)=2, (2,2)=4, (1,3)=3, (3,3)=4, (2,1)=5: the upper
	// (1,2)=2 and the lower (2,1)=5 disagree.
	conflicting, err := core.NewSparseCOO(
		mustInts(t, []int64{0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 1, 3, 3, 3, 2, 1}, 8, 2),
		floatsToArray([]float64{4, 1, 4, 2, 4, 3, 4, 5}, []int{8}),
		[]int{4, 4})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	if _, err := NewSparseCholesky(conflicting, SparseOrderingNatural); err == nil || !strings.Contains(err.Error(), "counterpart") {
		t.Fatalf("conflicting upper entry: %v", err)
	}
	// Not positive definite: the zero diagonal has no square root.
	sing := triDiagCOO(t, 4, 1, 0, 1)
	if _, err := NewSparseCholesky(sing, SparseOrderingNatural); err == nil || !strings.Contains(err.Error(), "positive definite") {
		t.Fatalf("zero diagonal: %v", err)
	}
	// Non-finite stored value.
	bad, err := core.NewSparseCOO(
		mustInts(t, []int64{0, 0, 1, 1}, 2, 2),
		floatsToArray([]float64{4, math.NaN()}, []int{2}),
		[]int{2, 2})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	if _, err := NewSparseCholesky(bad, SparseOrderingNatural); err == nil || !strings.Contains(err.Error(), "finite") {
		t.Fatalf("NaN entry: %v", err)
	}
	// Rectangular input.
	rect, err := core.NewSparseCOO(
		mustInts(t, []int64{0, 0}, 1, 2),
		floatsToArray([]float64{1}, []int{1}),
		[]int{1, 2})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	if _, err := NewSparseCholesky(rect, SparseOrderingNatural); err == nil {
		t.Fatal("a rectangular matrix was accepted")
	}
	// Solve-side refusals.
	f, err := NewSparseCholesky(good, SparseOrderingNatural)
	if err != nil {
		t.Fatalf("NewSparseCholesky: %v", err)
	}
	if _, err := f.Solve(core.New(core.Float, 3, 3)); err == nil {
		t.Fatal("a rank 2 right hand side was accepted")
	}
	if _, err := f.Solve(floatsToArray([]float64{1, 2, 3}, []int{3})); err == nil {
		t.Fatal("a short right hand side was accepted")
	}
}

// TestSparseCholeskyIsDeterministic factors the same matrix twice and
// requires the factor's stored values to come out bit for bit equal,
// the contract every Tensor entry point carries.
func TestSparseCholeskyIsDeterministic(t *testing.T) {
	coo := gridLaplacianCOO(t, 9, 9)
	f1, err := NewSparseCholesky(coo, SparseOrderingReverseCuthillMcKee)
	if err != nil {
		t.Fatalf("first factorisation: %v", err)
	}
	f2, err := NewSparseCholesky(coo, SparseOrderingReverseCuthillMcKee)
	if err != nil {
		t.Fatalf("second factorisation: %v", err)
	}
	if f1.NNZ() != f2.NNZ() {
		t.Fatalf("factor sizes differ: %d vs %d", f1.NNZ(), f2.NNZ())
	}
	for i := range f1.values {
		if f1.values[i] != f2.values[i] {
			t.Fatalf("value %d differs: %.17g vs %.17g", i, f1.values[i], f2.values[i])
		}
	}
	if slices.Compare(f1.perm, f2.perm) != 0 {
		t.Fatal("permutations differ")
	}
}

// TestCSCCanonicalisation checks CSCFromCOO's contract beside
// CSRFromCOO's: duplicates sum, explicit zeros drop, every column's
// row indices are sorted and unique, and the transpose round trips
// agree entry for entry with the direct conversions.
func TestCSCCanonicalisation(t *testing.T) {
	// Entries: (2,0)=3, (2,0)=1 duplicate, (1,1)=5, (0,2)=7 explicit
	// zero, (0,2)=2.
	coo, err := core.NewSparseCOO(
		mustInts(t, []int64{2, 0, 2, 0, 1, 1, 0, 2, 0, 2}, 5, 2),
		floatsToArray([]float64{3, 1, 5, 0, 2}, []int{5}),
		[]int{3, 3})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	csc, err := CSCFromCOO(coo)
	if err != nil {
		t.Fatalf("CSCFromCOO: %v", err)
	}
	if csc.NNZ() != 3 {
		t.Fatalf("nnz %d after merging and zero drop, want 3", csc.NNZ())
	}
	if csc.Values[0] != 4 {
		t.Fatalf("duplicates did not sum: column 0 first value %g", csc.Values[0])
	}
	for j := range csc.Cols {
		rows := csc.RowIdx[csc.ColStart[j]:csc.ColStart[j+1]]
		if !slices.IsSorted(rows) {
			t.Fatalf("column %d rows not sorted: %v", j, rows)
		}
	}
	csr, err := CSRFromCOO(coo)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	round, err := csc.ToCSR()
	if err != nil {
		t.Fatalf("ToCSR: %v", err)
	}
	if slices.Compare(csr.RowStart, round.RowStart) != 0 ||
		slices.Compare(csr.ColIdx, round.ColIdx) != 0 ||
		slices.Compare(csr.Values, round.Values) != 0 {
		t.Fatal("CSC to CSR round trip disagrees with the direct conversion")
	}
	back, err := csr.ToCSC()
	if err != nil {
		t.Fatalf("ToCSC: %v", err)
	}
	if slices.Compare(csc.ColStart, back.ColStart) != 0 ||
		slices.Compare(csc.RowIdx, back.RowIdx) != 0 ||
		slices.Compare(csc.Values, back.Values) != 0 {
		t.Fatal("CSR to CSC round trip disagrees with the direct conversion")
	}
	// MatVec over the CSC must answer what the CSR answers.
	x := floatsToArray([]float64{1, -2, 3}, []int{3})
	ycsr, err := csr.MatVec(x)
	if err != nil {
		t.Fatalf("CSR MatVec: %v", err)
	}
	ycsc, err := csc.MatVec(x)
	if err != nil {
		t.Fatalf("CSC MatVec: %v", err)
	}
	for i := range ycsr.Len() {
		if ycsr.FloatAt(i) != ycsc.FloatAt(i) {
			t.Fatalf("MatVec row %d: csr %.17g vs csc %.17g", i, ycsr.FloatAt(i), ycsc.FloatAt(i))
		}
	}
	if _, err := csc.MatVec(floatsToArray([]float64{1, 2}, []int{2})); err == nil {
		t.Fatal("a wrong-length vector was accepted")
	}
}

func mustInts(t *testing.T, vals []int64, shape ...int) *core.Array {
	t.Helper()
	a, err := core.FromInts(vals, shape...)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	return a
}