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

package linalg

import (
	"math"
	"strings"
	"testing"

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

// unitVectorCOO builds a coordinate vector with a single non-zero.
func unitVector(t *testing.T, n, pos int, val float64) *core.Array {
	t.Helper()
	vals := make([]float64, n)
	vals[pos] = val
	return floatsToArray(vals, []int{n})
}

// denseSparseNorm returns the largest absolute difference of two
// factor states: the diagonal and the stored values in order.
func factorDelta(f, g *SparseCholesky) float64 {
	worst := 0.0
	for i := range f.n {
		if d := math.Abs(f.diag[i] - g.diag[i]); d > worst {
			worst = d
		}
	}
	for i := range f.values {
		if d := math.Abs(f.values[i] - g.values[i]); d > worst {
			worst = d
		}
	}
	return worst
}

// addRankOneCOO returns the coordinate matrix of A + x·xᵀ, x given in
// the stored coordinates, by summing the modified entries on top of
// the original ones.
func addRankOneCOO(t *testing.T, a *core.SparseCOO, x *core.Array) *core.SparseCOO {
	t.Helper()
	nnz := a.Indices.Shape()[0]
	entries := make(map[[2]int]float64)
	for i := range nnz {
		r := int(a.Indices.RawInts()[i*2])
		c := int(a.Indices.RawInts()[i*2+1])
		entries[[2]int{r, c}] += a.Values.FloatAt(i)
	}
	for i := range x.Len() {
		if x.FloatAt(i) == 0 {
			continue
		}
		for j := range x.Len() {
			if x.FloatAt(j) == 0 {
				continue
			}
			entries[[2]int{i, j}] += x.FloatAt(i) * x.FloatAt(j)
		}
	}
	idx := make([]int64, 0, 2*len(entries))
	vals := make([]float64, 0, len(entries))
	for k, v := range entries {
		idx = append(idx, int64(k[0]), int64(k[1]))
		vals = append(vals, v)
	}
	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)}), a.Shape)
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	return coo
}

// TestSparseCholeskyUpdateDowndateRoundTrip pins the round trip on the
// banded case the contract names: a path-graph Laplacian, whose
// natural and reverse Cuthill-McKee factors are both banded, with
// supported updates whose sweeps stay on the stored pattern. Update
// then downdate must restore the factor, and the updated factor must
// solve like a fresh factorisation of the modified matrix.
func TestSparseCholeskyUpdateDowndateRoundTrip(t *testing.T) {
	coo := gridLaplacianCOO(t, 12, 1)
	n := 12
	csr, err := CSRFromCOO(coo)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	b, err := csr.MatVec(floatsToArray([]float64{1, -2, 3, -4, 5, -1, 2, -3, 4, -5, 1, 0}, []int{n}))
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	for _, ordering := range []SparseOrdering{SparseOrderingNatural, SparseOrderingReverseCuthillMcKee} {
		for name, x := range map[string]*core.Array{
			"single":   unitVector(t, n, 3, 0.5),
			"adjacent": floatsToArray([]float64{0, 0, 0, 0, 0, 0, 0, 0.5, 0.25, 0, 0, 0}, []int{n}),
		} {
			f, err := NewSparseCholesky(coo, ordering)
			if err != nil {
				t.Fatalf("NewSparseCholesky: %v", err)
			}
			diagBefore := append([]float64(nil), f.diag...)
			valsBefore := append([]float64(nil), f.values...)
			if err := f.Update(x); err != nil {
				t.Fatalf("%s ordering %d: update: %v", name, ordering, err)
			}
			// The updated factor solves like the refactored matrix.
			refactor, err := NewSparseCholesky(addRankOneCOO(t, coo, x), ordering)
			if err != nil {
				t.Fatalf("refactor: %v", err)
			}
			sModified, err := f.Solve(b)
			if err != nil {
				t.Fatalf("solve on the updated factor: %v", err)
			}
			sRefactor, err := refactor.Solve(b)
			if err != nil {
				t.Fatalf("solve on the refactored matrix: %v", err)
			}
			scale := 0.0
			for i := range n {
				if v := math.Abs(sRefactor.FloatAt(i)); v > scale {
					scale = v
				}
			}
			for i := range n {
				if math.Abs(sModified.FloatAt(i)-sRefactor.FloatAt(i)) > 1e-9*scale {
					t.Fatalf("%s ordering %d: updated solve disagrees with the refactored one at %d",
						name, ordering, i)
				}
			}
			// Downdate restores the factor to round-off.
			if err := f.Downdate(x); err != nil {
				t.Fatalf("%s ordering %d: downdate: %v", name, ordering, err)
			}
			diagAfter := append([]float64(nil), f.diag...)
			valsAfter := append([]float64(nil), f.values...)
			worst := 0.0
			for i := range diagBefore {
				worst = math.Max(worst, math.Abs(diagAfter[i]-diagBefore[i]))
			}
			for i := range valsBefore {
				worst = math.Max(worst, math.Abs(valsAfter[i]-valsBefore[i]))
			}
			if worst > 1e-12 {
				t.Fatalf("%s ordering %d: round trip lost %.3g on the stored factor", name, ordering, worst)
			}
		}
	}
}

// TestSparseCholeskyUpdateMatchesRefactorisation pins the numeric
// truth of the update: on the two-dimensional grid with a
// single-coordinate update the sweep stays on the pattern, and every
// stored entry must equal the fresh factorisation of the explicitly
// modified matrix, whose pattern the diagonal change cannot alter.
func TestSparseCholeskyUpdateMatchesRefactorisation(t *testing.T) {
	coo := gridLaplacianCOO(t, 8, 8)
	n := 64
	x := unitVector(t, n, 30, 0.5)
	for _, ordering := range []SparseOrdering{SparseOrderingNatural, SparseOrderingReverseCuthillMcKee} {
		f, err := NewSparseCholesky(coo, ordering)
		if err != nil {
			t.Fatalf("NewSparseCholesky: %v", err)
		}
		if err := f.Update(x); err != nil {
			t.Fatalf("update: %v", err)
		}
		refactor, err := NewSparseCholesky(addRankOneCOO(t, coo, x), ordering)
		if err != nil {
			t.Fatalf("refactor: %v", err)
		}
		if f.NNZ() != refactor.NNZ() {
			t.Fatalf("the update changed the fill: %d vs %d", f.NNZ(), refactor.NNZ())
		}
		if d := factorDelta(f, refactor); d > 1e-9 {
			t.Fatalf("the updated factor differs from the refactored one by %.3g", d)
		}
	}
}

// TestSparseCholeskyUpdateDowndateExact pins bit-for-bit restoration
// on the engineered pattern-stable example: a banded matrix whose
// stored factor carries the integers 3 and 5, updated by the
// Pythagorean vector 4·e₀, where hypot(3, 4) = 5 and the rotations
// round exactly. Update turns every 3 into a 5 and every 5 into a 3;
// the downdate turns them back without losing a bit.
func TestSparseCholeskyUpdateDowndateExact(t *testing.T) {
	const n = 6
	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 {
		if i == 0 {
			add(0, 0, 9)
		} else {
			add(i, i, 34)
		}
		if i+1 < n {
			add(i, i+1, 15)
			add(i+1, i, 15)
		}
	}
	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)
	}
	f, err := NewSparseCholesky(coo, SparseOrderingNatural)
	if err != nil {
		t.Fatalf("NewSparseCholesky: %v", err)
	}
	diagBefore := append([]float64(nil), f.diag...)
	valsBefore := append([]float64(nil), f.values...)
	for i := range diagBefore {
		if diagBefore[i] != 3 {
			t.Fatalf("the engineered factor's diagonal is %.4g, want 3", diagBefore[i])
		}
	}
	x := unitVector(t, n, 0, 4)
	if err := f.Update(x); err != nil {
		t.Fatalf("update: %v", err)
	}
	for i := range f.diag {
		if f.diag[i] != 5 {
			t.Fatalf("after the update diag[%d] = %.17g, want exactly 5", i, f.diag[i])
		}
	}
	for i := range f.values {
		if f.values[i] != 3 {
			t.Fatalf("after the update values[%d] = %.17g, want exactly 3", i, f.values[i])
		}
	}
	if err := f.Downdate(x); err != nil {
		t.Fatalf("downdate: %v", err)
	}
	for i := range f.diag {
		if f.diag[i] != diagBefore[i] {
			t.Fatalf("the downdate lost diag[%d]: %.17g vs %.17g", i, f.diag[i], diagBefore[i])
		}
	}
	for i := range f.values {
		if f.values[i] != valsBefore[i] {
			t.Fatalf("the downdate lost values[%d]: %.17g vs %.17g", i, f.values[i], valsBefore[i])
		}
	}
}

func TestSparseCholeskyUpdateRefusals(t *testing.T) {
	coo := gridLaplacianCOO(t, 12, 1)
	// The update whose support spreads beyond a stored column needs
	// fill the pattern does not hold: refused, factor untouched.
	f, err := NewSparseCholesky(coo, SparseOrderingNatural)
	if err != nil {
		t.Fatalf("NewSparseCholesky: %v", err)
	}
	before := append([]float64(nil), f.values...)
	spread := unitVector(t, 12, 0, 1)
	spread.RawFloats()[5] = 1
	err = f.Update(spread)
	if err == nil || !strings.Contains(err.Error(), "pattern does not hold") {
		t.Fatalf("a pattern-violating update was accepted: %v", err)
	}
	for i := range f.values {
		if f.values[i] != before[i] {
			t.Fatalf("the refused update moved values[%d]", i)
		}
	}
	// The downdate that loses positive dominance: 3·e₀ against a unit
	// diagonal.
	if err := f.Downdate(unitVector(t, 12, 0, 3)); err == nil || !strings.Contains(err.Error(), "positive definite") {
		t.Fatalf("a downdate outside the cone was accepted: %v", err)
	}
	// A downdate whose square leaves the float64 range.
	huge, err := core.NewSparseCOO(mustInts(t, []int64{0, 0}, 1, 2), floatsToArray([]float64{1e308}, []int{1}), []int{1, 1})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	hf, err := NewSparseCholesky(huge, SparseOrderingNatural)
	if err != nil {
		t.Fatalf("NewSparseCholesky: %v", err)
	}
	if err := hf.Update(unitVector(t, 1, 0, 1e154)); err != nil {
		t.Fatalf("huge update: %v", err)
	}
	if err := hf.Downdate(unitVector(t, 1, 0, 1e154)); err == nil || !strings.Contains(err.Error(), "range") {
		t.Fatalf("an out-of-range downdate was accepted: %v", err)
	}
	// Validation.
	if err := f.Update(unitVector(t, 11, 0, 1)); err == nil {
		t.Fatal("a short vector was accepted")
	}
	if err := f.Downdate(core.New(core.Float, 3, 4)); err == nil {
		t.Fatal("a rank-2 vector was accepted")
	}
	if err := f.Update(core.New(core.Complex, 12)); err == nil {
		t.Fatal("a complex vector was accepted")
	}
	// A non-finite entry.
	nan := unitVector(t, 12, 2, 1)
	nan.RawFloats()[2] = math.NaN()
	if err := f.Update(nan); err == nil || !strings.Contains(err.Error(), "finite") {
		t.Fatalf("a NaN entry was accepted: %v", err)
	}
}