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

// randomDense builds a deterministic m×n matrix with entries in [−1, 1)
// from the seeded generator, so no test randomness leaks.
func randomDense(t *testing.T, m, n int, seed int64) *core.Array {
	t.Helper()
	g := core.NewGenerator(seed)
	vals := make([]float64, m*n)
	for i := range vals {
		vals[i] = float64(g.Next()%2000)/1000 - 1
	}
	return floatsToArray(vals, []int{m, n})
}

func TestRRQRFactorisation(t *testing.T) {
	const m, n = 14, 8
	a := randomDense(t, m, n, 11)
	q, r, perm, rank, err := RRQR(a)
	if err != nil {
		t.Fatalf("RRQR: %v", err)
	}
	if rank != n {
		t.Fatalf("rank %d for a full-rank %d×%d matrix", rank, m, n)
	}
	if !rrqrValidPerm(perm) {
		t.Fatalf("the permutation %v is not a permutation", perm)
	}
	// A·P = Q·R to machine precision.
	worst, aMax := 0.0, 0.0
	for i := range m {
		for j := range n {
			ap := a.FloatAt(i*n + perm[j])
			qr := 0.0
			for l := range m {
				qr += q.FloatAt(i*m+l) * r.FloatAt(l*n+j)
			}
			if d := math.Abs(ap - qr); d > worst {
				worst = d
			}
			if v := math.Abs(a.FloatAt(i*n + j)); v > aMax {
				aMax = v
			}
		}
	}
	if worst > 1e-12*math.Max(1, aMax)*float64(m) {
		t.Fatalf("A·P = Q·R off by %.3g", worst)
	}
	// Q orthogonal.
	for i := range m {
		for j := range m {
			s := 0.0
			for l := range m {
				s += q.FloatAt(l*m+i) * q.FloatAt(l*m+j)
			}
			want := 0.0
			if i == j {
				want = 1
			}
			if math.Abs(s-want) > 1e-13 {
				t.Fatalf("QᵀQ[%d,%d] = %.16g, want %.16g", i, j, s, want)
			}
		}
	}
	// R upper triangular.
	for i := range m {
		for j := range n {
			if j < i && math.Abs(r.FloatAt(i*n+j)) > 1e-12 {
				t.Fatalf("R[%d,%d] = %.3g below the diagonal", i, j, r.FloatAt(i*n+j))
			}
		}
	}
	// The |R| diagonal decays monotonically, the pivoting guarantee.
	for i := 1; i < n; i++ {
		hi := math.Abs(r.FloatAt(i*n + i))
		lo := math.Abs(r.FloatAt((i-1)*n + i - 1))
		if hi > lo*(1+1e-12) {
			t.Fatalf("|R[%d,%d]| = %.6g exceeds |R[%d,%d]| = %.6g", i, i, hi, i-1, i-1, lo)
		}
	}
	// RRQRRank is the same count without the orthogonal factor.
	rank2, err := RRQRRank(a)
	if err != nil {
		t.Fatalf("RRQRRank: %v", err)
	}
	if rank2 != n {
		t.Fatalf("RRQRRank = %d, want %d", rank2, n)
	}
}

func TestRRQRRankDetection(t *testing.T) {
	// A hand-built 6×4: three independent columns and a fourth that is
	// exactly 2·c0 + c1 in float arithmetic.
	const m, n = 6, 4
	g := core.NewGenerator(5)
	raw := make([]float64, m*n)
	for i := range m {
		raw[i*n+0] = float64(g.Next()%100)/50 - 1
		raw[i*n+1] = float64(g.Next()%100)/50 - 1
		raw[i*n+2] = float64(g.Next()%100)/50 - 1
		raw[i*n+3] = 2*raw[i*n+0] + raw[i*n+1]
	}
	a := floatsToArray(raw, []int{m, n})
	_, _, perm, rank, err := RRQR(a)
	if err != nil {
		t.Fatalf("RRQR: %v", err)
	}
	if rank != 3 {
		t.Fatalf("rank %d for a matrix whose fourth column is exactly dependent", rank)
	}
	if !rrqrValidPerm(perm) {
		t.Fatalf("the permutation %v is not a permutation", perm)
	}
	// The SVD count must agree, independently.
	svdRank, err := MatrixRank(a, 0)
	if err != nil {
		t.Fatalf("MatrixRank: %v", err)
	}
	if svdRank != 3 {
		t.Fatalf("the SVD counts rank %d where the pivoted diagonal shows 3", svdRank)
	}
	// An exact zero column is the cleanest possible direction.
	zero := floatsToArray([]float64{
		1, 0, 2, 0,
		2, 1, 1, 0,
		1, 1, 3, 0,
		2, 0, 1, 0,
		1, 2, 2, 0,
		3, 1, 4, 0,
	}, []int{m, n})
	rankZero, err := RRQRRank(zero)
	if err != nil {
		t.Fatalf("RRQRRank: %v", err)
	}
	if rankZero != 3 {
		t.Fatalf("rank %d for a matrix with an exact zero column", rankZero)
	}
	// The zero matrix has no rank at all.
	empty := floatsToArray(make([]float64, 3*2), []int{3, 2})
	rankEmpty, err := RRQRRank(empty)
	if err != nil {
		t.Fatalf("RRQRRank: %v", err)
	}
	if rankEmpty != 0 {
		t.Fatalf("rank %d for the zero matrix", rankEmpty)
	}
}

func TestSolveRRQRMinimumNorm(t *testing.T) {
	// Rank-deficient and consistent: the minimum-norm answer, checked
	// against the SVD's Pinverse and, on a hand-checkable 3×3, against
	// the exact (0, 1, 1).
	const m, n = 6, 4
	g := core.NewGenerator(9)
	raw := make([]float64, m*n)
	for i := range m {
		raw[i*n+0] = float64(g.Next()%100)/50 - 1
		raw[i*n+1] = float64(g.Next()%100)/50 - 1
		raw[i*n+2] = float64(g.Next()%100)/50 - 1
		raw[i*n+3] = 2*raw[i*n+0] + raw[i*n+1]
	}
	a := floatsToArray(raw, []int{m, n})
	xTrue := make([]float64, n)
	for i := range n {
		xTrue[i] = math.Cos(0.3*float64(i)) + float64(i%3)
	}
	b := core.New(core.Float, m)
	for i := range m {
		s := 0.0
		for j := range n {
			s += raw[i*n+j] * xTrue[j]
		}
		b.RawFloats()[i] = s
	}
	x, err := SolveRRQR(a, b)
	if err != nil {
		t.Fatalf("SolveRRQR: %v", err)
	}
	pinv, err := Pinverse(a, 0)
	if err != nil {
		t.Fatalf("Pinverse: %v", err)
	}
	worst := 0.0
	for i := range n {
		s := 0.0
		for j := range m {
			s += pinv.FloatAt(i*m+j) * b.FloatAt(j)
		}
		if d := math.Abs(x.FloatAt(i) - s); d > worst {
			worst = d
		}
	}
	if worst > 1e-8 {
		t.Fatalf("the pivoted answer misses the SVD minimum norm by %.3g", worst)
	}
	// The hand-checkable 3×3 system from the sparse solver's own pin.
	tiny := floatsToArray([]float64{1, 0, 1, 0, 1, 1, 1, 1, 2}, []int{3, 3})
	tb := mustFloats(t, []float64{1, 2, 3}, 3)
	xt, err := SolveRRQR(tiny, tb)
	if err != nil {
		t.Fatalf("SolveRRQR: %v", err)
	}
	want := []float64{0, 1, 1}
	for i := range want {
		if math.Abs(xt.FloatAt(i)-want[i]) > 1e-10 {
			t.Fatalf("x[%d] = %.12g, want the minimum-norm %.12g", i, xt.FloatAt(i), want[i])
		}
	}
	// Full rank and inconsistent: the ordinary least-squares answer.
	full := randomDense(t, 5, 3, 13)
	fb := mustFloats(t, []float64{1, -1, 2, 0.5, 3}, 5)
	xf, err := SolveRRQR(full, fb)
	if err != nil {
		t.Fatalf("SolveRRQR: %v", err)
	}
	ref, err := LeastSquares(full, fb)
	if err != nil {
		t.Fatalf("LeastSquares: %v", err)
	}
	for i := range 3 {
		if math.Abs(xf.FloatAt(i)-ref.FloatAt(i)) > 1e-9 {
			t.Fatalf("full-rank solve x[%d] = %.12g, want %.12g", i, xf.FloatAt(i), ref.FloatAt(i))
		}
	}
}

func TestSolveRRQRErrors(t *testing.T) {
	wide := floatsToArray([]float64{1, 2, 3, 4, 5, 6}, []int{2, 3})
	b2 := mustFloats(t, []float64{1, 2}, 2)
	if _, err := SolveRRQR(wide, b2); err == nil {
		t.Fatal("an underdetermined system was accepted")
	}
	complexA, err := core.FromComplexes([]complex128{1, 0, 0, 1}, 2, 2)
	if err != nil {
		t.Fatalf("FromComplexes: %v", err)
	}
	if _, err := SolveRRQR(complexA, b2); err == nil {
		t.Fatal("a complex matrix was accepted")
	}
	good := randomDense(t, 4, 3, 17)
	complexB, err := core.FromComplexes([]complex128{1, 2, 3}, 3)
	if err != nil {
		t.Fatalf("FromComplexes: %v", err)
	}
	if _, err := SolveRRQR(good, complexB); err == nil {
		t.Fatal("a complex right-hand side was accepted")
	}
	if _, err := SolveRRQR(good, mustFloats(t, []float64{1, 2}, 2)); err == nil {
		t.Fatal("a short right-hand side was accepted")
	}
	if _, err := SolveRRQR(good, core.New(core.Float, 2, 2)); err == nil {
		t.Fatal("a rank-2 right-hand side was accepted")
	}
}

func TestRRQRErrors(t *testing.T) {
	oneD := mustFloats(t, []float64{1, 2, 3}, 3)
	if _, _, _, _, err := RRQR(oneD); err == nil {
		t.Fatal("a rank-1 input was accepted")
	}
	complexA, err := core.FromComplexes([]complex128{1, 0, 0, 1}, 2, 2)
	if err != nil {
		t.Fatalf("FromComplexes: %v", err)
	}
	if _, _, _, _, err := RRQR(complexA); err == nil {
		t.Fatal("a complex matrix was accepted")
	}
	wide := floatsToArray([]float64{1, 2, 3, 4, 5, 6}, []int{2, 3})
	if _, _, _, _, err := RRQR(wide); err == nil {
		t.Fatal("an underdetermined matrix was accepted")
	}
	if _, err := RRQRRank(oneD); err == nil {
		t.Fatal("RRQRRank accepted a rank-1 input")
	}
	if _, err := RRQRRank(complexA); err == nil {
		t.Fatal("RRQRRank accepted a complex matrix")
	}
	if _, err := RRQRRank(wide); err == nil {
		t.Fatal("RRQRRank accepted an underdetermined matrix")
	}
	empty := floatsToArray([]float64{1}, []int{1, 0})
	if _, _, _, _, err := RRQR(empty); err == nil {
		t.Fatal("an empty matrix was accepted")
	}
	// The permutation helper's contract.
	if !rrqrValidPerm([]int{2, 0, 1}) {
		t.Fatal("a rotation was rejected as a permutation")
	}
	if rrqrValidPerm([]int{0, 0, 1}) {
		t.Fatal("a repeated index passed as a permutation")
	}
	if rrqrValidPerm([]int{0, 2}) {
		t.Fatal("an out-of-range index passed as a permutation")
	}
}