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

// overdeterminedLSFixture builds a deterministic sparse overdetermined
// system: thirty rows, twelve columns, three stored entries per row,
// a right-hand side taken from a known x through the matrix plus a
// small inconsistent part, and its dense twin for the reference
// solvers.
func overdeterminedLSFixture(t *testing.T) (coo *core.SparseCOO, b *core.Array, dense *core.Array) {
	t.Helper()
	const m, n = 30, 12
	g := core.NewGenerator(7)
	idx := make([]int64, 0, 3*m+2)
	vals := make([]float64, 0, 3*m+2)
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	for i := range m {
		add(i, i%n, 1+float64(g.Next()%100)/200)
		add(i, (i*7+3)%n, -1+float64(g.Next()%100)/100)
		add(i, (i*13+5)%n, float64(g.Next()%100)/100)
	}
	add(0, 0, 3)
	add(1, 2, 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{m, n})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	csr, err := CSRFromCOO(coo)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	xTrue := make([]float64, n)
	for i := range n {
		xTrue[i] = math.Sin(0.7*float64(i)) + float64(i%5)*0.3 - 0.6
	}
	b, err = csr.MatVec(floatsToArray(xTrue, []int{n}))
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	for i := range m {
		b.RawFloats()[i] += 0.01 * math.Cos(float64(i))
	}
	denseVals := make([]float64, m*n)
	for i := range len(vals) {
		denseVals[int(idx[2*i])*n+int(idx[2*i+1])] += vals[i]
	}
	return coo, b, floatsToArray(denseVals, []int{m, n})
}

// lsAchievedNorms recomputes ‖b − A·x‖ and ‖Aᵀ(b − A·x)‖ through the
// public sparse surface, independently of the solver's own operator.
func lsAchievedNorms(t *testing.T, coo *core.SparseCOO, b, x *core.Array) (rNorm, arNorm float64) {
	t.Helper()
	csr, err := CSRFromCOO(coo)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	ax, err := csr.MatVec(x)
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	r := core.New(core.Float, b.Len())
	for i := range b.Len() {
		r.RawFloats()[i] = b.FloatAt(i) - ax.FloatAt(i)
	}
	rNorm = norm2F64(vectorF64(r, b.Len()))
	at, err := csr.Transpose().MatVec(r)
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	arNorm = norm2F64(vectorF64(at, at.Len()))
	return rNorm, arNorm
}

func TestSpLSQRMatchesDenseLeastSquares(t *testing.T) {
	coo, b, dense := overdeterminedLSFixture(t)
	ref, err := LeastSquares(dense, b)
	if err != nil {
		t.Fatalf("LeastSquares: %v", err)
	}
	x, info, err := SpLSQR(coo, b, 1e-12, 0, 0)
	if err != nil {
		t.Fatalf("SpLSQR: %v", err)
	}
	worst, refMax := 0.0, 0.0
	for i := range ref.Len() {
		if d := math.Abs(x.FloatAt(i) - ref.FloatAt(i)); d > worst {
			worst = d
		}
		if v := math.Abs(ref.FloatAt(i)); v > refMax {
			refMax = v
		}
	}
	if worst > 1e-9*refMax {
		t.Fatalf("LSQR disagrees with the dense solve by %.3g (scale %.3g)", worst, refMax)
	}
	if !info.Converged {
		t.Fatalf("LSQR stopped without convergence: %+v", info)
	}
	if info.Criterion != LeastSquaresResidual && info.Criterion != LeastSquaresNormal {
		t.Fatalf("criterion %q is not a residual test", info.Criterion)
	}
	if info.Iterations > 2*12 {
		t.Fatalf("LSQR took %d steps for a 12-column Krylov space", info.Iterations)
	}
	// The achieved quantities must be the explicit truth, not the
	// in-loop estimates.
	rNorm, arNorm := lsAchievedNorms(t, coo, b, x)
	if math.Abs(rNorm-info.ResidualNorm) > 1e-9*info.ResidualNorm {
		t.Fatalf("reported residual %.3g does not match the explicit %.3g", info.ResidualNorm, rNorm)
	}
	if math.Abs(arNorm-info.NormalResidual) > 1e-9*info.NormalResidual {
		t.Fatalf("reported normal residual %.3g does not match the explicit %.3g", info.NormalResidual, arNorm)
	}
}

func TestSpLSMRMatchesDenseLeastSquares(t *testing.T) {
	coo, b, dense := overdeterminedLSFixture(t)
	ref, err := LeastSquares(dense, b)
	if err != nil {
		t.Fatalf("LeastSquares: %v", err)
	}
	x, info, err := SpLSMR(coo, b, 1e-12, 0, 0)
	if err != nil {
		t.Fatalf("SpLSMR: %v", err)
	}
	worst, refMax := 0.0, 0.0
	for i := range ref.Len() {
		if d := math.Abs(x.FloatAt(i) - ref.FloatAt(i)); d > worst {
			worst = d
		}
		if v := math.Abs(ref.FloatAt(i)); v > refMax {
			refMax = v
		}
	}
	if worst > 1e-9*refMax {
		t.Fatalf("LSMR disagrees with the dense solve by %.3g (scale %.3g)", worst, refMax)
	}
	if !info.Converged || info.Criterion == "" {
		t.Fatalf("LSMR stopped without convergence: %+v", info)
	}
	// LSMR's own minimisation target moves monotonically.
	for i := 1; i < len(info.normalEstimates); i++ {
		if info.normalEstimates[i] > info.normalEstimates[i-1] {
			t.Fatalf("LSMR normal-equations estimate rose at step %d: %.6g after %.6g",
				i, info.normalEstimates[i], info.normalEstimates[i-1])
		}
	}
}

func TestSpLSQREstimateTrajectoriesMonotone(t *testing.T) {
	coo, b, _ := overdeterminedLSFixture(t)
	// LSQR's residual estimate is |φ̄|, shrank every step by |sn| ≤ 1;
	// LSMR's normal estimate is |ζ̄|, shrank by |s̄| ≤ 1. Both are
	// monotone by construction and must stay so in float.
	_, lsInfo, err := SpLSQR(coo, b, 1e-12, 0, 0)
	if err != nil {
		t.Fatalf("SpLSQR: %v", err)
	}
	for i := 1; i < len(lsInfo.residualEstimates); i++ {
		if lsInfo.residualEstimates[i] > lsInfo.residualEstimates[i-1] {
			t.Fatalf("LSQR residual estimate rose at step %d: %.6g after %.6g",
				i, lsInfo.residualEstimates[i], lsInfo.residualEstimates[i-1])
		}
	}
	_, lsmrInfo, err := SpLSMR(coo, b, 1e-12, 0, 0)
	if err != nil {
		t.Fatalf("SpLSMR: %v", err)
	}
	for i := 1; i < len(lsmrInfo.normalEstimates); i++ {
		if lsmrInfo.normalEstimates[i] > lsmrInfo.normalEstimates[i-1] {
			t.Fatalf("LSMR normal estimate rose at step %d: %.6g after %.6g",
				i, lsmrInfo.normalEstimates[i], lsmrInfo.normalEstimates[i-1])
		}
	}
}

// TestSpLeastSquaresMinimumNorm pins the hand-checkable rank-deficient
// consistent system A = [[1,0,1],[0,1,1],[1,1,2]], b = (1,2,3): the
// solution family is (1−t, 2−t, t) and the minimum-norm member is
// (0, 1, 1), which both solvers must answer from a zero start.
func TestSpLeastSquaresMinimumNorm(t *testing.T) {
	indices, err := core.FromInts([]int64{0, 0, 0, 2, 1, 1, 1, 2, 2, 0, 2, 1, 2, 2}, 7, 2)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	coo, err := core.NewSparseCOO(indices, floatsToArray([]float64{1, 1, 1, 1, 1, 1, 2}, []int{7}), []int{3, 3})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	csr, err := CSRFromCOO(coo)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	b, err := csr.MatVec(floatsToArray([]float64{1, 2, 0}, []int{3}))
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	want := []float64{0, 1, 1}
	for name, run := range map[string]func(*core.SparseCOO, *core.Array) (*core.Array, *LeastSquaresInfo, error){
		"LSQR": func(a *core.SparseCOO, bb *core.Array) (*core.Array, *LeastSquaresInfo, error) {
			return SpLSQR(a, bb, 1e-13, 0, 0)
		},
		"LSMR": func(a *core.SparseCOO, bb *core.Array) (*core.Array, *LeastSquaresInfo, error) {
			return SpLSMR(a, bb, 1e-13, 0, 0)
		},
	} {
		x, info, err := run(coo, b)
		if err != nil {
			t.Fatalf("%s: %v", name, err)
		}
		for i := range want {
			if math.Abs(x.FloatAt(i)-want[i]) > 1e-10 {
				t.Fatalf("%s: x[%d] = %.12g, want the minimum-norm %.12g", name, i, x.FloatAt(i), want[i])
			}
		}
		if !info.Converged {
			t.Fatalf("%s: not converged: %+v", name, info)
		}
	}
	// The same answer must agree with the SVD's minimum-norm solution.
	dense := floatsToArray([]float64{1, 0, 1, 0, 1, 1, 1, 1, 2}, []int{3, 3})
	pinv, err := Pinverse(dense, 0)
	if err != nil {
		t.Fatalf("Pinverse: %v", err)
	}
	ref := core.New(core.Float, 3)
	for i := range 3 {
		s := 0.0
		for j := range 3 {
			s += pinv.FloatAt(i*3+j) * b.FloatAt(j)
		}
		ref.RawFloats()[i] = s
	}
	for i := range 3 {
		if math.Abs(ref.FloatAt(i)-want[i]) > 1e-12 {
			t.Fatalf("Pinverse reference %.12g disagrees with the hand solution %.12g", ref.FloatAt(i), want[i])
		}
	}
}

// TestSpLeastSquaresIllConditioned pins convergence with the criterion
// recorded on a system whose diagonal decays by four orders: the
// conditional estimate stays under the limit and a residual test
// ends the iteration.
func TestSpLeastSquaresIllConditioned(t *testing.T) {
	const n = 10
	idx := make([]int64, 0, 3*n)
	vals := make([]float64, 0, 3*n)
	for i := range n {
		idx = append(idx, int64(i), int64(i))
		vals = append(vals, math.Pow(10, -0.45*float64(i)))
		if i+1 < n {
			idx = append(idx, int64(i), int64(i+1), int64(i+1), int64(i))
			vals = append(vals, 1e-7, 1e-7)
		}
	}
	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)
	}
	csr, err := CSRFromCOO(coo)
	if err != nil {
		t.Fatalf("CSRFromCOO: %v", err)
	}
	b, err := csr.MatVec(floatsToArray([]float64{1, 1, 1, 1, 1, 1, 1, 1, 1, 1}, []int{n}))
	if err != nil {
		t.Fatalf("MatVec: %v", err)
	}
	for name, run := range map[string]func(*core.SparseCOO, *core.Array) (*core.Array, *LeastSquaresInfo, error){
		"LSQR": func(a *core.SparseCOO, bb *core.Array) (*core.Array, *LeastSquaresInfo, error) {
			return SpLSQR(a, bb, 1e-9, 0, 0)
		},
		"LSMR": func(a *core.SparseCOO, bb *core.Array) (*core.Array, *LeastSquaresInfo, error) {
			return SpLSMR(a, bb, 1e-9, 0, 0)
		},
	} {
		x, info, err := run(coo, b)
		if err != nil {
			t.Fatalf("%s: %v", name, err)
		}
		if !info.Converged || info.Criterion == "" {
			t.Fatalf("%s: ill-conditioned system ended without a recorded criterion: %+v", name, info)
		}
		if info.Condition <= 0 || info.MatrixNorm <= 0 {
			t.Fatalf("%s: estimates not reported: %+v", name, info)
		}
		rNorm, _ := lsAchievedNorms(t, coo, b, x)
		if rNorm > 1e-8 {
			t.Fatalf("%s: achieved residual %.3g is too coarse", name, rNorm)
		}
	}
}

func TestSpLeastSquaresErrors(t *testing.T) {
	coo, b, _ := overdeterminedLSFixture(t)
	// Underdetermined systems are refused, in the dense surface's own
	// words.
	small, err := core.NewSparseCOO(
		mustInts(t, []int64{0, 0, 0, 1}, 2, 2),
		floatsToArray([]float64{1, 1}, []int{2}),
		[]int{2, 3})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	for name, run := range map[string]func(*core.SparseCOO, *core.Array) (*core.Array, *LeastSquaresInfo, error){
		"LSQR": func(a *core.SparseCOO, bb *core.Array) (*core.Array, *LeastSquaresInfo, error) {
			return SpLSQR(a, bb, 0, 0, 0)
		},
		"LSMR": func(a *core.SparseCOO, bb *core.Array) (*core.Array, *LeastSquaresInfo, error) {
			return SpLSMR(a, bb, 0, 0, 0)
		},
	} {
		if _, _, err := run(small, mustFloats(t, []float64{1, 2}, 2)); err == nil || !strings.Contains(err.Error(), "overdetermined") {
			t.Fatalf("%s: underdetermined system accepted: %v", name, err)
		}
	}
	// Complex inputs.
	complexValues, err := core.FromComplexes([]complex128{1}, 1)
	if err != nil {
		t.Fatalf("FromComplexes: %v", err)
	}
	complexCOO, err := core.NewSparseCOO(mustInts(t, []int64{0, 0}, 1, 2), complexValues, []int{1, 1})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	if _, _, err := SpLSQR(complexCOO, mustFloats(t, []float64{1}, 1), 0, 0, 0); err == nil {
		t.Fatal("a complex matrix was accepted")
	}
	if _, _, err := SpLSQR(coo, mustFloats(t, []float64{1, 2}, 2), 0, 0, 0); err == nil {
		t.Fatal("a short right-hand side was accepted")
	}
	if _, _, err := SpLSQR(coo, core.New(core.Float, 3, 3), 0, 0, 0); err == nil {
		t.Fatal("a rank-2 right-hand side was accepted")
	}
	// A budget that runs out with every tolerance unmet is the budget's
	// own fault: the error names the steps, not a drift no estimate
	// committed. A tolerance of 1e-14 with two steps fires no criterion
	// (the residual is still 1.59), and neither does 1e-300 in one.
	if _, _, err := SpLSQR(coo, b, 1e-14, 2, 0); err == nil ||
		!strings.Contains(err.Error(), "no convergence in 2 steps") ||
		strings.Contains(err.Error(), "drifted") {
		t.Fatalf("exhausted budget misreported: %v", err)
	}
	if _, _, err := SpLSQR(coo, b, 1e-300, 1, 0); err == nil ||
		!strings.Contains(err.Error(), "no convergence in 1 steps") ||
		strings.Contains(err.Error(), "drifted") {
		t.Fatalf("exhausted budget misreported: %v", err)
	}
	if _, _, err := SpLSMR(coo, b, 1e-300, 1, 0); err == nil ||
		!strings.Contains(err.Error(), "no convergence in 1 steps") ||
		strings.Contains(err.Error(), "drifted") {
		t.Fatalf("exhausted budget misreported: %v", err)
	}
	// A zero right-hand side answers the exact zero without a step.
	zero := floatsToArray(make([]float64, b.Len()), []int{b.Len()})
	x, info, err := SpLSQR(coo, zero, 0, 0, 0)
	if err != nil {
		t.Fatalf("zero right-hand side: %v", err)
	}
	for i := range x.Len() {
		if x.FloatAt(i) != 0 {
			t.Fatalf("zero right-hand side answered %.3g", x.FloatAt(i))
		}
	}
	if info.Criterion != "" || !info.Converged || info.Iterations != 0 {
		t.Fatalf("zero right-hand side info: %+v", info)
	}
	// A non-empty matrix with no column-space component of b answers
	// the exact zero as well: both columns lie along (1,0) and
	// b = (0,1) is orthogonal to the column space.
	null, err := core.NewSparseCOO(mustInts(t, []int64{0, 0, 0, 1}, 2, 2), floatsToArray([]float64{1, 2}, []int{2}), []int{2, 2})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	ortho := mustFloats(t, []float64{0, 1}, 2)
	x, info, err = SpLSQR(null, ortho, 0, 0, 0)
	if err != nil {
		t.Fatalf("orthogonal right-hand side: %v", err)
	}
	for i := range x.Len() {
		if x.FloatAt(i) != 0 {
			t.Fatalf("orthogonal right-hand side answered %.3g", x.FloatAt(i))
		}
	}
	if info.ResidualNorm <= 0 {
		t.Fatalf("the achieved residual of an orthogonal system vanished: %+v", info)
	}
}

// TestFinishLeastSquaresDriftRefusal pins the drift guard at its own
// gate: a fired criterion whose recomputed norms do not carry it is
// refused with the empty string, while the same answer at a tolerance
// it genuinely meets is carried. The end-to-end budget pins live in
// TestSpLeastSquaresErrors; the drift itself needs the estimate and
// the truth to disagree, which is settled here without a recursion.
func TestFinishLeastSquaresDriftRefusal(t *testing.T) {
	coo, err := core.NewSparseCOO(mustInts(t, []int64{0, 0}, 1, 2), mustFloats(t, []float64{1}, 1), []int{1, 1})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	op, err := newLSQROperator(coo)
	if err != nil {
		t.Fatalf("newLSQROperator: %v", err)
	}
	// x = 0.5 against A = [1], b = [1] leaves both the residual and
	// the normal residual at 0.5 with ‖b‖ = 1.
	if got := finishLeastSquares("SpLSQR", LeastSquaresResidual, false, 3, op, []float64{1}, []float64{0.5}, 1, 1, 1e-14); got != "" {
		t.Fatalf("a drifted criterion carried as %q", got)
	}
	if got := finishLeastSquares("SpLSQR", LeastSquaresResidual, false, 3, op, []float64{1}, []float64{0.5}, 1, 1, 0.4); got != LeastSquaresResidual {
		t.Fatalf("an honestly met criterion refused: %q", got)
	}
}

// TestSpLSQREstimateSeriesIdentities pins what each estimate series
// carries, not merely that it shrinks. LSQR's residual estimate tracks
// ‖b − A·x‖ and its normal-equations estimate tracks ‖Aᵀ(b − A·x)‖, so
// on a converged fit the first ends beside the explicitly recomputed
// info.ResidualNorm while the second sits orders of magnitude below it.
// Filling the residual series with the normal estimate, or the reverse,
// keeps both series non-increasing and fails here.
func TestSpLSQREstimateSeriesIdentities(t *testing.T) {
	coo, b, _ := overdeterminedLSFixture(t)
	_, info, err := SpLSQR(coo, b, 1e-12, 0, 0)
	if err != nil {
		t.Fatalf("SpLSQR: %v", err)
	}
	re, ne := info.residualEstimates, info.normalEstimates
	if len(re) != info.Iterations || len(ne) != info.Iterations {
		t.Fatalf("estimate series of %d and %d against %d recorded steps", len(re), len(ne), info.Iterations)
	}
	if len(re) == 0 {
		t.Fatal("no estimates recorded")
	}
	lastRe, lastNe := re[len(re)-1], ne[len(ne)-1]
	if dev := math.Abs(lastRe - info.ResidualNorm); dev > 0.5*info.ResidualNorm {
		t.Fatalf("the residual estimate ends at %.6g against the explicit residual %.6g, want the two on the same scale",
			lastRe, info.ResidualNorm)
	}
	if lastNe > 1e-3*info.ResidualNorm {
		t.Fatalf("the normal-equations estimate ends at %.6g against a residual of %.6g, want the normal residual's own, much smaller scale",
			lastNe, info.ResidualNorm)
	}
	if lastRe == lastNe {
		t.Fatalf("both series end at %.6g, want the residual and the normal-equations estimates", lastRe)
	}
}