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

package stats

import (
	"math"
	"testing"

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

// mustDesign builds a design matrix from row-major values.
func mustDesign(t *testing.T, vals []float64, rows, cols int) *core.Array {
	t.Helper()
	a, err := core.FromFloats(vals, rows, cols)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return a
}

// TestLinearRegressionExactModel pins the inference on a noiseless
// linear model: the coefficients recover, R² is 1, σ̂² is 0 and the
// perfect fit is flagged as such (t infinite, p zero).
func TestLinearRegressionExactModel(t *testing.T) {
	// y = 2 + 3x over x = 1..6.
	design := mustDesign(t, []float64{
		1, 1, 1, 2, 1, 3, 1, 4, 1, 5, 1, 6,
	}, 6, 2)
	yv := make([]float64, 6)
	for i := range 6 {
		yv[i] = 2 + 3*float64(i+1)
	}
	y, err := core.FromFloats(yv, 6)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	res, err := LinearRegression(design, y)
	if err != nil {
		t.Fatalf("LinearRegression: %v", err)
	}
	if math.Abs(res.Coefficients[0]-2) > 1e-12 || math.Abs(res.Coefficients[1]-3) > 1e-12 {
		t.Fatalf("coefficients = (%.10f, %.10f), want (2, 3)", res.Coefficients[0], res.Coefficients[1])
	}
	if math.Abs(res.RSquared-1) > 1e-12 {
		t.Fatalf("R² = %.12f, want 1", res.RSquared)
	}
	if res.ResidualVariance > 1e-24 {
		t.Fatalf("σ̂² = %.3g, want 0", res.ResidualVariance)
	}
	if res.FPValue > 1e-12 {
		t.Fatalf("F p-value = %.3g, want 0", res.FPValue)
	}
}

// TestLinearRegressionNoisy pins the inference against hand-computed
// formulas on a small noisy dataset.
func TestLinearRegressionNoisy(t *testing.T) {
	design := mustDesign(t, []float64{
		1, 0, 1, 1, 1, 2, 1, 3, 1, 4, 1, 5,
	}, 6, 2)
	yv := []float64{1.2, 2.1, 2.9, 4.2, 4.8, 6.3}
	y, _ := core.FromFloats(yv, 6)
	res, err := LinearRegression(design, y)
	if err != nil {
		t.Fatalf("LinearRegression: %v", err)
	}
	n, p := 6.0, 2.0
	// Reference computations from first principles.
	xv := []float64{0, 1, 2, 3, 4, 5}
	sx := 0.0
	sy := 0.0
	sxy := 0.0
	sxx := 0.0
	for i := range 6 {
		sx += xv[i]
		sy += yv[i]
		sxy += xv[i] * yv[i]
		sxx += xv[i] * xv[i]
	}
	slope := (n*sxy - sx*sy) / (n*sxx - sx*sx)
	intercept := (sy - slope*sx) / n
	if math.Abs(res.Coefficients[1]-slope) > 1e-12 || math.Abs(res.Coefficients[0]-intercept) > 1e-12 {
		t.Fatalf("coefficients = (%.10f, %.10f), want (%.10f, %.10f)",
			res.Coefficients[0], res.Coefficients[1], intercept, slope)
	}
	rss := 0.0
	for i := range 6 {
		r := yv[i] - (intercept + slope*xv[i])
		rss += r * r
		if math.Abs(res.Residuals[i]-r) > 1e-12 {
			t.Fatalf("residual[%d] = %.10f, want %.10f", i, res.Residuals[i], r)
		}
	}
	if math.Abs(res.ResidualVariance-rss/(n-p)) > 1e-12 {
		t.Fatalf("σ̂² = %.10f, want %.10f", res.ResidualVariance, rss/(n-p))
	}
	// SE of the slope: σ̂·sqrt(1/Sxx_c) with Sxx_c = Σ(x−x̄)².
	xbar := sx / n
	sxxc := 0.0
	for _, v := range xv {
		sxxc += (v - xbar) * (v - xbar)
	}
	seSlope := math.Sqrt(rss/(n-p)) / math.Sqrt(sxxc)
	// NaN would slip through a plain tolerance comparison, so it gets
	// its own gate before the value check.
	if math.IsNaN(res.StandardErrors[1]) || math.IsNaN(res.TStatistics[1]) || math.IsNaN(res.PValues[1]) {
		t.Fatalf("non-finite inference: SE=%.6g t=%.6g p=%.6g",
			res.StandardErrors[1], res.TStatistics[1], res.PValues[1])
	}
	if math.Abs(res.StandardErrors[1]-seSlope) > 1e-12 {
		t.Fatalf("SE(slope) = %.10f, want %.10f", res.StandardErrors[1], seSlope)
	}
	if math.Abs(res.TStatistics[1]-slope/seSlope) > 1e-10 {
		t.Fatalf("t(slope) = %.10f, want %.10f", res.TStatistics[1], slope/seSlope)
	}
	// Two-sided p from the t CDF itself (independent re-entry).
	pv, err := twoSidedT(res.TStatistics[1], 4)
	if err != nil {
		t.Fatalf("twoSidedT: %v", err)
	}
	if math.Abs(res.PValues[1]-pv) > 1e-14 {
		t.Fatalf("p(slope) = %.10f, want %.10f", res.PValues[1], pv)
	}
	// F = t² for the single-slope model.
	if math.Abs(res.FStatistic-res.TStatistics[1]*res.TStatistics[1]) > 1e-8 {
		t.Fatalf("F = %.8f, want t² = %.8f", res.FStatistic, res.TStatistics[1]*res.TStatistics[1])
	}
	// F p equals the slope's t p for p = 2.
	if math.Abs(res.FPValue-res.PValues[1]) > 1e-10 {
		t.Fatalf("F p = %.10f, want %.10f", res.FPValue, res.PValues[1])
	}
}

// TestLinearRegressionTwoPredictors pins a multi-column design against
// a brute-force normal-equation solve in the test.
func TestLinearRegressionTwoPredictors(t *testing.T) {
	g := core.NewGenerator(17)
	const n = 40
	vals := make([]float64, n*3)
	yv := make([]float64, n)
	for i := range n {
		a := g.NormalUnit()
		b := g.NormalUnit()
		vals[i*3] = 1
		vals[i*3+1] = a
		vals[i*3+2] = b
		yv[i] = -1.5 + 2*a - 0.7*b + 0.3*g.NormalUnit()
	}
	design := mustDesign(t, vals, n, 3)
	y, _ := core.FromFloats(yv, n)
	res, err := LinearRegression(design, y)
	if err != nil {
		t.Fatalf("LinearRegression: %v", err)
	}
	if math.Abs(res.Coefficients[0]+1.5) > 0.2 || math.Abs(res.Coefficients[1]-2) > 0.2 ||
		math.Abs(res.Coefficients[2]+0.7) > 0.2 {
		t.Fatalf("coefficients = %v, want roughly (-1.5, 2, -0.7)", res.Coefficients)
	}
	if res.RSquared < 0.9 {
		t.Fatalf("R² = %.3f, the model explains more than that", res.RSquared)
	}
	if res.FPValue > 1e-10 {
		t.Fatalf("F p = %.3g, the model is significant", res.FPValue)
	}
	if res.DModel != 2 || res.DResidual != 37 {
		t.Fatalf("degrees of freedom (%d, %d), want (2, 37)", res.DModel, res.DResidual)
	}
}

// TestLinearRegressionErrors pins the input gates.
func TestLinearRegressionErrors(t *testing.T) {
	x := mustDesign(t, []float64{1, 1, 1, 1}, 2, 2)
	y, _ := core.FromFloats([]float64{1, 2}, 2)
	if _, err := LinearRegression(x, y); err == nil {
		t.Error("n == p accepted")
	}
	if _, err := LinearRegression(x, y); err == nil {
		t.Error("square design accepted")
	}
	bad := mustDesign(t, []float64{1, 2, 2, 4, 1, 2}, 3, 2)
	y3, _ := core.FromFloats([]float64{1, 2, 3}, 3)
	if _, err := LinearRegression(bad, y3); err == nil {
		t.Error("rank-deficient design accepted")
	}
}