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

package stats

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

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

// pcaFixture builds a seeded (n, 4) observation array with correlated
// columns of clearly different scales, the ordinary material a PCA
// runs on.
func pcaFixture(t *testing.T, n int, seed int64) *core.Array {
	t.Helper()
	g := core.NewGenerator(31)
	x0 := make([]float64, n)
	vals := make([]float64, 0, 4*n)
	for i := range n {
		x0[i] = g.NormalUnit()
	}
	for i := range n {
		x1 := 0.8*x0[i] + 0.6*g.NormalUnit()
		x2 := -0.5*x0[i] + g.NormalUnit()
		x3 := 0.3 * g.NormalUnit()
		vals = append(vals, x0[i], x1, x2, x3)
	}
	return mustFromFloats(t, vals, n, 4)
}

// TestPCALongAxis puts a two-cluster anisotropic cloud under the
// decomposition: two Gaussian blobs strung along a thirty-degree axis
// must come back with the first component along that axis, the
// measured angle against the truth, and nearly all the variance on it.
func TestPCALongAxis(t *testing.T) {
	const n = 80
	g := core.NewGenerator(29)
	const theta = math.Pi / 6
	cos, sin := math.Cos(theta), math.Sin(theta)
	vals := make([]float64, 0, 4*n)
	for shift := 0.0; shift <= 10; shift += 10 {
		for range n {
			u := 5 * g.NormalUnit()
			v := 0.5 * g.NormalUnit()
			vals = append(vals,
				u*cos-v*sin+shift*cos,
				u*sin+v*cos+shift*sin)
		}
	}
	design := mustFromFloats(t, vals, 2*n, 2)
	res, err := PCA(design)
	if err != nil {
		t.Fatalf("PCA: %v", err)
	}
	// The angle of a component's axis, read off its two loadings and
	// folded into (−π/2, π/2], where the fixed sign convention leaves
	// it. Loadings entry (j, k) sits at j·p+k.
	angleOf := func(k int) float64 {
		a := math.Atan2(res.Loadings.FloatAt(1*2+k), res.Loadings.FloatAt(0*2+k))
		if a > math.Pi/2 {
			a -= math.Pi
		}
		if a <= -math.Pi/2 {
			a += math.Pi
		}
		return a
	}
	// Axis directions are defined only up to a half turn, so the
	// short axis's angle is folded like the first before comparing.
	first := angleOf(0)
	second := angleOf(1)
	secondWant := theta + math.Pi/2
	if secondWant > math.Pi/2 {
		secondWant -= math.Pi
	}
	t.Logf("first axis at %.4f rad against %.4f, second at %.4f against %.4f, explaining %.4f of the variance",
		first, theta, second, secondWant, res.ExplainedVarianceRatio[0])
	if math.Abs(first-theta) > 0.05 {
		t.Fatalf("the first component sits at %.4f rad, want the long axis %.4f", first, theta)
	}
	if math.Abs(second-secondWant) > 0.05 {
		t.Fatalf("the second component sits at %.4f rad, want the short axis %.4f", second, secondWant)
	}
	if res.ExplainedVarianceRatio[0] < 0.98 {
		t.Fatalf("the long axis explains only %.4f of the variance", res.ExplainedVarianceRatio[0])
	}
}

// TestPCAVariancesAndLoadings pins the spectral accounting: the
// eigenvalues sum to the covariance's trace, the ratios sum to one,
// they arrive in falling order, the loadings are orthonormal as
// columns, and the loadings and eigenvalues rebuild the covariance
// they came from.
func TestPCAVariancesAndLoadings(t *testing.T) {
	const n, p = 60, 4
	a := pcaFixture(t, n, 31)
	res, err := PCA(a)
	if err != nil {
		t.Fatalf("PCA: %v", err)
	}
	// The trace, computed here straight from the data.
	means := make([]float64, p)
	for j := range p {
		s := 0.0
		for i := range n {
			s += a.FloatAt(i*p + j)
		}
		means[j] = s / float64(n)
	}
	trace := 0.0
	for j := range p {
		s := 0.0
		for i := range n {
			d := a.FloatAt(i*p+j) - means[j]
			s += d * d
		}
		trace += s / float64(n-1)
	}
	total := 0.0
	for k := range p {
		total += res.ExplainedVariance[k]
	}
	if math.Abs(total-trace) > 1e-10*math.Max(1, trace) {
		t.Fatalf("the eigenvalues sum to %.12g, want the trace %.12g", total, trace)
	}
	ratioSum := 0.0
	for k := range p {
		ratioSum += res.ExplainedVarianceRatio[k]
		if k > 0 && res.ExplainedVariance[k] > res.ExplainedVariance[k-1] {
			t.Fatalf("the variances are not descending at %d", k)
		}
	}
	if math.Abs(ratioSum-1) > 1e-12 {
		t.Fatalf("the ratios sum to %.16g, want 1", ratioSum)
	}
	// Orthonormal columns: LᵀL is the identity.
	for j := range p {
		for k := j; k < p; k++ {
			dot := 0.0
			for i := range p {
				dot += res.Loadings.FloatAt(i*p+j) * res.Loadings.FloatAt(i*p+k)
			}
			want := 0.0
			if j == k {
				want = 1
			}
			if math.Abs(dot-want) > 1e-10 {
				t.Fatalf("loadings %d and %d have inner product %.4g, want %.4g", j, k, dot, want)
			}
		}
	}
	// The covariance rebuilds: L·D·Lᵀ against the entries the package
	// computed.
	cov, err := CovarianceMatrix(a)
	if err != nil {
		t.Fatalf("CovarianceMatrix: %v", err)
	}
	for i := range p {
		for j := range p {
			s := 0.0
			for k := range p {
				s += res.ExplainedVariance[k] * res.Loadings.FloatAt(i*p+k) * res.Loadings.FloatAt(j*p+k)
			}
			if math.Abs(s-cov.FloatAt(i*p+j)) > 1e-9*math.Max(1, math.Abs(cov.FloatAt(i*p+j))) {
				t.Fatalf("the rebuilt covariance entry (%d, %d) is %.12g, want %.12g", i, j, s, cov.FloatAt(i*p+j))
			}
		}
	}
}

// TestPCAScoresWhitenRoundTrip pins the transforms: the scores are
// the centred observations times the loadings, the whitened data is
// the scores scaled by the components' standard deviations, the
// covariance of the whitened data is the identity, and unwhitening
// returns the centred observations.
func TestPCAScoresWhitenRoundTrip(t *testing.T) {
	const n, p = 60, 4
	a := pcaFixture(t, n, 31)
	res, err := PCA(a)
	if err != nil {
		t.Fatalf("PCA: %v", err)
	}
	// Scores against their definition.
	for i := range n {
		for k := range p {
			s := 0.0
			for j := range p {
				s += (a.FloatAt(i*p+j) - res.Mean[j]) * res.Loadings.FloatAt(j*p+k)
			}
			if math.Abs(s-res.Scores.FloatAt(i*p+k)) > 1e-9 {
				t.Fatalf("score (%d, %d) is %.12g, want the centred row times the loading %.12g",
					i, k, res.Scores.FloatAt(i*p+k), s)
			}
		}
	}
	// Whitened against the scores scaled by the component scales, and
	// with identity covariance.
	z, err := res.Whiten(a)
	if err != nil {
		t.Fatalf("Whiten: %v", err)
	}
	for i := range n {
		for k := range p {
			want := res.Scores.FloatAt(i*p+k) / math.Sqrt(res.ExplainedVariance[k])
			if math.Abs(z.FloatAt(i*p+k)-want) > 1e-8 {
				t.Fatalf("whitened (%d, %d) is %.12g, want the scaled score %.12g",
					i, k, z.FloatAt(i*p+k), want)
			}
		}
	}
	zcov, err := CovarianceMatrix(z)
	if err != nil {
		t.Fatalf("CovarianceMatrix of the whitened data: %v", err)
	}
	for i := range p {
		for j := range p {
			want := 0.0
			if i == j {
				want = 1
			}
			if math.Abs(zcov.FloatAt(i*p+j)-want) > 1e-9 {
				t.Fatalf("the whitened covariance (%d, %d) is %.6g, want %.6g", i, j, zcov.FloatAt(i*p+j), want)
			}
		}
	}
	// Unwhitening returns the observations themselves: the round trip
	// closes exactly.
	back, err := res.Unwhiten(z)
	if err != nil {
		t.Fatalf("Unwhiten: %v", err)
	}
	for i := range n {
		for j := range p {
			if math.Abs(back.FloatAt(i*p+j)-a.FloatAt(i*p+j)) > 1e-9 {
				t.Fatalf("the round trip returned %.12g at (%d, %d), want the observation %.12g",
					back.FloatAt(i*p+j), i, j, a.FloatAt(i*p+j))
			}
		}
	}
}

// TestPCATransposedSpectrum pins the transform consistency across the
// transposed problem: for a square matrix centred along both axes the
// covariance of the rows and the covariance of the columns are the
// Gram pair XXᵀ and XᵀX, which share their spectrum exactly, and the
// decomposition must see the same eigenvalues from either side.
func TestPCATransposedSpectrum(t *testing.T) {
	const n = 6
	g := core.NewGenerator(37)
	vals := make([]float64, 0, n*n)
	for range n * n {
		vals = append(vals, g.NormalUnit())
	}
	// Centre along the columns and then along the rows, so both
	// readings of the matrix describe the same centred scatter.
	for i := range n {
		mean := 0.0
		for j := range n {
			mean += vals[i*n+j]
		}
		mean /= float64(n)
		for j := range n {
			vals[i*n+j] -= mean
		}
	}
	for j := range n {
		mean := 0.0
		for i := range n {
			mean += vals[i*n+j]
		}
		mean /= float64(n)
		for i := range n {
			vals[i*n+j] -= mean
		}
	}
	a := mustFromFloats(t, vals, n, n)
	transposed := make([]float64, 0, n*n)
	for i := range n {
		for j := range n {
			transposed = append(transposed, vals[j*n+i])
		}
	}
	at := mustFromFloats(t, transposed, n, n)
	res, err := PCA(a)
	if err != nil {
		t.Fatalf("PCA: %v", err)
	}
	resT, err := PCA(at)
	if err != nil {
		t.Fatalf("PCA of the transposed data: %v", err)
	}
	for k := range n {
		if math.Abs(res.ExplainedVariance[k]-resT.ExplainedVariance[k]) > 1e-8 {
			t.Fatalf("eigenvalue %d: %.10g from the rows, %.10g from the columns",
				k, res.ExplainedVariance[k], resT.ExplainedVariance[k])
		}
	}
}

// TestPCASignConvention pins the orientation rule directly on the
// helper: every eigenvector row turns so its largest-magnitude entry
// is positive, the first index winning a tie, and a row already
// oriented stays untouched.
func TestPCASignConvention(t *testing.T) {
	// Row 0 ties at 0.6 across indices 0 and 1, index 0 negative: the
	// first index wins, so the row flips. Row 1's largest entry is
	// −0.9: it flips. Row 2's largest entry is 0.7: it stays.
	v := [][]float64{
		{-0.6, 0.6, 0.1},
		{0.2, -0.9, 0.4},
		{0.1, 0.7, -0.2},
	}
	fixEigenSigns(v)
	want := [][]float64{
		{0.6, -0.6, -0.1},
		{-0.2, 0.9, -0.4},
		{0.1, 0.7, -0.2},
	}
	for i := range 3 {
		if !slices.Equal(v[i], want[i]) {
			t.Fatalf("orientation wrong at row %d: %v, want %v", i, v[i], want[i])
		}
	}
	// And on a real fit: every component's heaviest loading positive.
	a := pcaFixture(t, 40, 31)
	res, err := PCA(a)
	if err != nil {
		t.Fatalf("PCA: %v", err)
	}
	for k := range 4 {
		worst, index := 0.0, 0
		for i := range 4 {
			if magnitude := math.Abs(res.Loadings.FloatAt(i*4 + k)); magnitude > worst {
				worst = magnitude
				index = i
			}
		}
		if res.Loadings.FloatAt(index*4+k) < 0 {
			t.Fatalf("component %d is oriented against the convention", k)
		}
	}
}

// TestPCAValidation refuses the inputs without a decomposition and
// withholds the whitening transforms where they do not exist.
func TestPCAValidation(t *testing.T) {
	good := pcaFixture(t, 20, 31)
	if _, err := PCA(mustFromFloats(t, []float64{1, 2, 3}, 3)); err == nil || !strings.Contains(err.Error(), "2-D array") {
		t.Fatalf("a rank 1 array: got %v, want the rank refusal", err)
	}
	if _, err := PCA(mustFromFloats(t, []float64{1, 2}, 1, 2)); err == nil || !strings.Contains(err.Error(), "at least two observations") {
		t.Fatalf("a single observation: got %v, want the observation floor refusal", err)
	}
	if _, err := PCA(core.New(core.Complex, 4, 2)); err == nil || !strings.Contains(err.Error(), "complex observations") {
		t.Fatalf("complex observations: got %v, want the complex refusal", err)
	}
	if _, err := PCA(mustFromFloats(t, []float64{1, 2, math.NaN(), 4, 5, 6, 7, 8}, 4, 2)); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("non-finite observations: got %v, want the non-finite refusal", err)
	}
	constant := mustFromFloats(t, []float64{1, 2, 1, 2, 1, 2, 1, 2}, 4, 2)
	if _, err := PCA(constant); err == nil || !strings.Contains(err.Error(), "no variance") {
		t.Fatalf("data with no variance: got %v, want the variance refusal", err)
	}
	// A duplicated column: the decomposition stands, the whitening
	// transforms do not exist and are withheld.
	singular := mustFromFloats(t, []float64{
		1, 1, 2,
		2, 2, 1,
		3, 3, 0,
		4, 4, 1,
		5, 5, 2,
		6, 6, 3,
	}, 6, 3)
	res, err := PCA(singular)
	if err != nil {
		t.Fatalf("PCA on a singular covariance: %v", err)
	}
	if res.Whitening != nil || res.Unwhitening != nil {
		t.Fatalf("a rank-deficient fit published whitening transforms")
	}
	if _, err := res.Whiten(singular); err == nil || !strings.Contains(err.Error(), "rank deficient") {
		t.Fatalf("Whiten on a rank-deficient fit: got %v, want the rank-deficiency refusal", err)
	}
	if _, err := res.Unwhiten(singular); err == nil || !strings.Contains(err.Error(), "rank deficient") {
		t.Fatalf("Unwhiten on a rank-deficient fit: got %v, want the rank-deficiency refusal", err)
	}
	// Shape and content checks on the transforms.
	fit, err := PCA(good)
	if err != nil {
		t.Fatalf("PCA: %v", err)
	}
	if _, err := fit.Whiten(mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6}, 3, 2)); err == nil || !strings.Contains(err.Error(), "columns, the fit") {
		t.Fatalf("Whiten with a column mismatch: got %v, want the column refusal", err)
	}
	if _, err := fit.Whiten(mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6, 7, math.NaN(), 9, 10, 11, 12, 13, 14, 15, 16}, 4, 4)); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("Whiten with non-finite observations: got %v, want the non-finite refusal", err)
	}
	if _, err := fit.Unwhiten(mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6}, 3, 2)); err == nil || !strings.Contains(err.Error(), "columns, the fit") {
		t.Fatalf("Unwhiten with a column mismatch: got %v, want the column refusal", err)
	}
	var noFit *PCAResult
	if _, err := noFit.Whiten(good); err == nil || !strings.Contains(err.Error(), "no fit to whiten") {
		t.Fatalf("a nil fit whitened: got %v, want the nil-fit refusal", err)
	}
	if _, err := noFit.Unwhiten(good); err == nil || !strings.Contains(err.Error(), "no fit to unwhiten") {
		t.Fatalf("a nil fit unwhitened: got %v, want the nil-fit refusal", err)
	}
	complexData := core.New(core.Complex, 4, 4)
	if _, err := fit.Whiten(complexData); err == nil || !strings.Contains(err.Error(), "complex observations") {
		t.Fatalf("Whiten with complex observations: got %v, want the complex refusal", err)
	}
	if _, err := fit.Unwhiten(complexData); err == nil || !strings.Contains(err.Error(), "complex observations") {
		t.Fatalf("Unwhiten with complex observations: got %v, want the complex refusal", err)
	}
	// Integer observations reach the transforms through the widening
	// accessor instead of a raw float payload, and whiten back out
	// identically. Five rows of four generic columns keep the
	// covariance full rank.
	integers := mustFromInts(t, []int64{
		1, 2, 3, 4,
		2, 4, 6, 3,
		3, 6, 2, 9,
		4, 3, 8, 1,
		5, 7, 1, 2,
	}, 5, 4)
	intFit, err := PCA(integers)
	if err != nil {
		t.Fatalf("PCA on integer observations: %v", err)
	}
	z, err := intFit.Whiten(integers)
	if err != nil {
		t.Fatalf("Whiten on integer observations: %v", err)
	}
	if _, err := intFit.Unwhiten(z); err != nil {
		t.Fatalf("Unwhiten on integer observations: %v", err)
	}
}