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

// TestSVDReconstruction checks that A = U · Σ · Vᵀ reconstructs A
// for several rank profiles (full-rank square, tall, wide, and a
// rank-1 matrix where the one-sided Jacobi prototype failed).
func TestSVDReconstruction(t *testing.T) {
	cases := []struct {
		name    string
		vals    []float64
		m, n    int
		hasSwap bool
	}{
		{
			name: "square_full_rank",
			vals: []float64{1, 2, 3, 4, 5, 6, 7, 8, 9},
			m:    3, n: 3,
		},
		{
			name: "tall_rank2",
			vals: []float64{1, 1, 1, 2, 2, 2, 3, 3, 3, 1, 2, 3},
			m:    4, n: 3,
		},
		{
			name: "wide_rank2",
			vals: []float64{1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2},
			m:    3, n: 4,
			hasSwap: true,
		},
		{
			name: "rank1",
			vals: []float64{2, 4, 6, 8, 10, 12},
			m:    3, n: 2,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			a := mustFromFloats(t, tc.vals, tc.m, tc.n)
			uOut, sigma, vt, err := SVD(a)
			if err != nil {
				t.Fatalf("SVD: %v", err)
			}
			// Thin SVD: U is (m, k), Vᵀ is (k, n) where k = min(m, n).
			k := min(tc.m, tc.n)
			wantU := [2]int{tc.m, k}
			wantVt := [2]int{k, tc.n}
			if got := uOut.Shape(); got[0] != wantU[0] || got[1] != wantU[1] {
				t.Fatalf("U shape: got %v want %v", got, wantU)
			}
			if got := vt.Shape(); got[0] != wantVt[0] || got[1] != wantVt[1] {
				t.Fatalf("Vᵀ shape: got %v want %v", got, wantVt)
			}
			wantSigma := k
			if got := sigma.Shape()[0]; got != wantSigma {
				t.Fatalf("sigma shape: got %d want %d", got, wantSigma)
			}
			aRecon := matmulSVDReconstruct(uOut, sigma, vt, tc.m, tc.n)
			if !matricesClose(aRecon, denseFloats(a, tc.m, tc.n), tc.m, tc.n, 1e-9) {
				t.Fatalf("SVD reconstruction [%s]: max diff > 1e-9\norig=%v\nrecon=%v",
					tc.name, a, aRecon)
			}
			if !matricesThinOrthogonal(uOut.RawFloats(), uOut.Shape()[0], U_SVD_COLS(uOut), 1e-9) {
				t.Errorf("U not column-orthonormal in %s (shape=%v)", tc.name, uOut.Shape())
			}
			for i := 1; i < wantSigma; i++ {
				if sigma.RawFloats()[i-1] < sigma.RawFloats()[i] {
					t.Errorf("Σ not sorted descending: %v", sigma.RawFloats())
					break
				}
			}
		})
	}
}

// TestSVDRankDeficient checks the SVD on a deliberately rank-deficient
// matrix, the case the one-sided Jacobi prototype failed on.
func TestSVDRankDeficient(t *testing.T) {
	u := []float64{1, 2, 3, 4}
	v := []float64{1, -1, 1}
	a := outerProduct(u, v, 4, 3)
	uOut, sigma, vt, err := SVD(a)
	if err != nil {
		t.Fatalf("SVD: %v", err)
	}
	if len(sigma.RawFloats()) != 3 {
		t.Fatalf("σ shape: %v", sigma.Shape())
	}
	if !(sigma.RawFloats()[0] > 1e-9) {
		t.Errorf("σ₀: got %g, want > 1e-9", sigma.RawFloats()[0])
	}
	if math.Abs(sigma.RawFloats()[1]) > 1e-9 || math.Abs(sigma.RawFloats()[2]) > 1e-9 {
		t.Errorf("σ₁ and σ₂ should be ~0, got %g %g",
			sigma.RawFloats()[1], sigma.RawFloats()[2])
	}
	aRecon := matmulSVDReconstruct(uOut, sigma, vt, 4, 3)
	if !matricesClose(aRecon, denseFloats(a, 4, 3), 4, 3, 1e-9) {
		t.Fatalf("rank-deficient SVD reconstruction: max diff > 1e-9\norig=%v\nrecon=%v",
			a, aRecon)
	}
}

// TestEigenDiagonal checks that a diagonal matrix returns the
// diagonal entries as eigenvalues.
func TestEigenDiagonal(t *testing.T) {
	a := mustFromFloats(t, []float64{
		3, 0, 0,
		0, 1, 0,
		0, 0, 5,
	}, 3, 3)
	vals, vecs, err := Eigen(a)
	if err != nil {
		t.Fatalf("Eigen: %v", err)
	}
	if got := vals.Shape()[0]; got != 3 {
		t.Fatalf("vals shape: %v", vals.Shape())
	}
	want := []float64{1, 3, 5}
	for i, w := range want {
		if math.Abs(vals.RawFloats()[i]-w) > 1e-9 {
			t.Errorf("eigenvalue[%d]: got %g want %g", i, vals.RawFloats()[i], w)
		}
	}
	recon := eigenReconstruct(vecs, vals, 3)
	wantMat := denseFloats(a, 3, 3)
	if !matricesClose(recon, wantMat, 3, 3, 1e-9) {
		t.Fatalf("Eigen reconstruction: max diff > 1e-9\nwant=%v\ngot=%v",
			wantMat, recon)
	}
}

// TestEigenAsymmetric verifies the asymmetric input path errors.
func TestEigenAsymmetric(t *testing.T) {
	a := mustFromFloats(t, []float64{
		1, 2,
		3, 4,
	}, 2, 2)
	if _, _, err := Eigen(a); err == nil {
		t.Fatal("Eigen: expected error on asymmetric matrix")
	}
}

// TestPinverse checks Moore-Penrose pseudoinverse on the canonical
// example and verifies A · A⁺ · A = A on a rank-deficient matrix.
func TestPinverse(t *testing.T) {
	a := mustFromFloats(t, []float64{1, 2, 3, 4}, 2, 2)
	inv, err := Pinverse(a, 0)
	if err != nil {
		t.Fatalf("Pinverse: %v", err)
	}
	if !matricesClose(denseFloats(inv, 2, 2), []float64{-2, 1, 1.5, -0.5}, 2, 2, 1e-9) {
		t.Fatalf("Pinverse of invertible: got %v", inv)
	}
	u := []float64{1, 2, 3}
	v := []float64{2, -1}
	aRank1 := outerProduct(u, v, 3, 2)
	pinv, err := Pinverse(aRank1, 0)
	if err != nil {
		t.Fatalf("Pinverse(rank1): %v", err)
	}
	_ = pinv
	recon := matmulMul(matmulMul(denseFloats(aRank1, 3, 2), denseFloats(pinv, 2, 3), 3, 2, 3), denseFloats(aRank1, 3, 2), 3, 3, 2)
	want := denseFloats(aRank1, 3, 2)
	if !matricesClose(recon, want, 3, 2, 1e-9) {
		t.Fatalf("A · A⁺ · A != A for rank-1 input")
	}
}

// TestPinverseWideThinShapes pins the aspect-ratio plumbing of the
// pseudoinverse: a wide (m < n) input used to read the thin Vᵀ as an
// (n, n) matrix and panic past its payload. Values ride on the SVD
// track; the shapes must hold on their own.
func TestPinverseWideThinShapes(t *testing.T) {
	diag := mustFromFloats(t, []float64{
		3, 0, 0,
		0, 2, 0,
	}, 2, 3)
	inv, err := Pinverse(diag, 0)
	if err != nil {
		t.Fatalf("Pinverse wide diagonal: %v", err)
	}
	if inv.Shape()[0] != 3 || inv.Shape()[1] != 2 {
		t.Fatalf("Pinverse wide diagonal shape: %v, want (3, 2)", inv.Shape())
	}

	general := mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6}, 2, 3)
	inv2, err := Pinverse(general, 0)
	if err != nil {
		t.Fatalf("Pinverse wide general: %v", err)
	}
	if inv2.Shape()[0] != 3 || inv2.Shape()[1] != 2 {
		t.Fatalf("Pinverse wide general shape: %v, want (3, 2)", inv2.Shape())
	}
}

// TestPinverseWideDiagonalValues checks the wide-matrix values against
// the exact pseudoinverse once the SVD track lands.
func TestPinverseWideDiagonalValues(t *testing.T) {
	a := mustFromFloats(t, []float64{
		3, 0, 0,
		0, 2, 0,
	}, 2, 3)
	inv, err := Pinverse(a, 0)
	if err != nil {
		t.Fatalf("Pinverse wide: %v", err)
	}
	want := []float64{1.0 / 3, 0, 0, 0.5, 0, 0}
	got := denseFloats(inv, 3, 2)
	for i, w := range want {
		if diff := got[i] - w; diff > 1e-12 || diff < -1e-12 {
			t.Fatalf("Pinverse wide [%d]: got %v, want %v", i, got[i], w)
		}
	}
}

// TestMatrixRankProperties checks rank on a couple of cases.
func TestMatrixRankProperties(t *testing.T) {
	// Identity padded with a zero row: rows [1,0],[0,1],[0,0].
	a := mustFromFloats(t, []float64{1, 0, 0, 1, 0, 0}, 3, 2)
	if got, err := MatrixRank(a, 0); err != nil || got != 2 {
		t.Errorf("MatrixRank(full-rank): got %d err %v, want 2", got, err)
	}
	// Every row a multiple of [1,2]: rank 1.
	rank1 := mustFromFloats(t, []float64{1, 2, 2, 4, 3, 6}, 3, 2)
	if got, err := MatrixRank(rank1, 0); err != nil || got != 1 {
		t.Errorf("MatrixRank(rank-1): got %d err %v, want 1", got, err)
	}
}

// TestCondProperties checks condition number on a diagonal matrix
// (singular values are the diagonal entries).
func TestCondProperties(t *testing.T) {
	a := mustFromFloats(t, []float64{
		3, 0,
		0, 0.5,
	}, 2, 2)
	c, err := Cond(a, 0)
	if err != nil {
		t.Fatalf("Cond: %v", err)
	}
	if math.Abs(c-6) > 1e-9 {
		t.Errorf("Cond([[3, 0], [0, 0.5]]): got %g want 6", c)
	}
	sing := mustFromFloats(t, []float64{
		2, 0,
		0, 0,
	}, 2, 2)
	c, err = Cond(sing, 0)
	if err != nil {
		t.Fatalf("Cond(singular): %v", err)
	}
	if !math.IsInf(c, 1) {
		t.Errorf("Cond(singular): got %g want +Inf", c)
	}
}

// TestSVDWideTransposed checks that a wide matrix (m < n) decomposes
// correctly with the transpose swap.
func TestSVDWideTransposed(t *testing.T) {
	a := mustFromFloats(t, []float64{
		1, 2, 3, 4,
		5, 6, 7, 8,
		9, 10, 11, 12,
	}, 3, 4)
	uOut, sigma, vt, err := SVD(a)
	if err != nil {
		t.Fatalf("SVD: %v", err)
	}
	if uOut.Shape()[0] != 3 || uOut.Shape()[1] != 3 {
		t.Fatalf("U shape: got %v want (3, 3)", uOut.Shape())
	}
	if vt.Shape()[0] != 3 || vt.Shape()[1] != 4 {
		t.Fatalf("Vᵀ shape: got %v want (3, 4)", vt.Shape())
	}
	if sigma.Shape()[0] != 3 {
		t.Fatalf("Σ shape: got %v want (3,)", sigma.Shape())
	}
	recon := matmulSVDReconstruct(uOut, sigma, vt, 3, 4)
	if !matricesClose(recon, denseFloats(a, 3, 4), 3, 4, 1e-9) {
		t.Fatalf("wide SVD reconstruction: max diff > 1e-9")
	}
}

// --- helpers (test-only) ---

// U_SVD_COLS returns the number of columns of a thin SVD U
// (min(m, n)), the rank-n bound the orthogonal columns reach.
func U_SVD_COLS(u *core.Array) int {
	return u.Shape()[1]
}

func matmulSVDReconstruct(u, sigma, vt *core.Array, m, n int) []float64 {
	// Thin SVD: U is (m, k); Σ is (k,); Vᵀ is (k, n) where k = min(m, n).
	k := min(m, n)
	uMat := denseFloats(u, m, k)
	vtMat := denseFloats(vt, k, n)
	sVals := sigma.RawFloats()
	out := make([]float64, m*n)
	for i := range m {
		for j := range n {
			s := 0.0
			for kk := range k {
				s += uMat[i*k+kk] * sVals[kk] * vtMat[kk*n+j]
			}
			out[i*n+j] = s
		}
	}
	return out
}

func eigenReconstruct(q, vals *core.Array, n int) []float64 {
	qMat := denseFloats(q, n, n)
	out := make([]float64, n*n)
	for i := range n {
		for j := range n {
			s := 0.0
			for k := range n {
				s += qMat[i*n+k] * vals.RawFloats()[k] * qMat[j*n+k]
			}
			out[i*n+j] = s
		}
	}
	return out
}

func outerProduct(u, v []float64, m, n int) *core.Array {
	out := make([]float64, m*n)
	for i := range m {
		for j := range n {
			out[i*n+j] = u[i] * v[j]
		}
	}
	return floatsToArray(out, []int{m, n})
}

func matmulMul(a, b []float64, rows, mid, cols int) []float64 {
	out := make([]float64, rows*cols)
	for i := range rows {
		for j := range cols {
			s := 0.0
			for k := range mid {
				s += a[i*mid+k] * b[k*cols+j]
			}
			out[i*cols+j] = s
		}
	}
	return out
}

func matricesClose(a, b []float64, m, n int, tol float64) bool {
	if len(a) != m*n || len(b) != m*n {
		return false
	}
	for i := 0; i < m*n; i++ {
		if math.Abs(a[i]-b[i]) > tol {
			return false
		}
	}
	return true
}

// matricesThinOrthogonal checks that a (m, n) matrix a satisfies
// aᵀ a = I_n. Used for the thin SVD factor U.
func matricesThinOrthogonal(a []float64, m, n int, tol float64) bool {
	for i := range n {
		for j := range n {
			s := 0.0
			want := 0.0
			if i == j {
				want = 1
			}
			for k := range m {
				s += a[k*n+i] * a[k*n+j]
			}
			if math.Abs(s-want) > tol {
				return false
			}
		}
	}
	return true
}

// TestLeastSquaresCollinear pins the rank guard: exactly collinear
// columns are refused, and near-collinear ones whose R pivot sits at
// the rounding floor of max|R| are refused too, where the old
// exact-zero check returned a silently huge x.
func TestLeastSquaresCollinear(t *testing.T) {
	// Exactly collinear, integer dtype: the historical case.
	ints, err := core.FromInts([]int64{1, 2, 2, 4}, 2, 2)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	if _, err := LeastSquares(ints, mustFloats(t, []float64{3, 7})); err == nil {
		t.Fatal("exactly collinear integer columns: want an error")
	}
	// Near-collinear float columns: the second column differs from the
	// first by 2^-45, so its R pivot (2^-45) sits below the relative
	// rank floor n·eps·max|R| = 2·eps·2^10, where back-substitution
	// would amplify it into a ~1e13 answer. Powers of two keep every
	// elimination step exact, so the trip is deterministic.
	a := mustFloats(t, []float64{
		1024, 1024,
		0, math.Ldexp(1, -45),
		0, 0,
	}, 3, 2)
	if _, err := LeastSquares(a, mustFloats(t, []float64{1, 1, 1})); err == nil {
		t.Fatal("near-collinear columns: want an error")
	}
	// A well-conditioned system of the same shape still solves.
	good := mustFloats(t, []float64{
		1, 1,
		1, 2,
		1, 3,
	}, 3, 2)
	x, err := LeastSquares(good, mustFloats(t, []float64{2, 3, 4}))
	if err != nil {
		t.Fatalf("LeastSquares full-rank: %v", err)
	}
	if math.Abs(x.FloatAt(0)-1) > 1e-9 || math.Abs(x.FloatAt(1)-1) > 1e-9 {
		t.Fatalf("x = (%.12g, %.12g), want (1, 1)", x.FloatAt(0), x.FloatAt(1))
	}
}

// TestCondZeroMatrix pins the zero-matrix contract: no usable inverse
// direction, so the condition number is +Inf, not 0.
func TestCondZeroMatrix(t *testing.T) {
	c, err := Cond(mustFloats(t, []float64{0, 0, 0, 0}, 2, 2), 0)
	if err != nil {
		t.Fatalf("Cond: %v", err)
	}
	if !math.IsInf(c, 1) {
		t.Fatalf("Cond(zero matrix) = %g, want +Inf", c)
	}
}

// TestCholeskyUpdateRejectsUpperTriangle pins the factor validation:
// a nonzero strict upper triangle is refused, matching what the error
// message has always claimed.
func TestCholeskyUpdateRejectsUpperTriangle(t *testing.T) {
	l := mustFloats(t, []float64{
		2, 1,
		0, 1,
	}, 2, 2)
	if _, err := CholeskyUpdate(l, mustFloats(t, []float64{1, 1})); err == nil {
		t.Fatal("CholeskyUpdate with a nonzero upper triangle: want an error")
	}
	if _, err := CholeskyDowndate(l, mustFloats(t, []float64{1, 1})); err == nil {
		t.Fatal("CholeskyDowndate with a nonzero upper triangle: want an error")
	}
}

// TestHouseholderVectorIntoLargeScale pins the overflow-safe reflector:
// entries near 1e154 keep a finite norm and beta instead of squaring
// their way to +Inf.
func TestHouseholderVectorIntoLargeScale(t *testing.T) {
	x := []float64{1e154, 1e154, 1e154}
	dst := make([]float64, 3)
	hh := householderVectorInto(dst, x)
	if math.IsInf(hh.beta, 0) || math.IsNaN(hh.beta) {
		t.Fatalf("beta = %v, want finite", hh.beta)
	}
	for i, v := range hh.v {
		if math.IsInf(v, 0) || math.IsNaN(v) {
			t.Fatalf("v[%d] = %v, want finite", i, v)
		}
	}
	// The zero vector still answers a zero-beta identity reflector.
	empty := householderVectorInto(make([]float64, 2), []float64{0, 0})
	if empty.beta != 0 {
		t.Fatalf("beta = %v for a zero vector, want 0", empty.beta)
	}
}

// TestEigenRefusesNonFinite pins that a poisoned matrix never reads as
// symmetric: the guard's comparison cannot see a NaN difference, so the
// entry is refused outright, the way the sparse sibling's symmetry
// check refuses it.
func TestEigenRefusesNonFinite(t *testing.T) {
	nan := mustFromFloats(t, []float64{math.NaN(), 0, 0, 1}, 2, 2)
	if _, _, err := Eigen(nan); err == nil || !strings.Contains(err.Error(), "not finite") {
		t.Fatalf("Eigen(NaN): %v", err)
	}
	inf := mustFromFloats(t, []float64{math.Inf(1), 0, 0, 1}, 2, 2)
	if _, _, err := Eigen(inf); err == nil || !strings.Contains(err.Error(), "not finite") {
		t.Fatalf("Eigen(Inf): %v", err)
	}
}