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

package stats

import (
	"fmt"
	"math"
	"testing"

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

// The randomised and cross-checked tests that guard the pair counting of
// KendallTau, the covariance diagonals of the regression entries and the
// pairwise-slope walk of Theil-Sen against the implementations they
// replace, followed by the benchmarks of those paths.

// randSource is a deterministic linear congruential generator for the
// randomised checks: a fixed recurrence makes every failing case
// reproducible from its seed alone.
type randSource struct{ state uint64 }

// next returns the next draw in [0, 1).
func (r *randSource) next() float64 {
	r.state = r.state*6364136223846793005 + 1442695040888963407
	return float64(r.state>>11) / float64(uint64(1)<<53)
}

// intn returns the next draw in [0, n).
func (r *randSource) intn(n int) int { return int(r.next() * float64(n)) }

// bruteForceTally counts the pairs of two paired samples the way the
// definition reads: every ordered pair of positions, its two differences
// multiplied, the sign of the product deciding the side, and the ties of
// a sample counted over the same pair walk. It is the oracle the
// merge-sort pair counting is checked against.
func bruteForceTally(xs, ys []float64) pairTally {
	var tally pairTally
	for i := 1; i < len(xs); i++ {
		for j := 0; j < i; j++ {
			switch prod := (xs[i] - xs[j]) * (ys[i] - ys[j]); {
			case prod > 0:
				tally.concordant++
			case prod < 0:
				tally.discordant++
			}
			if xs[i] == xs[j] {
				tally.tiedX++
			}
			if ys[i] == ys[j] {
				tally.tiedY++
			}
		}
	}
	return tally
}

// bruteForceTau is the O(n²) τ-b of the definition, the oracle the
// returned correlation is checked against.
func bruteForceTau(xs, ys []float64) float64 {
	tally := bruteForceTally(xs, ys)
	n := float64(len(xs))
	n0 := n * (n - 1) / 2
	return (float64(tally.concordant) - float64(tally.discordant)) /
		math.Sqrt((n0-float64(tally.tiedX))*(n0-float64(tally.tiedY)))
}

// TestKendallTauCountsMatchBruteForce checks the merge-sort pair
// counting against the pair enumeration it replaces. The values are
// drawn from a handful of levels so that ties in either sample are the
// common case, which is where the block scan and the tie accounting
// carry the work; continuous draws, degenerate ranges and a constant
// second sample cover the rest. Both sides are exact integers, so the
// comparison is on the returned bits rather than on a tolerance.
func TestKendallTauCountsMatchBruteForce(t *testing.T) {
	src := &randSource{state: 0x9e3779b97f4a7c15}
	levels := []float64{-1, 0, 1, 2, 3}
	for trial := range 400 {
		n := 2 + src.intn(38)
		xs := make([]float64, n)
		ys := make([]float64, n)
		// A third of the draws are continuous, the rest come from the
		// levels, so both samples are tie-heavy on average.
		draw := func() float64 {
			if src.intn(3) == 0 {
				return src.next()*4 - 2
			}
			return levels[src.intn(len(levels))]
		}
		for i := range n {
			xs[i] = draw()
			ys[i] = draw()
		}
		if trial%4 == 0 {
			// The pairing is perfectly monotone, its τ exactly 1.
			ys = append([]float64(nil), xs...)
		}
		if trial%9 == 0 {
			// A constant chunk inside otherwise varying samples: the
			// widest blocks of equal values the scan has to handle.
			for i := n / 3; i < 2*n/3; i++ {
				xs[i] = xs[0]
			}
		}
		want := bruteForceTally(xs, ys)
		got := tallyPairs(append([]float64(nil), xs...), append([]float64(nil), ys...))
		if got != want {
			t.Fatalf("trial %d, n %d: pair counts %+v, want %+v\nx = %v\ny = %v", trial, n, got, want, xs, ys)
		}
		if want.tiedX == int64(n*(n-1)/2) || want.tiedY == int64(n*(n-1)/2) {
			continue // a constant sample has no τ-b to compare
		}
		gotTau, err := KendallTau(mustFloats(t, xs), mustFloats(t, ys))
		if err != nil {
			t.Fatalf("trial %d, n %d: KendallTau: %v", trial, n, err)
		}
		wantTau := bruteForceTau(xs, ys)
		if gotTau != wantTau {
			t.Fatalf("trial %d, n %d: τ = %v (%#x), want %v (%#x)\nx = %v\ny = %v",
				trial, n, gotTau, math.Float64bits(gotTau), wantTau, math.Float64bits(wantTau), xs, ys)
		}
	}
}

// invertDiagonal returns the diagonal of the inverse of a dense matrix
// by Gauss-Jordan elimination with partial pivoting: an independent
// route to the covariance diagonals the fits read through the shared LU
// solve. It is accurate to rounding, which is all the cross-checks below
// need from it.
func invertDiagonal(m [][]float64) []float64 {
	p := len(m)
	aug := make([][]float64, p)
	for i := range p {
		aug[i] = make([]float64, 2*p)
		copy(aug[i], m[i])
		aug[i][p+i] = 1
	}
	for k := range p {
		piv := k
		for i := k + 1; i < p; i++ {
			if math.Abs(aug[i][k]) > math.Abs(aug[piv][k]) {
				piv = i
			}
		}
		aug[k], aug[piv] = aug[piv], aug[k]
		d := aug[k][k]
		for j := range 2 * p {
			aug[k][j] /= d
		}
		for i := range p {
			if i == k {
				continue
			}
			f := aug[i][k]
			if f == 0 {
				continue
			}
			for j := range 2 * p {
				aug[i][j] -= f * aug[k][j]
			}
		}
	}
	diag := make([]float64, p)
	for i := range p {
		diag[i] = aug[i][p+i]
	}
	return diag
}

// copyMatrix copies a p×p system, the shape SolveSystem factors in place.
func copyMatrix(m [][]float64) [][]float64 {
	out := make([][]float64, len(m))
	for i, row := range m {
		out[i] = append([]float64(nil), row...)
	}
	return out
}

// TestSolveSystemUnitColumnsMatchPerColumnSolves checks the shared
// factorisation the covariance diagonals read against the
// per-coefficient solves it replaced: the same normal-equations matrix,
// once against all p unit columns and once against each column alone, on
// a positive-definite system built from a deterministic design. The
// diagonal entries are compared bit for bit.
func TestSolveSystemUnitColumnsMatchPerColumnSolves(t *testing.T) {
	const p = 9
	src := &randSource{state: 0xda3e39cb94b95bdb}
	m := make([][]float64, p)
	for i := range p {
		m[i] = make([]float64, p)
	}
	row := make([]float64, p)
	for range 40 {
		row[0] = 1
		for c := 1; c < p; c++ {
			row[c] = src.next() - 0.5
		}
		for i := range p {
			for j := range p {
				m[i][j] += row[i] * row[j]
			}
		}
	}
	units := make([][]float64, p)
	for j := range p {
		units[j] = make([]float64, p)
		units[j][j] = 1
	}
	together, err := base.SolveSystem("test", copyMatrix(m), units)
	if err != nil {
		t.Fatal(err)
	}
	for j := range p {
		e := make([]float64, p)
		e[j] = 1
		alone, err := base.SolveSystem("test", copyMatrix(m), [][]float64{e})
		if err != nil {
			t.Fatal(err)
		}
		if got, want := together[j][j], alone[0][j]; got != want {
			t.Fatalf("diagonal %d of the batched solve = %v (%#x), want %v (%#x) from the single-column solve",
				j, got, math.Float64bits(got), want, math.Float64bits(want))
		}
	}
}

// TestLogisticRegressionWaldStandardErrors rebuilds the Fisher
// information at the reported fit from the reported probabilities and
// inverts it by an independent elimination: the Wald standard errors are
// that inverse's diagonal, so the check covers the shared factorisation
// the entry now reads them through on a design no digest pins.
func TestLogisticRegressionWaldStandardErrors(t *testing.T) {
	const n, p = 300, 6
	src := &randSource{state: 0x2545f4914f6cdd1d}
	design := make([]float64, n*p)
	response := make([]float64, n)
	for r := range n {
		eta := 0.3
		for c := range p {
			v := src.next() - 0.5
			design[r*p+c] = v
			eta += 0.8 * v
		}
		design[r*p] = 1
		if src.next() < 1/(1+math.Exp(-eta)) {
			response[r] = 1
		}
	}
	res, err := LogisticRegression(mustFloats(t, design, n, p), mustFloats(t, response, n))
	if err != nil {
		t.Fatal(err)
	}
	fisher := make([][]float64, p)
	for i := range p {
		fisher[i] = make([]float64, p)
	}
	for r := range n {
		w := res.Fitted[r] * (1 - res.Fitted[r])
		for i := range p {
			for j := range p {
				fisher[i][j] += design[r*p+i] * design[r*p+j] * w
			}
		}
	}
	diag := invertDiagonal(fisher)
	for j := range p {
		want := math.Sqrt(diag[j])
		if rel := math.Abs(res.StandardErrors[j]-want) / want; rel > 1e-9 {
			t.Fatalf("standard error %d = %.15g, want %.15g (relative %.3g)",
				j, res.StandardErrors[j], want, rel)
		}
	}
}

// TestHuberRegressionStandardErrors does the same for the Huber fit:
// σ²·(XᵀWX)⁻¹ with σ the reported robust scale and W the reported final
// weights, its diagonal checked against an independent elimination of
// the weighted normal equations the fit writes its own weights into.
func TestHuberRegressionStandardErrors(t *testing.T) {
	const n, p = 300, 6
	src := &randSource{state: 0x853c49e6748fea9b}
	design := make([]float64, n*p)
	response := make([]float64, n)
	for r := range n {
		design[r*p] = 1
		response[r] = 1
		for c := 1; c < p; c++ {
			v := src.next() - 0.5
			design[r*p+c] = v
			response[r] += float64(c) * 0.5 * v
		}
		response[r] += 0.2 * (src.next() - 0.5)
	}
	response[7] += 5 // the gross point the robust fit exists for
	res, err := HuberRegression(mustFloats(t, design, n, p), mustFloats(t, response, n))
	if err != nil {
		t.Fatal(err)
	}
	if !(res.Scale > 0) {
		t.Fatalf("the robust scale collapsed to %g on a contaminated sample", res.Scale)
	}
	wxx := make([][]float64, p)
	for i := range p {
		wxx[i] = make([]float64, p)
	}
	for r := range n {
		w := res.Weights[r]
		for i := range p {
			for j := range p {
				wxx[i][j] += design[r*p+i] * design[r*p+j] * w
			}
		}
	}
	diag := invertDiagonal(wxx)
	for j := range p {
		want := math.Sqrt(res.Scale * res.Scale * diag[j])
		if rel := math.Abs(res.StandardErrors[j]-want) / want; rel > 1e-9 {
			t.Fatalf("standard error %d = %.15g, want %.15g (relative %.3g)",
				j, res.StandardErrors[j], want, rel)
		}
	}
}

// TestTheilSenParallelPairs checks the pairwise-slope walk against the
// serial enumeration it replaces, element for element and bit for bit,
// on a sample large enough to take the parallel path and with repeated
// predictors, which are the pairs the walk has to skip. The reference
// walk reproduces the original append order, so a slope landing in the
// wrong buffer slot moves the median and fails the test.
func TestTheilSenParallelPairs(t *testing.T) {
	const n = 1400
	src := &randSource{state: 0xc0ffee123456789a}
	xs := make([]float64, n)
	ys := make([]float64, n)
	for i := range n {
		// A tenth of a unit apart, so a tenth of the pairs share a
		// predictor and carry no slope.
		xs[i] = math.Floor(src.next()*300) / 10
		ys[i] = 2*xs[i] + src.next()
	}
	gotIntercept, gotSlope, err := TheilSenRegression(mustFloats(t, xs), mustFloats(t, ys))
	if err != nil {
		t.Fatal(err)
	}
	slopes := make([]float64, 0, n*(n-1)/2)
	for i := range n {
		for j := i + 1; j < n; j++ {
			if dx := xs[j] - xs[i]; dx != 0 {
				slopes = append(slopes, (ys[j]-ys[i])/dx)
			}
		}
	}
	slope := medianSlice(slopes)
	intercepts := make([]float64, n)
	for i := range n {
		intercepts[i] = ys[i] - slope*xs[i]
	}
	intercept := medianSlice(intercepts)
	if gotSlope != slope || gotIntercept != intercept {
		t.Fatalf("Theil-Sen = (%v, %v), want (%v, %v) from the serial walk",
			gotIntercept, gotSlope, intercept, slope)
	}
}

// regDesign builds a deterministic (n, p) design with an intercept column
// and covariates in (−0.5, 0.5), plus the row-major values for the
// cross-checks that rebuild a normal-equations matrix from them.
func regDesign(tb testing.TB, n, p int, seed uint64) (*core.Array, []float64) {
	tb.Helper()
	src := &randSource{state: seed}
	vals := make([]float64, n*p)
	for r := range n {
		vals[r*p] = 1
		for c := 1; c < p; c++ {
			vals[r*p+c] = src.next() - 0.5
		}
	}
	a, err := core.FromFloats(vals, n, p)
	if err != nil {
		tb.Fatal(err)
	}
	return a, vals
}

// BenchmarkKendallTau measures the rank correlation at a size where the
// pair enumeration it replaced costs milliseconds: 8.4 million pairs,
// counted by the merge sort in a fraction of that.
func BenchmarkKendallTau(b *testing.B) {
	const n = 4096
	src := &randSource{state: 0x123456789abcdef}
	xv := make([]float64, n)
	yv := make([]float64, n)
	for i := range n {
		// Rounded draws, so ties in both samples are part of the walk.
		xv[i] = math.Floor(src.next()*512) / 512
		yv[i] = math.Floor(src.next()*512) / 512
	}
	x, err := core.FromFloats(xv, n)
	if err != nil {
		b.Fatal(err)
	}
	y, err := core.FromFloats(yv, n)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := KendallTau(x, y); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkLinearRegressionCovariance measures the least-squares fit with
// its inference on a wide design, where the covariance diagonal is a
// visible share of the work: p is large enough that the p solves it needs
// outgrow the normal equations themselves.
func BenchmarkLinearRegressionCovariance(b *testing.B) {
	const n, p = 400, 48
	x, _ := regDesign(b, n, p, 0x5eed0001)
	src := &randSource{state: 0x5eed0002}
	yv := make([]float64, n)
	for i := range n {
		yv[i] = src.next()*2 - 1
	}
	y, err := core.FromFloats(yv, n)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := LinearRegression(x, y); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkPoissonRegressionInference measures the count fit with its
// Wald inference, the covariance diagonal of the Fisher information
// included.
func BenchmarkPoissonRegressionInference(b *testing.B) {
	const n, p = 600, 30
	x, vals := regDesign(b, n, p, 0x5eed0003)
	src := &randSource{state: 0x5eed0004}
	yv := make([]float64, n)
	for i := range n {
		eta := 0.4
		for c := range p {
			eta += 0.3 * vals[i*p+c]
		}
		mu := math.Exp(eta)
		yv[i] = math.Round(mu * (0.5 + src.next()))
	}
	y, err := core.FromFloats(yv, n)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := PoissonRegression(x, y); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkLogisticRegressionInference measures the binary fit with its
// Wald inference, the covariance diagonal of the Fisher information
// included.
func BenchmarkLogisticRegressionInference(b *testing.B) {
	const n, p = 1500, 30
	x, vals := regDesign(b, n, p, 0x5eed0005)
	src := &randSource{state: 0x5eed0006}
	yv := make([]float64, n)
	for i := range n {
		eta := 0.2 + 0.5*vals[i*p]
		for c := 1; c < p; c++ {
			eta += 0.5 * vals[i*p+c]
		}
		if src.next() < 1/(1+math.Exp(-eta)) {
			yv[i] = 1
		}
	}
	y, err := core.FromFloats(yv, n)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := LogisticRegression(x, y); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkHuberRegressionInference measures the robust fit with its
// standard errors: the reweighting rounds, their medians, and the
// weighted normal equations read for the covariance diagonal.
func BenchmarkHuberRegressionInference(b *testing.B) {
	const n, p = 800, 8
	x, vals := regDesign(b, n, p, 0x5eed0007)
	src := &randSource{state: 0x5eed0008}
	yv := make([]float64, n)
	for i := range n {
		yv[i] = 1 + 0.5*vals[i*p+1] - 0.3*vals[i*p+2] + 0.2*(src.next()-0.5)
	}
	yv[11] += 6
	y, err := core.FromFloats(yv, n)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := HuberRegression(x, y); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkTheilSenRegression measures the median-of-slopes fit: the
// pairwise walk over the pairs with distinct predictors and the median
// taken over them.
func BenchmarkTheilSenRegression(b *testing.B) {
	const n = 512
	src := &randSource{state: 0x5eed0009}
	xv := make([]float64, n)
	yv := make([]float64, n)
	for i := range n {
		xv[i] = src.next()*10 - 5
		yv[i] = 1 + 2*xv[i] + src.next()
	}
	x, err := core.FromFloats(xv, n)
	if err != nil {
		b.Fatal(err)
	}
	y, err := core.FromFloats(yv, n)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := TheilSenRegression(x, y); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkLassoPath measures the warm-started regularisation path over
// the documented hundred-lambda grid, on an elastic net mixing that
// exercises the standardised coordinate descent. The design carries no
// intercept column: the penalty standardises every column, and a
// constant one has no slope to report.
func BenchmarkLassoPath(b *testing.B) {
	const n, p = 400, 12
	src := &randSource{state: 0x5eed000a}
	xv := make([]float64, n*p)
	yv := make([]float64, n)
	for i := range n {
		for c := range p {
			xv[i*p+c] = src.next() - 0.5
		}
		yv[i] = 1 + 0.8*xv[i*p+1] - 0.5*xv[i*p+2] + 0.1*(src.next()-0.5)
	}
	x, err := core.FromFloats(xv, n, p)
	if err != nil {
		b.Fatal(err)
	}
	y, err := core.FromFloats(yv, n)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := LassoPath(x, y, 1); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkKendallTauScaling reports the merge-sort pair count against
// the sample size: the linear-ish growth the count replaced the
// quadratic walk with.
func BenchmarkKendallTauScaling(b *testing.B) {
	for _, n := range []int{512, 2048, 8192} {
		src := &randSource{state: 0x51ed + uint64(n)}
		xv := make([]float64, n)
		yv := make([]float64, n)
		for i := range n {
			xv[i] = math.Floor(src.next()*512) / 512
			yv[i] = math.Floor(src.next()*512) / 512
		}
		x, err := core.FromFloats(xv, n)
		if err != nil {
			b.Fatal(err)
		}
		y, err := core.FromFloats(yv, n)
		if err != nil {
			b.Fatal(err)
		}
		b.Run(fmt.Sprintf("n=%d", n), func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				if _, err := KendallTau(x, y); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}

// BenchmarkTheilSenRegressionCap runs the fit at the observation cap,
// where the pairwise slope list holds millions of entries and the middle
// order statistic is taken by selection rather than by a full sort.
func BenchmarkTheilSenRegressionCap(b *testing.B) {
	const n = TheilSenMaxObservations
	src := &randSource{state: 0x5eed1234}
	xv := make([]float64, n)
	yv := make([]float64, n)
	for i := range n {
		xv[i] = src.next()*10 - 5
		yv[i] = 1 + 2*xv[i] + src.next()
	}
	x, err := core.FromFloats(xv, n)
	if err != nil {
		b.Fatal(err)
	}
	y, err := core.FromFloats(yv, n)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := TheilSenRegression(x, y); err != nil {
			b.Fatal(err)
		}
	}
}