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

// TestPoissonRegressionRecoversCoefficients fits a generated count
// response whose truth is known: with 1500 samples the Newton fit must
// land within a few standard errors of the generating coefficients,
// the Wald statistics must flag the slope, and the fitted means must
// increase in the direction of the true slope.
func TestPoissonRegressionRecoversCoefficients(t *testing.T) {
	g := core.NewGenerator(7)
	const n = 1500
	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
		counts, err := PoissonDraws(g, 1, math.Exp(0.3+0.9*xv))
		if err != nil {
			t.Fatalf("PoissonDraws: %v", err)
		}
		y.RawFloats()[i] = counts.FloatAt(0)
	}
	res, err := PoissonRegression(design, y)
	if err != nil {
		t.Fatalf("PoissonRegression: %v", err)
	}
	if !res.Converged {
		t.Fatal("the fit reported no convergence")
	}
	if math.Abs(res.Coefficients[0]-0.3) > 4*res.StandardErrors[0] {
		t.Fatalf("intercept = %.4g (%.4g SE), outside four SEs of 0.3",
			res.Coefficients[0], res.StandardErrors[0])
	}
	if math.Abs(res.Coefficients[1]-0.9) > 4*res.StandardErrors[1] {
		t.Fatalf("slope = %.4g (%.4g SE), outside four SEs of 0.9",
			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])
	}
	if !(res.Fitted[n-1] > res.Fitted[0]) {
		t.Fatalf("fitted means not increasing: %g then %g", res.Fitted[0], res.Fitted[n-1])
	}
}

// TestPoissonRegressionInterceptOnlyClosedForm pins the one case with
// a closed-form maximum likelihood estimate: with an intercept-only
// design the estimate is log of the mean count and its standard error
// is 1/sqrt(n·mean).
func TestPoissonRegressionInterceptOnlyClosedForm(t *testing.T) {
	const n = 40
	design := core.New(core.Float, n, 1)
	y := core.New(core.Float, n)
	total := 0.0
	for i := range n {
		design.RawFloats()[i] = 1
		y.RawFloats()[i] = float64(i % 7)
		total += float64(i % 7)
	}
	meanY := total / n
	res, err := PoissonRegression(design, y)
	if err != nil {
		t.Fatalf("PoissonRegression: %v", err)
	}
	if math.Abs(res.Coefficients[0]-math.Log(meanY)) > 1e-9 {
		t.Fatalf("intercept = %.15g, want log(%.15g) = %.15g",
			res.Coefficients[0], meanY, math.Log(meanY))
	}
	wantSE := 1 / math.Sqrt(n*meanY)
	if math.Abs(res.StandardErrors[0]-wantSE) > 1e-9*wantSE {
		t.Fatalf("SE = %.15g, want %.15g", res.StandardErrors[0], wantSE)
	}
	for i := range n {
		if math.Abs(res.Fitted[i]-meanY) > 1e-9 {
			t.Fatalf("fitted[%d] = %.15g, want %.15g", i, res.Fitted[i], meanY)
		}
	}
	// The score must vanish at the optimum.
	score := 0.0
	for i := range n {
		score += res.Fitted[i] - y.FloatAt(i)
	}
	if math.Abs(score) > 1e-6 {
		t.Fatalf("score at the optimum = %.3g, want 0", score)
	}
}

// TestPoissonRegressionStandardErrorsClosedForm checks the Wald
// inference against a closed-form 2x2 inverse of the Fisher
// information built from the fitted means: no solve, only the
// adjugate formula.
func TestPoissonRegressionStandardErrorsClosedForm(t *testing.T) {
	// Near-deterministic counts: large means make the rounding to
	// integers a relative perturbation under 1e-3, so the fit must
	// recover the generating coefficients tightly.
	const n = 25
	truth := [2]float64{0.2, 0.5}
	design := core.New(core.Float, n, 2)
	y := core.New(core.Float, n)
	for i := range n {
		xv := float64(i)
		design.RawFloats()[i*2] = 1
		design.RawFloats()[i*2+1] = xv
		y.RawFloats()[i] = math.Round(math.Exp(truth[0] + truth[1]*xv))
	}
	res, err := PoissonRegression(design, y)
	if err != nil {
		t.Fatalf("PoissonRegression: %v", err)
	}
	for j := range 2 {
		if math.Abs(res.Coefficients[j]-truth[j]) > 4e-3 {
			t.Fatalf("coefficient %d = %.6g, want %.6g within rounding noise",
				j, res.Coefficients[j], truth[j])
		}
	}
	// Fisher information from the fitted means, inverted by the
	// adjugate.
	a, b, c := 0.0, 0.0, 0.0
	for i := range n {
		x1 := design.FloatAt(i * 2)
		x2 := design.FloatAt(i*2 + 1)
		mu := res.Fitted[i]
		a += x1 * x1 * mu
		b += x1 * x2 * mu
		c += x2 * x2 * mu
	}
	det := a*c - b*b
	wantSE0 := math.Sqrt(c / det)
	wantSE1 := math.Sqrt(a / det)
	if math.Abs(res.StandardErrors[0]-wantSE0) > 1e-9*wantSE0 {
		t.Fatalf("SE0 = %.15g, want %.15g", res.StandardErrors[0], wantSE0)
	}
	if math.Abs(res.StandardErrors[1]-wantSE1) > 1e-9*wantSE1 {
		t.Fatalf("SE1 = %.15g, want %.15g", res.StandardErrors[1], wantSE1)
	}
	z := res.ZStatistics[0]
	// The p-value reference is computed independently of the tail
	// formula the fit uses: Simpson quadrature of the Gaussian density,
	// every summand positive, sharing nothing with NormalCDF or
	// math.Erfc. The intercept's z lands far past the z ≈ 8.3 where the
	// algebraic form 2·(1−Φ(z)) cancels to an exact zero, so this pin
	// fails against the cancelled form, which answers 0 here.
	wantP := gaussTailReference(math.Abs(z))
	if math.Abs(res.PValues[0]-wantP) > 1e-6*wantP {
		t.Fatalf("p-value = %.15g, want %.15g", res.PValues[0], wantP)
	}
	// The slope's z is in the hundreds: its tail underflows the float64
	// range entirely, and an exact zero is the correctly rounded answer
	// both the fit and the reference produce.
	if res.PValues[1] != 0 {
		t.Fatalf("p-value of the slope = %.15g, want the underflowed 0", res.PValues[1])
	}
}

// TestPoissonRegressionRefusals checks the response and design
// validation and the honest failure of a singular Fisher information
// from a duplicated column.
func TestPoissonRegressionRefusals(t *testing.T) {
	y4 := mustFromFloats(t, []float64{1, 0, 1, 0}, 4)
	if _, err := PoissonRegression(core.New(core.Float, 4), y4); err == nil {
		t.Fatal("a rank 1 design was accepted")
	}
	design := core.New(core.Float, 4, 2)
	for i := range 4 {
		design.RawFloats()[i*2] = 1
		design.RawFloats()[i*2+1] = float64(i)
	}
	rank2 := core.New(core.Float, 2, 2)
	if _, err := PoissonRegression(design, rank2); err == nil {
		t.Fatal("a rank 2 response was accepted")
	}
	if _, err := PoissonRegression(design, mustFromFloats(t, []float64{1, 0, 1}, 3)); err == nil {
		t.Fatal("a length mismatch was accepted")
	}
	if _, err := PoissonRegression(design, mustFromFloats(t, []float64{1, 0, 1, -2}, 4)); err == nil {
		t.Fatal("a negative count was accepted")
	}
	if _, err := PoissonRegression(design, mustFromFloats(t, []float64{1, 0, 1, 0.5}, 4)); err == nil {
		t.Fatal("a fractional count was accepted")
	}
	if _, err := PoissonRegression(design, mustFromFloats(t, []float64{1, 0, math.NaN(), 0}, 4)); err == nil {
		t.Fatal("a NaN count was accepted")
	}
	if _, err := PoissonRegression(design, mustFromFloats(t, []float64{1, 0, math.Inf(1), 0}, 4)); err == nil {
		t.Fatal("an infinite count was accepted")
	}
	bad := mustFromFloats(t, []float64{1, math.NaN(), 0, 1, 2, 0, 1, 2}, 4, 2)
	if _, err := PoissonRegression(bad, y4); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("PoissonRegression with a NaN design: %v", err)
	}
	// n <= p has no unique fit.
	square := core.New(core.Float, 2, 2)
	square.RawFloats()[0] = 1
	square.RawFloats()[3] = 1
	if _, err := PoissonRegression(square, mustFromFloats(t, []float64{1, 2}, 2)); err == nil {
		t.Fatal("n <= p was accepted")
	}
	// A duplicated column makes the Fisher information singular at the
	// start.
	dup := core.New(core.Float, 4, 2)
	for i := range 4 {
		dup.RawFloats()[i*2] = float64(i)
		dup.RawFloats()[i*2+1] = float64(i)
	}
	if _, err := PoissonRegression(dup, y4); err == nil || !strings.Contains(err.Error(), "singular") {
		t.Fatalf("a duplicated column: %v", err)
	}
}

// TestPoissonRegressionIsDeterministic runs the same fit twice and
// requires identical coefficients bit for bit, as every Tensor entry
// point promises.
func TestPoissonRegressionIsDeterministic(t *testing.T) {
	const n = 300
	build := func(seed int64) (*core.Array, *core.Array) {
		g := core.NewGenerator(seed)
		design := core.New(core.Float, n, 2)
		y := core.New(core.Float, n)
		for i := range n {
			xv := -1 + 2*g.Unit()
			design.RawFloats()[i*2] = 1
			design.RawFloats()[i*2+1] = xv
			counts, err := PoissonDraws(g, 1, math.Exp(0.4+0.6*xv))
			if err != nil {
				t.Fatalf("PoissonDraws: %v", err)
			}
			y.RawFloats()[i] = counts.FloatAt(0)
		}
		return design, y
	}
	d1, y1 := build(11)
	d2, y2 := build(11)
	r1, err := PoissonRegression(d1, y1)
	if err != nil {
		t.Fatalf("first fit: %v", err)
	}
	r2, err := PoissonRegression(d2, y2)
	if err != nil {
		t.Fatalf("second fit: %v", err)
	}
	for j := range r1.Coefficients {
		if r1.Coefficients[j] != r2.Coefficients[j] {
			t.Fatalf("coefficient %d differs: %.17g vs %.17g",
				j, r1.Coefficients[j], r2.Coefficients[j])
		}
	}
	if r1.LogLikelihood != r2.LogLikelihood {
		t.Fatalf("log likelihood differs: %.17g vs %.17g",
			r1.LogLikelihood, r2.LogLikelihood)
	}
}

// TestPoissonRegressionFloat32Design runs the fit over a float32
// design, the widened path the raw-slice fast path must agree with.
func TestPoissonRegressionFloat32Design(t *testing.T) {
	const n = 120
	f64 := core.New(core.Float, n, 2)
	f32 := core.New(core.Float32, n, 2)
	y := core.New(core.Float, n)
	for i := range n {
		xv := -2 + 4*float64(i)/float64(n-1)
		f64.RawFloats()[i*2] = 1
		f64.RawFloats()[i*2+1] = xv
		f32.SetFloatAt(i*2, 1)
		f32.SetFloatAt(i*2+1, xv)
		y.RawFloats()[i] = math.Round(math.Exp(0.2 + 0.3*xv))
	}
	ref, err := PoissonRegression(f64, y)
	if err != nil {
		t.Fatalf("float64 fit: %v", err)
	}
	got, err := PoissonRegression(f32, y)
	if err != nil {
		t.Fatalf("float32 fit: %v", err)
	}
	for j := range ref.Coefficients {
		if math.Abs(got.Coefficients[j]-ref.Coefficients[j]) > 1e-6 {
			t.Fatalf("coefficient %d: float32 %.12g vs float64 %.12g",
				j, got.Coefficients[j], ref.Coefficients[j])
		}
	}
}