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

// TestLogisticRegressionRecoversCoefficients fits a generated binary
// response whose truth is known: with 2000 samples the Newton fit
// must land within a few standard errors of the generating
// coefficients, the Wald statistics must match the coefficients, and
// the fitted probabilities must increase in the direction of the
// true slope.
func TestLogisticRegressionRecoversCoefficients(t *testing.T) {
	g := core.NewGenerator(7)
	const n = 2000
	design := core.New(core.Float, n, 2)
	y := core.New(core.Float, n)
	for i := range n {
		xv := -2 + 4*g.Unit()
		design.RawFloats()[i*2] = 1
		design.RawFloats()[i*2+1] = xv
		pr := 1 / (1 + math.Exp(-(0.5 + 1.5*xv)))
		bit := 0.0
		if g.Unit() < pr {
			bit = 1
		}
		y.RawFloats()[i] = bit
	}
	res, err := LogisticRegression(design, y)
	if err != nil {
		t.Fatalf("LogisticRegression: %v", err)
	}
	if !res.Converged {
		t.Fatal("the fit reported no convergence")
	}
	if math.Abs(res.Coefficients[0]-0.5) > 4*res.StandardErrors[0] {
		t.Fatalf("intercept = %.4g (%.4g SE), outside four SEs of 0.5",
			res.Coefficients[0], res.StandardErrors[0])
	}
	if math.Abs(res.Coefficients[1]-1.5) > 4*res.StandardErrors[1] {
		t.Fatalf("slope = %.4g (%.4g SE), outside four SEs of 1.5",
			res.Coefficients[1], res.StandardErrors[1])
	}
	if res.ZStatistics[1] <= 3 {
		t.Fatalf("slope z = %.4g, want a clearly non-zero effect", res.ZStatistics[1])
	}
	last := res.Fitted[n-1]
	first := res.Fitted[0]
	if !(last > first) {
		t.Fatalf("fitted probabilities not increasing: %g then %g", first, last)
	}
	// The maximised likelihood must beat the null model's.
	nullLike := float64(n) * math.Log(0.5)
	if res.LogLikelihood <= nullLike {
		t.Fatalf("log likelihood %.4g does not beat the null %.4g", res.LogLikelihood, nullLike)
	}
}

// TestLogisticRegressionRefusals checks the response and design
// guards, including the separable-data refusal.
func TestLogisticRegressionRefusals(t *testing.T) {
	design := core.New(core.Float, 4, 2)
	y := core.New(core.Float, 4)
	if _, err := LogisticRegression(core.New(core.Float, 4), y); err == nil {
		t.Fatal("rank-1 design accepted")
	}
	bad := core.New(core.Float, 4)
	bad.RawFloats()[2] = 0.5
	if _, err := LogisticRegression(design, bad); err == nil {
		t.Fatal("non-binary response accepted")
	}
	// Perfect separation: y = 1 exactly when x > 0 has no finite
	// optimum, and the run must say so instead of diverging.
	xsep := core.New(core.Float, 8, 2)
	sep := core.New(core.Float, 8)
	for i := range 8 {
		xsep.RawFloats()[i*2] = 1
		xsep.RawFloats()[i*2+1] = float64(i) - 3.5
		if i >= 4 {
			sep.RawFloats()[i] = 1
		}
	}
	if _, err := LogisticRegression(xsep, sep); err == nil {
		t.Fatal("separable data accepted")
	}
}

// TestMultivariateNormalDensity checks the density against the
// two-dimensional formula with a diagonal covariance, where the
// answer is a product of one-dimensional normals.
func TestMultivariateNormalDensity(t *testing.T) {
	mean := smallVector(t, []float64{1, -2})
	cov, err := core.FromFloats([]float64{4, 0, 0, 9}, 2, 2)
	if err != nil {
		t.Fatalf("cov: %v", err)
	}
	x := smallVector(t, []float64{3, 1})
	got, err := MultivariateNormalLogDensity(mean, cov, x)
	if err != nil {
		t.Fatalf("MultivariateNormalLogDensity: %v", err)
	}
	dx := float64(3-1) / 2
	dy := float64(1-(-2)) / 3
	want := math.Log(1/(2*math.Pi*6)) - 0.5*dx*dx - 0.5*dy*dy
	if math.Abs(got-want) > 1e-12 {
		t.Fatalf("log density = %.14g, want %.14g", got, want)
	}
	atMean, err := MultivariateNormalLogDensity(mean, cov, mean)
	if err != nil {
		t.Fatalf("density at the mean: %v", err)
	}
	if math.Abs(atMean-math.Log(1/(2*math.Pi*6))) > 1e-12 {
		t.Fatalf("density at the mean = %.14g, want the normaliser", atMean)
	}
	if _, err := MultivariateNormalLogDensity(mean, smallVector(t, []float64{1, 2, 3}), x); err == nil {
		t.Fatal("mismatched covariance accepted")
	}
	// A negative pivot must be refused by name.
	bad, _ := core.FromFloats([]float64{1, 0, 0, -4}, 2, 2)
	if _, err := MultivariateNormalLogDensity(mean, bad, x); err == nil {
		t.Fatal("indefinite covariance accepted")
	}
}

// TestMultivariateNormalDraws checks the sampler's moments: with
// 60000 draws the sample mean and covariance must sit close to the
// parameters, well inside the Monte Carlo error of the moment.
func TestMultivariateNormalDraws(t *testing.T) {
	g := core.NewGenerator(11)
	mean := smallVector(t, []float64{2, -1})
	cov, err := core.FromFloats([]float64{1, 0.5, 0.5, 4}, 2, 2)
	if err != nil {
		t.Fatalf("cov: %v", err)
	}
	const n = 60000
	draws, err := MultivariateNormalDraws(g, n, mean, cov)
	if err != nil {
		t.Fatalf("MultivariateNormalDraws: %v", err)
	}
	if draws.Shape()[0] != n || draws.Shape()[1] != 2 {
		t.Fatalf("shape %v, want [%d 2]", draws.Shape(), n)
	}
	m1, m2, c, v1, v2 := 0.0, 0.0, 0.0, 0.0, 0.0
	for i := range n {
		x := draws.FloatAt(i * 2)
		y := draws.FloatAt(i*2 + 1)
		m1 += x
		m2 += y
		v1 += x * x
		v2 += y * y
		c += x * y
	}
	m1 /= n
	m2 /= n
	if math.Abs(m1-2) > 0.03 || math.Abs(m2+1) > 0.05 {
		t.Fatalf("means = %.4f, %.4f, want 2, -1", m1, m2)
	}
	if v1/n-m1*m1 > 1.06 || v2/n-m2*m2 > 4.25 {
		t.Fatalf("variances = %.4f, %.4f, want about 1 and 4", v1/n-m1*m1, v2/n-m2*m2)
	}
	covHat := c/n - m1*m2
	if math.Abs(covHat-0.5) > 0.05 {
		t.Fatalf("covariance = %.4f, want 0.5", covHat)
	}
}

// TestKernelDensity checks the estimate on a standard normal sample:
// it must integrate to one over a wide grid, peak near the true mode
// and stay non-negative everywhere.
func TestKernelDensity(t *testing.T) {
	g := core.NewGenerator(23)
	const n = 3000
	sampleVals := make([]float64, n)
	for i := range n {
		sampleVals[i] = g.NormalUnit()
	}
	sample := smallVector(t, sampleVals)
	const lo, hi = -5.0, 5.0
	const grid = 400
	pointVals := make([]float64, grid)
	for i := range grid {
		pointVals[i] = lo + (hi-lo)*float64(i)/float64(grid-1)
	}
	points := smallVector(t, pointVals)
	density, err := KernelDensity(sample, 0, points)
	if err != nil {
		t.Fatalf("KernelDensity: %v", err)
	}
	vals := density.RawFloats()[:density.Len()]
	total := 0.0
	for i, v := range vals {
		if v < 0 {
			t.Fatalf("negative density at %g", pointVals[i])
		}
		if i > 0 {
			total += 0.5 * (v + vals[i-1]) * ((hi - lo) / (grid - 1))
		}
	}
	if math.Abs(total-1) > 0.01 {
		t.Fatalf("the estimate integrates to %.4f, want 1", total)
	}
	peak, peakAt := 0.0, 0.0
	for i, v := range vals {
		if v > peak {
			peak, peakAt = v, pointVals[i]
		}
	}
	if math.Abs(peakAt) > 0.25 {
		t.Fatalf("the estimate peaks at %.3f, want the true mode near 0", peakAt)
	}
	if peak < 0.3 || peak > 0.5 {
		t.Fatalf("peak height %.4f, want the normal's 0.399 within a KDE's honesty", peak)
	}
}

// TestLogisticRegressionRefusesNonFiniteDesign pins the finite-input
// gate: a NaN coefficient in the design used to flow through the
// sigmoid and the Newton step into a fit that reported convergence on
// an all-NaN result.
func TestLogisticRegressionRefusesNonFiniteDesign(t *testing.T) {
	design := core.New(core.Float, 4, 2)
	for i := range 8 {
		design.RawFloats()[i] = float64(i%4) + float64(i/4)
	}
	design.RawFloats()[5] = math.NaN()
	y := core.New(core.Float, 4)
	y.RawFloats()[0], y.RawFloats()[1] = 0, 1
	y.RawFloats()[2], y.RawFloats()[3] = 1, 0
	if _, err := LogisticRegression(design, y); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("LogisticRegression with a NaN design: %v", err)
	}
}

// TestKernelDensityViewReadsVisibleElements pins the dense payload
// bound on both sweeps: a rebased view shares its parent's payload,
// which runs past the view's own count, so the sample and the points
// are cut to the elements a caller can see. A non-finite value in the
// invisible tail must neither poison the estimate nor refuse the call,
// and the density must be the one the standalone sample gives, bit for
// bit.
func TestKernelDensityViewReadsVisibleElements(t *testing.T) {
	visible := []float64{0.5, 1.5, 2.5, 3.5}
	pointVals := []float64{0, 1, 2, 3, 4}
	dense := mustFromFloats(t, append(append([]float64(nil), visible...), math.NaN(), math.Inf(1)), 6)
	view, err := core.Slice(dense, 0, 0, len(visible))
	if err != nil {
		t.Fatalf("Slice: %v", err)
	}
	pointParent := mustFromFloats(t, append(append([]float64(nil), pointVals...), math.NaN(), math.Inf(-1)), len(pointVals)+2)
	pointView, err := core.Slice(pointParent, 0, 0, len(pointVals))
	if err != nil {
		t.Fatalf("Slice: %v", err)
	}
	got, err := KernelDensity(view, 0.5, pointView)
	if err != nil {
		t.Fatalf("KernelDensity over views with a non-finite tail: %v", err)
	}
	want, err := KernelDensity(mustFloats(t, visible, len(visible)), 0.5, mustFloats(t, pointVals, len(pointVals)))
	if err != nil {
		t.Fatalf("KernelDensity over the visible elements: %v", err)
	}
	if got.Len() != want.Len() {
		t.Fatalf("the estimate carries %d values, want %d", got.Len(), want.Len())
	}
	for i := range want.Len() {
		if gb, wb := math.Float64bits(got.FloatAt(i)), math.Float64bits(want.FloatAt(i)); gb != wb {
			t.Fatalf("density[%d] = %v (%#x) over the view, %v (%#x) over the visible elements",
				i, got.FloatAt(i), gb, want.FloatAt(i), wb)
		}
	}
	// The sample's own view is the whole parent's, tail included: the
	// guard must still refuse it by name.
	if _, err := KernelDensity(dense, 0.5, pointView); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("KernelDensity over the whole parent: %v, want the non-finite refusal", err)
	}
}