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

package stats

import (
	"math"
	"strings"
	"testing"

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

// TestGPInterpolatesNoiselessTrainingPoints fits a noiseless GP on its
// own training points: the posterior must reproduce every observation
// to rounding and carry no variance there.
func TestGPInterpolatesNoiselessTrainingPoints(t *testing.T) {
	g := core.NewGenerator(41)
	const n = 9
	train := core.New(core.Float, n, 1)
	y := core.New(core.Float, n)
	for i := range n {
		x := -2 + 0.5*float64(i)
		train.RawFloats()[i] = x
		y.RawFloats()[i] = math.Sin(x) + 0.1*x*g.Unit()
	}
	kernel, err := SquaredExponentialKernel(0.7)
	if err != nil {
		t.Fatalf("SquaredExponentialKernel: %v", err)
	}
	res, err := GaussianProcessRegression(kernel, train, y, 0, train)
	if err != nil {
		t.Fatalf("GaussianProcessRegression: %v", err)
	}
	for i := range n {
		if math.Abs(res.Mean[i]-y.FloatAt(i)) > 1e-9 {
			t.Fatalf("point %d: posterior mean %.12g against the observation %.12g",
				i, res.Mean[i], y.FloatAt(i))
		}
		if res.Variance[i] > 1e-12 {
			t.Fatalf("point %d: posterior variance %.3g at a noiseless training point", i, res.Variance[i])
		}
	}
	// The full posterior covariance vanishes on the training set too:
	// the conditioning has removed everything the prior had there.
	for i := range n {
		for j := range n {
			if math.Abs(res.Covariance[i*n+j]) > 1e-10 {
				t.Fatalf("posterior covariance (%d, %d) = %.3g at the training points, want zero",
					i, j, res.Covariance[i*n+j])
			}
		}
	}
	// The reported marginal likelihood must match the standalone
	// function bit for bit on the same inputs.
	mll, err := MarginalLogLikelihood(kernel, train, y, 0)
	if err != nil {
		t.Fatalf("MarginalLogLikelihood: %v", err)
	}
	if mll != res.LogLikelihood {
		t.Fatalf("the fit reported %.17g but the standalone %.17g", res.LogLikelihood, mll)
	}
}

// TestGPHugeLengthScaleFollowsGlobalMean pins the limiting behaviour
// of a giant length scale: when the prior cannot tell neighbouring
// inputs apart, the posterior mean collapses onto the constant that
// the likelihood alone supports, the sample mean of the training
// responses. A small noise keeps the Gram matrix well conditioned;
// with n = 30 and noise 1e-6 the constant fit sits within 1e-7 of the
// mean, far inside the 1e-4 tolerance.
func TestGPHugeLengthScaleFollowsGlobalMean(t *testing.T) {
	g := core.NewGenerator(43)
	const n = 30
	train := core.New(core.Float, n, 1)
	y := core.New(core.Float, n)
	total := 0.0
	for i := range n {
		train.RawFloats()[i] = float64(i) / float64(n-1)
		y.RawFloats()[i] = 3 + 0.4*g.NormalUnit()
		total += y.FloatAt(i)
	}
	meanY := total / n
	kernel, err := SquaredExponentialKernel(1e6)
	if err != nil {
		t.Fatalf("SquaredExponentialKernel: %v", err)
	}
	test := mustFloats(t, []float64{-5, 0.37, 12}, 3, 1)
	res, err := GaussianProcessRegression(kernel, train, y, 1e-6, test)
	if err != nil {
		t.Fatalf("GaussianProcessRegression: %v", err)
	}
	for i, m := range res.Mean {
		if math.Abs(m-meanY) > 1e-4 {
			t.Fatalf("test point %d: posterior mean %.8g against the global mean %.8g", i, m, meanY)
		}
	}
}

// TestGPVarianceFarFromDataEqualsPrior pins the uncertainty
// behaviour: a test point many length scales from every observation
// has learned nothing, so its posterior variance returns to the prior
// variance of 1 and its posterior covariance with the data region
// vanishes.
func TestGPVarianceFarFromDataEqualsPrior(t *testing.T) {
	const n = 6
	train := mustFloats(t, []float64{0, 0.2, 0.4, 0.6, 0.8, 1.0}, n, 1)
	y := mustFloats(t, []float64{1, -1, 0.5, -0.5, 2, -2}, n)
	kernel, err := SquaredExponentialKernel(0.3)
	if err != nil {
		t.Fatalf("SquaredExponentialKernel: %v", err)
	}
	// 30 is one hundred length scales from the nearest observation.
	test := mustFloats(t, []float64{0.5, 30}, 2, 1)
	res, err := GaussianProcessRegression(kernel, train, y, 0, test)
	if err != nil {
		t.Fatalf("GaussianProcessRegression: %v", err)
	}
	if v := res.Variance[1]; math.Abs(v-1) > 1e-10 {
		t.Fatalf("far posterior variance %.15g, want the prior 1", v)
	}
	// The far point's covariance with the data region: the cross
	// covariances are underflows, so the off-diagonal entry must be a
	// rounding dust.
	if c := res.Covariance[0*2+1]; math.Abs(c) > 1e-10 {
		t.Fatalf("far covariance with the data region %.3g, want a rounding dust", c)
	}
}

// TestMatern32MatchesSpectralForm holds the Matérn 3/2 closed form
// against its defining spectral equation: the density
// S(ω) = 4a³/(a² + ω²)², a = √3/L, inverts to the kernel through
// k(r) = (1/π)∫₀^∞ S(ω)·cos(ω·r) dω. The integral is evaluated by
// trapezoid out to five hundred decay widths, where the integrand's
// ω⁻⁴ tail has less than 1e-7 left to give.
func TestMatern32MatchesSpectralForm(t *testing.T) {
	const lengthScale = 1.3
	kernel, err := Matern32Kernel(lengthScale)
	if err != nil {
		t.Fatalf("Matern32Kernel: %v", err)
	}
	a := math.Sqrt(3) / lengthScale
	const omegaMax = 500
	const steps = 400000
	h := omegaMax / float64(steps)
	spectral := func(r float64) float64 {
		total := 0.0
		for s := range steps + 1 {
			omega := float64(s) * h
			w := 1.0
			if s == 0 || s == steps {
				w = 0.5
			}
			den := a*a + omega*omega
			total += w * 4 * a * a * a / (den * den) * math.Cos(omega*r)
		}
		return total * h / math.Pi
	}
	for _, r := range []float64{0, 0.7, lengthScale, 2.9} {
		want := kernel.Covariance([]float64{r}, []float64{0})
		got := spectral(r)
		if math.Abs(got-want) > 1e-6 {
			t.Fatalf("r = %g: closed form %.9g against the spectral inversion %.9g", r, want, got)
		}
	}
	// And the spot values of both Matérns: k(0) = 1 and the documented
	// closed forms at one length scale, r = L, where the shape
	// parameter a = √3 (respectively √5) exactly.
	m52, err := Matern52Kernel(lengthScale)
	if err != nil {
		t.Fatalf("Matern52Kernel: %v", err)
	}
	if kernel.Covariance([]float64{0}, []float64{0}) != 1 {
		t.Fatal("the Matern 3/2 kernel is not of unit amplitude")
	}
	if m52.Covariance([]float64{0}, []float64{0}) != 1 {
		t.Fatal("the Matern 5/2 kernel is not of unit amplitude")
	}
	a32 := math.Sqrt(3)
	want32 := (1 + a32) * math.Exp(-a32)
	if math.Abs(kernel.Covariance([]float64{lengthScale}, []float64{0})-want32) > 1e-12 {
		t.Fatalf("Matern 3/2 at r = L: %.15g, want %.15g",
			kernel.Covariance([]float64{lengthScale}, []float64{0}), want32)
	}
	a52 := math.Sqrt(5)
	want52 := (1 + a52 + a52*a52/3) * math.Exp(-a52)
	if math.Abs(m52.Covariance([]float64{lengthScale}, []float64{0})-want52) > 1e-12 {
		t.Fatalf("Matern 5/2 at r = L: %.15g, want %.15g",
			m52.Covariance([]float64{lengthScale}, []float64{0}), want52)
	}
}

// TestPeriodicKernelRepeats pins the period: a full period apart the
// kernel returns to 1, half a period apart with a short length scale
// it has decorrelated to rounding dust.
func TestPeriodicKernelRepeats(t *testing.T) {
	kernel, err := PeriodicKernel(0.2, 3)
	if err != nil {
		t.Fatalf("PeriodicKernel: %v", err)
	}
	if got := kernel.Covariance([]float64{7}, []float64{7 + 3}); math.Abs(got-1) > 1e-12 {
		t.Fatalf("one period apart the covariance is %.15g, want 1", got)
	}
	if got := kernel.Covariance([]float64{0}, []float64{1.5}); got > 1e-6 {
		t.Fatalf("half a period apart with a short scale the covariance is %.3g", got)
	}
}

// TestMarginalLogLikelihoodPrefersTrueHyperparameters draws one path
// of a Gaussian process with known hyperparameters over the house
// generator and requires the marginal likelihood to rank the true pair
// above badly mismatched ones, with a margin far beyond the rounding.
func TestMarginalLogLikelihoodPrefersTrueHyperparameters(t *testing.T) {
	g := core.NewGenerator(53)
	const n = 30
	const trueScale = 1.0
	const trueNoise = 0.05
	train := core.New(core.Float, n, 1)
	for i := range n {
		train.RawFloats()[i] = 5 * float64(i) / float64(n-1)
	}
	// The path: one multivariate normal draw over the grid under the
	// true kernel, plus the observation noise.
	trueKernel, err := SquaredExponentialKernel(trueScale)
	if err != nil {
		t.Fatalf("SquaredExponentialKernel: %v", err)
	}
	cov := core.New(core.Float, n, n)
	for i := range n {
		for j := range n {
			cov.RawFloats()[i*n+j] = trueKernel.Covariance(train.RawFloats()[i:i+1], train.RawFloats()[j:j+1])
		}
		// A hair of jitter on the diagonal: the Gram matrix of a long
		// length scale is positive definite by a margin that float64
		// rounding can eat on the way down, and the draw is fixture
		// construction, not the model, whose noise variance is fitted
		// separately below.
		cov.RawFloats()[i*n+i] += 1e-9
	}
	paths, err := MultivariateNormalDraws(g, 1, core.New(core.Float, n), cov)
	if err != nil {
		t.Fatalf("MultivariateNormalDraws: %v", err)
	}
	y := core.New(core.Float, n)
	for i := range n {
		y.RawFloats()[i] = paths.FloatAt(i) + trueNoise*g.NormalUnit()
	}
	truth, err := MarginalLogLikelihood(trueKernel, train, y, trueNoise)
	if err != nil {
		t.Fatalf("MarginalLogLikelihood: %v", err)
	}
	mismatched := []struct {
		scale float64
		noise float64
	}{
		{0.05, 5},
		{10, 1e-4},
	}
	for _, pair := range mismatched {
		wrongKernel, kerr := SquaredExponentialKernel(pair.scale)
		if kerr != nil {
			t.Fatalf("SquaredExponentialKernel(%g): %v", pair.scale, kerr)
		}
		wrong, err := MarginalLogLikelihood(wrongKernel, train, y, pair.noise)
		if err != nil {
			t.Fatalf("MarginalLogLikelihood at scale %g: %v", pair.scale, err)
		}
		if truth < wrong+5 {
			t.Fatalf("the true pair scored %.4g against (%g, %g) at %.4g, the margin is under 5",
				truth, pair.scale, pair.noise, wrong)
		}
	}
}

// TestGPRegressionMultiDimensional fits a three-dimensional input and
// requires the same exact interpolation, the path the Euclidean
// distance of every kernel takes over columns.
func TestGPRegressionMultiDimensional(t *testing.T) {
	g := core.NewGenerator(59)
	const n = 8
	train := core.New(core.Float, n, 3)
	y := core.New(core.Float, n)
	for i := range n {
		for j := range 3 {
			train.RawFloats()[i*3+j] = g.Unit()
		}
		y.RawFloats()[i] = g.NormalUnit()
	}
	kernel, err := Matern52Kernel(1.5)
	if err != nil {
		t.Fatalf("Matern52Kernel: %v", err)
	}
	res, err := GaussianProcessRegression(kernel, train, y, 0, train)
	if err != nil {
		t.Fatalf("GaussianProcessRegression: %v", err)
	}
	for i := range n {
		if math.Abs(res.Mean[i]-y.FloatAt(i)) > 1e-8 {
			t.Fatalf("point %d: mean %.10g against the observation %.10g", i, res.Mean[i], y.FloatAt(i))
		}
		if res.Variance[i] > 1e-10 {
			t.Fatalf("point %d: variance %.3g at a training point", i, res.Variance[i])
		}
	}
	// The posterior covariance is symmetric by construction; the
	// mirror must hold.
	for i := range n {
		for j := range n {
			if res.Covariance[i*n+j] != res.Covariance[j*n+i] {
				t.Fatalf("the posterior covariance is not symmetric at (%d, %d)", i, j)
			}
		}
	}
}

// TestGPValidationAndKernelRefusals checks the constructors and the
// entry points refuse what they must, including the singular Gram
// matrix that duplicated noiseless rows name.
func TestGPValidationAndKernelRefusals(t *testing.T) {
	kernel, kerr := SquaredExponentialKernel(1)
	if kerr != nil {
		t.Fatalf("SquaredExponentialKernel: %v", kerr)
	}
	train := mustFloats(t, []float64{0, 1, 2}, 3, 1)
	y := mustFloats(t, []float64{1, -1, 0.5}, 3)
	test := mustFloats(t, []float64{0.5}, 1, 1)

	for _, scale := range []float64{0, -1, math.Inf(1), math.NaN()} {
		if _, err := SquaredExponentialKernel(scale); err == nil {
			t.Fatalf("squared exponential accepted the scale %v", scale)
		}
		if _, err := Matern32Kernel(scale); err == nil {
			t.Fatalf("Matern 3/2 accepted the scale %v", scale)
		}
		if _, err := Matern52Kernel(scale); err == nil {
			t.Fatalf("Matern 5/2 accepted the scale %v", scale)
		}
		if _, err := PeriodicKernel(scale, 1); err == nil {
			t.Fatalf("periodic kernel accepted the length scale %v", scale)
		}
		if _, err := PeriodicKernel(1, scale); err == nil {
			t.Fatalf("periodic kernel accepted the period %v", scale)
		}
	}
	if _, err := GaussianProcessRegression(nil, train, y, 0, test); err == nil || !strings.Contains(err.Error(), "kernel is nil") {
		t.Fatalf("nil kernel: got %v, want the nil-kernel refusal", err)
	}
	for _, noise := range []float64{-0.1, math.Inf(-1), math.NaN()} {
		if _, err := GaussianProcessRegression(kernel, train, y, noise, test); err == nil || !strings.Contains(err.Error(), "noise variance") {
			t.Fatalf("noise variance %v: got %v, want the noise-variance refusal", noise, err)
		}
		if _, err := MarginalLogLikelihood(kernel, train, y, noise); err == nil || !strings.Contains(err.Error(), "noise variance") {
			t.Fatalf("marginal likelihood, noise variance %v: got %v, want the noise-variance refusal", noise, err)
		}
	}
	if _, err := GaussianProcessRegression(kernel, core.New(core.Float, 3), y, 0, test); err == nil || !strings.Contains(err.Error(), "must be rank 2") {
		t.Fatalf("rank-1 training design: got %v, want the rank refusal", err)
	}
	if _, err := GaussianProcessRegression(kernel, train, core.New(core.Float, 3, 1), 0, test); err == nil || !strings.Contains(err.Error(), "must be rank 1") {
		t.Fatalf("rank-2 response: got %v, want the rank refusal", err)
	}
	if _, err := GaussianProcessRegression(kernel, train, mustFloats(t, []float64{1, -1}, 2), 0, test); err == nil || !strings.Contains(err.Error(), "rows but the response") {
		t.Fatalf("short response: got %v, want the length refusal", err)
	}
	if _, err := GaussianProcessRegression(kernel, train, y, 0, core.New(core.Float, 2)); err == nil || !strings.Contains(err.Error(), "must be rank 2") {
		t.Fatalf("rank-1 test design: got %v, want the rank refusal", err)
	}
	wide := mustFloats(t, []float64{0.1, 0.2, 0.3, 0.4}, 2, 2)
	if _, err := GaussianProcessRegression(kernel, train, y, 0, wide); err == nil || !strings.Contains(err.Error(), "wide but the test design") {
		t.Fatalf("test design of the wrong width: got %v, want the width refusal", err)
	}
	sick := core.New(core.Float, 3, 1)
	sick.RawFloats()[2] = math.NaN()
	if _, err := GaussianProcessRegression(kernel, sick, y, 0, test); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("non-finite training design: got %v, want the non-finite refusal", err)
	}
	if _, err := GaussianProcessRegression(kernel, train, mustFloats(t, []float64{1, -1, math.Inf(1)}, 3), 0, test); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("non-finite response: got %v, want the non-finite refusal", err)
	}
	// Duplicated rows with no noise: the Gram matrix is singular and
	// the refusal names the row.
	dup := mustFloats(t, []float64{1, 1, 2}, 3, 1)
	_, err := GaussianProcessRegression(kernel, dup, y, 0, test)
	if err == nil || !strings.Contains(err.Error(), "positive definite") {
		t.Fatalf("duplicated noiseless rows: got %v, want a positive-definiteness refusal", err)
	}
	// A zero-column design has nothing to evaluate.
	if _, err := GaussianProcessRegression(kernel, core.New(core.Float, 3, 0), y, 0, test); err == nil || !strings.Contains(err.Error(), "at least one column") {
		t.Fatalf("zero-column design: got %v, want the column refusal", err)
	}
	if _, err := GaussianProcessRegression(kernel, core.New(core.Complex, 3, 1), y, 0, test); err == nil || !strings.Contains(err.Error(), "complex") {
		t.Fatalf("complex training design: got %v, want the complex refusal", err)
	}
	if _, err := GaussianProcessRegression(kernel, train, core.New(core.Complex, 3), 0, test); err == nil || !strings.Contains(err.Error(), "complex") {
		t.Fatalf("complex response: got %v, want the complex refusal", err)
	}
	// Integer inputs take the widening accessor path end to end:
	// design, response and the likelihood assembly.
	intTrain, ierr := core.FromInts([]int64{0, 10, 20}, 3, 1)
	if ierr != nil {
		t.Fatalf("FromInts: %v", ierr)
	}
	intY, ierr := core.FromInts([]int64{1, -2, 4}, 3)
	if ierr != nil {
		t.Fatalf("FromInts: %v", ierr)
	}
	intRes, err := GaussianProcessRegression(kernel, intTrain, intY, 0, intTrain)
	if err != nil {
		t.Fatalf("GaussianProcessRegression over int inputs: %v", err)
	}
	for i := range 3 {
		if math.Abs(intRes.Mean[i]-intY.FloatAt(i)) > 1e-8 {
			t.Fatalf("int input point %d: mean %.10g against the observation %.10g",
				i, intRes.Mean[i], intY.FloatAt(i))
		}
	}
	intMLL, err := MarginalLogLikelihood(kernel, intTrain, intY, 0)
	if err != nil {
		t.Fatalf("MarginalLogLikelihood over int inputs: %v", err)
	}
	if intMLL != intRes.LogLikelihood {
		t.Fatalf("the int-input likelihoods disagree: %.17g against %.17g", intRes.LogLikelihood, intMLL)
	}
}