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

// Regression-edge pins for the linear model: the statistics of a
// design without an intercept column, the tail of the coefficient
// p-values, and the exact-fit report.

// fact returns n!, small n only.
func fact(n int) float64 {
	r := 1.0
	for i := 2; i <= n; i++ {
		r *= float64(i)
	}
	return r
}

// binom returns the binomial coefficient, small n only.
func binom(n, k int) float64 { return fact(n) / (fact(k) * fact(n-k)) }

// tTailExactEven returns P(|T| > t) for even df in closed form. There,
// the incomplete beta has a = df/2, an integer, and b = 1/2, so the
// integral is elementary: nothing is approximated and nothing from the
// library is used.
func tTailExactEven(tv float64, df int) float64 {
	m := df / 2
	z := float64(df) / (float64(df) + tv*tv)
	// ∫_0^z u^(m−1)(1−u)^(−1/2) du, with u = 1 − s², is
	// 2∫_{√(1−z)}^{1} (1−s²)^(m−1) ds.
	s0 := math.Sqrt(1 - z)
	antiderivative := func(s float64) float64 {
		sum := 0.0
		for k := 0; k <= m-1; k++ {
			sum += binom(m-1, k) * math.Pow(-1, float64(k)) * math.Pow(s, float64(2*k+1)) / float64(2*k+1)
		}
		return sum
	}
	num := 2 * (antiderivative(1) - antiderivative(s0))
	lg, _ := math.Lgamma(float64(m))
	lgb, _ := math.Lgamma(0.5)
	lgs, _ := math.Lgamma(float64(m) + 0.5)
	beta := math.Exp(lg + lgb - lgs) // B(m, 1/2)
	return num / beta
}

// TestTwoSidedTAccuracy pins the coefficient tail against exact closed
// forms and against the far tail, where the previous 2·(1 − T_cdf)
// form lost every digit and returned an exact zero.
func TestTwoSidedTAccuracy(t *testing.T) {
	// Exact references: Cauchy (df = 1) and the df = 2 closed form.
	for _, tc := range []struct {
		t    float64
		want float64
	}{{0.5, 2 * math.Atan(1/0.5) / math.Pi}, {2, 2 * math.Atan(1.0/2) / math.Pi}} {
		got, err := twoSidedT(tc.t, 1)
		if err != nil {
			t.Fatal(err)
		}
		if math.Abs(got-tc.want) > 1e-13*tc.want {
			t.Fatalf("twoSidedT(%v, 1) = %.17g, want %.17g", tc.t, got, tc.want)
		}
	}
	for _, tc := range []struct {
		t    float64
		want float64
	}{{0.5, 1 - 0.5/math.Sqrt(2+0.25)}, {2, 1 - 2/math.Sqrt(6)}} {
		got, err := twoSidedT(tc.t, 2)
		if err != nil {
			t.Fatal(err)
		}
		if math.Abs(got-tc.want) > 1e-13*tc.want {
			t.Fatalf("twoSidedT(%v, 2) = %.17g, want %.17g", tc.t, got, tc.want)
		}
	}
	// Even degrees of freedom: the elementary closed form.
	for _, tc := range []struct {
		t  float64
		df int
	}{{2, 6}, {8, 6}, {0.5, 6}, {2, 10}, {8, 10}, {3, 20}} {
		got, err := twoSidedT(tc.t, tc.df)
		if err != nil {
			t.Fatalf("twoSidedT(%v, %d): %v", tc.t, tc.df, err)
		}
		want := tTailExactEven(tc.t, tc.df)
		if math.Abs(got-want) > 1e-9*want {
			t.Fatalf("twoSidedT(%v, %d) = %.17g, closed form says %.17g", tc.t, tc.df, got, want)
		}
	}
	// Large df, against a reference computed outside the library to
	// 60 digits: the tail of t(8 | df = 1e6) is
	// 1.2455063433202503e-15. The cancelling form returned 1.33227e-15
	// here, 7 % high, so this pins the accuracy rather than the order.
	const independentTail = 1.2455063433202503e-15
	got, err := twoSidedT(8, 1000000)
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(got-independentTail) > 1e-6*independentTail {
		t.Fatalf("twoSidedT(8, 1e6) = %.17g, the independent reference is %.17g", got, independentTail)
	}
	// The far tail must stay positive: the cancelling form returned 0
	// for t = 30, df = 100, where the true tail is 8.4e-52.
	far, err := twoSidedT(30, 100)
	if err != nil {
		t.Fatal(err)
	}
	if far <= 0 {
		t.Fatalf("twoSidedT(30, 100) = %v, want a positive tail", far)
	}
	if far > 1e-40 {
		t.Fatalf("twoSidedT(30, 100) = %v, want a tail near 8.4e-52", far)
	}
	// Exact value at the centre.
	if p, err := twoSidedT(0, 7); err != nil || p != 1 {
		t.Fatalf("twoSidedT(0, 7) = %v (err %v), want exactly 1", p, err)
	}
}

// TestLinearRegressionWithoutIntercept pins the statistics of a design
// with no constant column: the uncentred total sum of squares is the
// null model, and the model degrees of freedom are the column count.
func TestLinearRegressionWithoutIntercept(t *testing.T) {
	t.Run("single column", func(t *testing.T) {
		x := mustMatrix(t, []float64{1, 0, -1}, 3, 1)
		y := mustFloats(t, []float64{3, 1, 2}, 3)
		res, err := LinearRegression(x, y)
		if err != nil {
			t.Fatalf("LinearRegression: %v", err)
		}
		// beta = Σxy/Σx² = 1/2, rss = 13.5, Σy² = 14.
		if math.Abs(res.Coefficients[0]-0.5) > 1e-15 {
			t.Fatalf("slope = %v, want 0.5", res.Coefficients[0])
		}
		wantR2 := 1 - 13.5/14
		if math.Abs(res.RSquared-wantR2) > 1e-12 {
			t.Fatalf("R² = %v, want %v (uncentred)", res.RSquared, wantR2)
		}
		if res.DModel != 1 {
			t.Fatalf("DModel = %d, want 1", res.DModel)
		}
		if !(res.FStatistic > 0) || res.FPValue <= 0 || res.FPValue > 1 {
			t.Fatalf("F = %v, p = %v, want a positive statistic and a probability", res.FStatistic, res.FPValue)
		}
	})
	t.Run("two columns", func(t *testing.T) {
		x := mustMatrix(t, []float64{1, 0, 0, 1, -1, 1}, 3, 2)
		y := mustFloats(t, []float64{3, 1, 2}, 3)
		res, err := LinearRegression(x, y)
		if err != nil {
			t.Fatalf("LinearRegression: %v", err)
		}
		// beta = [5/3, 7/3], rss = 16/3, Σy² = 14.
		wantR2 := 1 - (16.0/3)/14
		if math.Abs(res.RSquared-wantR2) > 1e-12 {
			t.Fatalf("R² = %v, want %v", res.RSquared, wantR2)
		}
		if res.DModel != 2 {
			t.Fatalf("DModel = %d, want 2", res.DModel)
		}
	})
	t.Run("intercept unchanged", func(t *testing.T) {
		// The centred form still applies when a constant column is
		// present: an exact line fits perfectly.
		x := mustMatrix(t, []float64{1, 1, 1, 2, 1, 3}, 3, 2)
		y := mustFloats(t, []float64{3, 5, 7}, 3)
		res, err := LinearRegression(x, y)
		if err != nil {
			t.Fatalf("LinearRegression: %v", err)
		}
		if math.Abs(res.RSquared-1) > 1e-12 {
			t.Fatalf("R² = %v, want 1", res.RSquared)
		}
		if res.DModel != 1 {
			t.Fatalf("DModel = %d, want p−1 = 1", res.DModel)
		}
	})
}

// TestLinearRegressionExactFitReport pins the exact-fit report: zero
// standard errors mean an infinite statistic, not a zero one, and the
// p-value is zero rather than absent.
func TestLinearRegressionExactFitReport(t *testing.T) {
	x := mustMatrix(t, []float64{1, 1, 1, 2, 1, 3, 1, 4}, 4, 2)
	y := mustFloats(t, []float64{3, 5, 7, 9}, 4)
	res, err := LinearRegression(x, y)
	if err != nil {
		t.Fatalf("LinearRegression: %v", err)
	}
	if res.RSquared != 1 {
		t.Fatalf("R² = %v, want exactly 1", res.RSquared)
	}
	for j := range 2 {
		if res.StandardErrors[j] != 0 {
			t.Fatalf("se[%d] = %v, want 0", j, res.StandardErrors[j])
		}
		if !math.IsInf(res.TStatistics[j], 0) {
			t.Fatalf("t[%d] = %v, want ±Inf", j, res.TStatistics[j])
		}
		if res.PValues[j] != 0 {
			t.Fatalf("p[%d] = %v, want 0", j, res.PValues[j])
		}
	}
}

// TestLinearRegressionRefusesNonFinite pins the input contract: the
// other tests in the package refuse non-finite data and this one must
// not answer with a silent column of NaN.
func TestLinearRegressionRefusesNonFinite(t *testing.T) {
	x := mustMatrix(t, []float64{1, 1, 1, 2, 1, 3}, 3, 2)
	for _, bad := range []float64{math.NaN(), math.Inf(1), math.Inf(-1)} {
		y := mustFloats(t, []float64{3, bad, 7}, 3)
		if _, err := LinearRegression(x, y); err == nil {
			t.Fatalf("expected an error for the response holding %v", bad)
		} else if !strings.Contains(err.Error(), "non-finite") {
			t.Fatalf("error = %v, want a non-finite refusal", err)
		}
	}
	xb := mustMatrix(t, []float64{1, 1, 1, 2, 1, math.NaN()}, 3, 2)
	if _, err := LinearRegression(xb, mustFloats(t, []float64{3, 5, 7}, 3)); err == nil {
		t.Fatal("expected an error for a design holding NaN")
	}
}

// mustMatrix builds an (r, c) float64 array.
func mustMatrix(t *testing.T, vals []float64, r, c int) *core.Array {
	t.Helper()
	a, err := core.FromFloats(vals, r, c)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return a
}

// TestLinearRegressionTinyScaleInference pins the inference of a response
// on a scale whose squared residuals fall below the subnormal floor: the
// plain residual and total sums of squares read zero there, and the fit
// used to report the evidence backwards, R² of 1 with an infinite t and
// p = 0 beside an F of zero with p = 1. The response y = [0, 0, e] over
// x = 1, 2, 3 keeps every least-squares quantity exactly representable
// while both sums of squares underflow: slope e/2, residual sum e²/6,
// total 2e²/3, (XᵀX)⁻¹₁₁ = 1/2, so SE(slope) = e/(2√3), t = √3,
// p = 1/3, F = 3 and R² = 3/4, all closed fractions.
func TestLinearRegressionTinyScaleInference(t *testing.T) {
	const e = 1e-200
	x := mustMatrix(t, []float64{1, 1, 1, 2, 1, 3}, 3, 2)
	y := mustFloats(t, []float64{0, 0, e}, 3)
	res, err := LinearRegression(x, y)
	if err != nil {
		t.Fatalf("LinearRegression: %v", err)
	}
	if math.Abs(res.Coefficients[1]-e/2) > 1e-12*e/2 {
		t.Fatalf("slope = %.17g, want %.17g", res.Coefficients[1], e/2)
	}
	wantSE := e / (2 * math.Sqrt(3))
	se := res.StandardErrors[1]
	if !(se > 0) || math.IsInf(se, 0) {
		t.Fatalf("slope standard error = %g beside nonzero residuals, want %.17g", se, wantSE)
	}
	if math.Abs(se-wantSE) > 1e-12*wantSE {
		t.Fatalf("slope standard error = %.17g, want %.17g", se, wantSE)
	}
	if math.Abs(res.TStatistics[1]-math.Sqrt(3)) > 1e-12 {
		t.Fatalf("t = %.17g, want √3", res.TStatistics[1])
	}
	if math.Abs(res.PValues[1]-1.0/3) > 1e-12 {
		t.Fatalf("p = %.17g, want 1/3", res.PValues[1])
	}
	if math.Abs(res.RSquared-0.75) > 1e-12 {
		t.Fatalf("R² = %.17g, want 3/4", res.RSquared)
	}
	if math.Abs(res.AdjustedRSquared-0.5) > 1e-12 {
		t.Fatalf("adjusted R² = %.17g, want 1/2", res.AdjustedRSquared)
	}
	if math.Abs(res.FStatistic-3) > 1e-11 {
		t.Fatalf("F = %.17g, want 3", res.FStatistic)
	}
	if math.Abs(res.FPValue-1.0/3) > 1e-11 {
		t.Fatalf("F p-value = %.17g, want 1/3", res.FPValue)
	}

	// Unit weights are the same fit, weighted statistics included.
	w := mustFloats(t, []float64{1, 1, 1}, 3)
	wres, err := WeightedLinearRegression(x, y, w)
	if err != nil {
		t.Fatalf("WeightedLinearRegression: %v", err)
	}
	if math.Abs(wres.TStatistics[1]-math.Sqrt(3)) > 1e-12 {
		t.Fatalf("weighted t = %.17g, want √3", wres.TStatistics[1])
	}
	if math.Abs(wres.RSquared-0.75) > 1e-12 {
		t.Fatalf("weighted R² = %.17g, want 3/4", wres.RSquared)
	}
	if math.Abs(wres.FStatistic-3) > 1e-11 {
		t.Fatalf("weighted F = %.17g, want 3", wres.FStatistic)
	}
	if math.Abs(wres.FPValue-1.0/3) > 1e-11 {
		t.Fatalf("weighted F p-value = %.17g, want 1/3", wres.FPValue)
	}

	// A response with a live scale beside the tiny spread keeps the same
	// behaviour: the slope's own rounding leaves residuals near its last
	// ulp, and the standard error must stay representable and the t
	// finite rather than answer an exact fit the residuals contradict.
	const a = 1e-160
	step := math.Nextafter(3*a, math.Inf(1)) - 3*a
	y2 := mustFloats(t, []float64{a, 2 * a, 3*a + step}, 3)
	res2, err := LinearRegression(x, y2)
	if err != nil {
		t.Fatalf("LinearRegression: %v", err)
	}
	if se2 := res2.StandardErrors[1]; !(se2 > 0) || math.IsInf(se2, 0) {
		t.Fatalf("slope standard error = %g beside nonzero residuals, want a representable value", se2)
	}
	if math.IsInf(res2.TStatistics[1], 0) {
		t.Fatalf("t = %g beside nonzero residuals, want a finite statistic", res2.TStatistics[1])
	}
	if res2.FPValue > 1e-10 || res2.PValues[1] > 1e-10 {
		t.Fatalf("p = %g, F p = %g, want the far tail both", res2.PValues[1], res2.FPValue)
	}
}

// TestLinearRegressionTinyScaleExactLine pins the fully underflowed
// corner: an exact line at a scale where both the residual and the total
// sums of squares fall below the subnormal floor. The t statistics
// already answer the exact fit with infinite evidence; the F test must
// agree with them instead of reporting zero evidence.
func TestLinearRegressionTinyScaleExactLine(t *testing.T) {
	const a = 1e-300
	x := mustMatrix(t, []float64{1, 1, 1, 2, 1, 3}, 3, 2)
	y := mustFloats(t, []float64{a, 2 * a, 3 * a}, 3)
	res, err := LinearRegression(x, y)
	if err != nil {
		t.Fatalf("LinearRegression: %v", err)
	}
	if res.FPValue > 1e-10 {
		t.Fatalf("F p-value = %g on an exact line at a tiny scale, want the far tail beside the infinite t", res.FPValue)
	}
	if res.PValues[1] > 1e-10 {
		t.Fatalf("slope p-value = %g, want the exact-fit report", res.PValues[1])
	}
}