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

package stats

import (
	"math"
	"slices"
	"strings"
	"testing"

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

// huberFixture builds the seeded contaminated fixture the Huber pins
// run on: sixty observations of a clean linear model, then ten percent
// of the responses pushed thirty units off the trend. It returns the
// design with the constant column first, the clean response and the
// contaminated response.
func huberFixture(t *testing.T) (*core.Array, *core.Array, *core.Array) {
	t.Helper()
	const n = 60
	g := core.NewGenerator(11)
	designVals := make([]float64, 0, 2*n)
	clean := make([]float64, n)
	for i := range n {
		x := g.NormalUnit()
		clean[i] = 2 + 3*x + 0.5*g.NormalUnit()
		designVals = append(designVals, 1, x)
	}
	design := mustFromFloats(t, designVals, n, 2)
	contaminated := append([]float64(nil), clean...)
	for i := 0; i < n; i += 10 {
		contaminated[i] += 100
	}
	return design, mustFromFloats(t, clean, n), mustFromFloats(t, contaminated, n)
}

// TestHuberResistsOutliers is the pin that gives the estimator its
// reason to exist: on the contaminated fixture the Huber fit lands
// beside the clean-data ordinary least squares answer, while plain
// least squares on the same contaminated data is dragged visibly off
// it. Both margins are measured, not asserted as anecdotes.
func TestHuberResistsOutliers(t *testing.T) {
	design, clean, contaminated := huberFixture(t)
	cleanOLS, err := LinearRegression(design, clean)
	if err != nil {
		t.Fatalf("LinearRegression on the clean data: %v", err)
	}
	if math.Abs(cleanOLS.Coefficients[0]-2) > 0.2 || math.Abs(cleanOLS.Coefficients[1]-3) > 0.2 {
		t.Fatalf("the clean-data OLS is off the truth: (%.4f, %.4f)", cleanOLS.Coefficients[0], cleanOLS.Coefficients[1])
	}
	olsCont, err := LinearRegression(design, contaminated)
	if err != nil {
		t.Fatalf("LinearRegression on the contaminated data: %v", err)
	}
	huber, err := HuberRegression(design, contaminated)
	if err != nil {
		t.Fatalf("HuberRegression: %v", err)
	}
	if !huber.Converged {
		t.Fatalf("the Huber fit did not converge")
	}
	dragged := math.Abs(olsCont.Coefficients[1] - cleanOLS.Coefficients[1])
	resisted := math.Abs(huber.Coefficients[1] - cleanOLS.Coefficients[1])
	t.Logf("slope: clean OLS %.4f, contaminated OLS %.4f (drag %.4f), contaminated Huber %.4f (drag %.4f)",
		cleanOLS.Coefficients[1], olsCont.Coefficients[1], dragged, huber.Coefficients[1], resisted)
	if dragged < 0.3 {
		t.Fatalf("the contamination failed to drag OLS in the first place: drag %.4f", dragged)
	}
	if resisted > 0.15 {
		t.Fatalf("the Huber slope sits %.4f from the clean-data answer, want at most 0.15", resisted)
	}
	if d := math.Abs(huber.Coefficients[0] - cleanOLS.Coefficients[0]); d > 0.15 {
		t.Fatalf("the Huber intercept sits %.4f from the clean-data answer, want at most 0.15", d)
	}
}

// TestHuberWeightsBandTaper pins the weight function itself: exactly
// one strictly inside the band and on its boundary, the band over the
// magnitude outside it, and no weight anywhere else on the number
// line. The fixture's contaminated observations must then be the ones
// the fit tapered hardest.
func TestHuberWeightsBandTaper(t *testing.T) {
	// The function, on exact values.
	if huberWeight(0.4, 1) != 1 {
		t.Fatalf("a residual inside the band does not carry weight one")
	}
	if huberWeight(1, 1) != 1 {
		t.Fatalf("a residual on the band boundary does not carry weight one")
	}
	if huberWeight(-1, 1) != 1 {
		t.Fatalf("the band is not symmetric: a negative boundary residual tapered")
	}
	if w := huberWeight(2, 1); w != 0.5 {
		t.Fatalf("a residual at twice the band carries %g, want 0.5", w)
	}
	if w := huberWeight(-3, 1.5); w != 0.5 {
		t.Fatalf("a negative residual at twice the band carries %g, want 0.5", w)
	}
	if huberWeight(0, 0) != 1 {
		t.Fatalf("a collapsed band must keep a zero residual at weight one")
	}
	if huberWeight(1, 0) != 0 {
		t.Fatalf("a collapsed band must give a nonzero residual zero weight")
	}
	// The taper decreases strictly beyond the band.
	previous := 1.0
	for magnitude := 1.5; magnitude < 20; magnitude += 0.5 {
		w := huberWeight(magnitude, 1)
		if w >= previous {
			t.Fatalf("the taper rose at magnitude %.1f: %g after %g", magnitude, w, previous)
		}
		previous = w
	}
	// The fit's weights: the six contaminated observations are the six
	// lightest.
	design, _, contaminated := huberFixture(t)
	huber, err := HuberRegression(design, contaminated)
	if err != nil {
		t.Fatalf("HuberRegression: %v", err)
	}
	weights := append([]float64(nil), huber.Weights...)
	sorted := append([]float64(nil), weights...)
	slices.Sort(sorted)
	lightest := map[float64]bool{}
	for _, w := range sorted[:6] {
		lightest[w] = true
	}
	for i := 0; i < 60; i += 10 {
		if !lightest[huber.Weights[i]] {
			t.Fatalf("contaminated observation %d carries weight %.4f, not among the six lightest", i, huber.Weights[i])
		}
	}
}

// TestHuberExactFit walks the collapsed-scale corner: a constant
// response is reproduced exactly by the unpenalised fit, every
// residual is zero, the robust scale is zero, and the fit reports the
// exactness instead of dividing by it.
func TestHuberExactFit(t *testing.T) {
	design := mustFromFloats(t, []float64{1, 0, 1, 1, 1, 2, 1, 3, 1, 4}, 5, 2)
	y := mustFromFloats(t, []float64{2.5, 2.5, 2.5, 2.5, 2.5}, 5)
	res, err := HuberRegression(design, y)
	if err != nil {
		t.Fatalf("HuberRegression: %v", err)
	}
	if !res.Converged {
		t.Fatalf("the exact fit did not report convergence")
	}
	if res.Scale != 0 {
		t.Fatalf("the exact fit reports scale %g, want 0", res.Scale)
	}
	if math.Abs(res.Coefficients[0]-2.5) > 1e-9 || math.Abs(res.Coefficients[1]) > 1e-9 {
		t.Fatalf("the exact fit moved: (%.12g, %.12g), want (2.5, 0)", res.Coefficients[0], res.Coefficients[1])
	}
	for i := range 5 {
		if res.Weights[i] != 1 {
			t.Fatalf("an exact fit tapered observation %d to weight %g", i, res.Weights[i])
		}
		if math.Abs(res.Residuals[i]) > 1e-12 {
			t.Fatalf("the exact fit left residual %g at observation %d", res.Residuals[i], i)
		}
		if res.StandardErrors[0] != 0 {
			t.Fatalf("the exact fit reports standard error %g, want 0", res.StandardErrors[0])
		}
	}
}

// TestTheilSenExactLine pins the estimator on an outlier-free line:
// every pairwise slope of y = 2 + 3x is exactly 3 in floating point,
// so their median is exactly 3 and the intercept exactly 2, bit for
// bit.
func TestTheilSenExactLine(t *testing.T) {
	const n = 10
	xs := make([]float64, n)
	ys := make([]float64, n)
	for i := range n {
		xs[i] = float64(i)
		ys[i] = 2 + 3*float64(i)
	}
	x := mustFromFloats(t, xs, n)
	y := mustFromFloats(t, ys, n)
	intercept, slope, err := TheilSenRegression(x, y)
	if err != nil {
		t.Fatalf("TheilSenRegression: %v", err)
	}
	if slope != 3 {
		t.Fatalf("slope = %.17g, want exactly 3", slope)
	}
	if intercept != 2 {
		t.Fatalf("intercept = %.17g, want exactly 2", intercept)
	}
}

// TestTheilSenSurvivesBrokenPoint breaks one observation of an exact
// line by a thousand units: the least squares slope is destroyed by
// the leverage, while the Theil-Sen slope and intercept stay exactly
// on the line, because the broken point's slopes sit outside the
// median window.
func TestTheilSenSurvivesBrokenPoint(t *testing.T) {
	const n, broken = 25, 7
	xs := make([]float64, n)
	ys := make([]float64, n)
	for i := range n {
		xs[i] = float64(i)
		ys[i] = 1 + 0.5*float64(i)
	}
	ys[broken] += 1000
	x := mustFromFloats(t, xs, n)
	y := mustFromFloats(t, ys, n)
	design := mustFromFloats(t, []float64{1, 0, 1, 1, 1, 2, 1, 3, 1, 4, 1, 5, 1, 6, 1, 7, 1, 8, 1, 9, 1, 10, 1, 11, 1, 12, 1, 13, 1, 14, 1, 15, 1, 16, 1, 17, 1, 18, 1, 19, 1, 20, 1, 21, 1, 22, 1, 23, 1, 24}, n, 2)
	ols, err := LinearRegression(design, y)
	if err != nil {
		t.Fatalf("LinearRegression: %v", err)
	}
	dragged := math.Abs(ols.Coefficients[1] - 0.5)
	t.Logf("one broken point drags the OLS slope to %.4f; Theil-Sen holds it", ols.Coefficients[1])
	if dragged < 3 {
		t.Fatalf("the broken point failed to destroy OLS in the first place: drag %.4f", dragged)
	}
	intercept, slope, err := TheilSenRegression(x, y)
	if err != nil {
		t.Fatalf("TheilSenRegression: %v", err)
	}
	if slope != 0.5 {
		t.Fatalf("slope = %.17g under one broken point, want exactly 0.5", slope)
	}
	if intercept != 1 {
		t.Fatalf("intercept = %.17g under one broken point, want exactly 1", intercept)
	}
}

// TestTheilSenValidation refuses the inputs the estimator cannot
// answer exactly: wrong shapes, non-finite samples, degenerate
// predictors and samples past the exactness contract.
func TestTheilSenValidation(t *testing.T) {
	x := mustFromFloats(t, []float64{0, 1, 2, 3}, 4)
	y := mustFromFloats(t, []float64{0, 1, 2, 3}, 4)
	if _, _, err := TheilSenRegression(mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6, 7, 8}, 4, 2), y); err == nil || !strings.Contains(err.Error(), "the predictor must be rank 1") {
		t.Fatalf("a rank 2 predictor: got %v, want the rank refusal", err)
	}
	if _, _, err := TheilSenRegression(x, mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6}, 3, 2)); err == nil || !strings.Contains(err.Error(), "the response must be rank 1") {
		t.Fatalf("a rank 2 response: got %v, want the rank refusal", err)
	}
	if _, _, err := TheilSenRegression(x, mustFromFloats(t, []float64{1, 2, 3}, 3)); err == nil || !strings.Contains(err.Error(), "samples but the response") {
		t.Fatalf("a length mismatch: got %v, want the length refusal", err)
	}
	if _, _, err := TheilSenRegression(x, core.New(core.Complex, 4)); err == nil || !strings.Contains(err.Error(), "complex") {
		t.Fatalf("a complex response: got %v, want the complex refusal", err)
	}
	// Two agreeing observations reach the three-observation floor the
	// simple model needs for a median of pairwise slopes.
	if _, _, err := TheilSenRegression(mustFromFloats(t, []float64{0, 1}, 2), mustFromFloats(t, []float64{0, 1}, 2)); err == nil || !strings.Contains(err.Error(), "at least three observations") {
		t.Fatalf("two observations: got %v, want the three-observation floor", err)
	}
	if _, _, err := TheilSenRegression(mustFromFloats(t, []float64{2, 2, 2, 2}, 4), y); err == nil || !strings.Contains(err.Error(), "does not vary") {
		t.Fatalf("a constant predictor: got %v, want the no-slope refusal", err)
	}
	if _, _, err := TheilSenRegression(mustFromFloats(t, []float64{0, 1, math.NaN(), 3}, 4), y); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("a non-finite predictor: got %v, want the non-finite refusal", err)
	}
	if _, _, err := TheilSenRegression(x, mustFromFloats(t, []float64{0, 1, 2, math.Inf(1)}, 4)); err == nil || !strings.Contains(err.Error(), "non-finite") {
		t.Fatalf("a non-finite response: got %v, want the non-finite refusal", err)
	}
	if _, err := HuberRegression(x, y); err == nil || !strings.Contains(err.Error(), "must be rank 2") {
		t.Fatalf("a rank 1 design: got %v, want the rank refusal", err)
	}
	// One past the exactness contract: the refusal names the cost.
	big := make([]float64, TheilSenMaxObservations+1)
	for i := range big {
		big[i] = float64(i)
	}
	_, _, err := TheilSenRegression(mustFromFloats(t, big, len(big)), mustFromFloats(t, big, len(big)))
	if err == nil || !strings.Contains(err.Error(), "exactness contract") {
		t.Fatalf("a sample past the exactness contract: got %v, want the exactness refusal", err)
	}
	// At the contract itself the fit runs.
	if _, _, err := TheilSenRegression(mustFromFloats(t, big[:TheilSenMaxObservations], TheilSenMaxObservations), mustFromFloats(t, big[:TheilSenMaxObservations], TheilSenMaxObservations)); err != nil {
		t.Fatalf("a sample at the contract was refused: %v", err)
	}
}

// TestRobustEstimatorsOnIntegerArrays exercises the widening
// accessor's fallback paths: an integer design and response reach the
// estimators through FloatAt rather than a raw float payload, and the
// fits must answer exactly as they do on the widened floats.
func TestRobustEstimatorsOnIntegerArrays(t *testing.T) {
	design := mustFromInts(t, []int64{1, 0, 1, 1, 1, 2, 1, 3, 1, 4, 1, 5, 1, 6}, 7, 2)
	y := mustFromInts(t, []int64{2, 3, 4, 5, 6, 7, 8}, 7)
	huber, err := HuberRegressionTuned(design, y, DefaultHuberTuning)
	if err != nil {
		t.Fatalf("HuberRegression on integer input: %v", err)
	}
	if math.Abs(huber.Coefficients[0]-2) > 1e-9 || math.Abs(huber.Coefficients[1]-1) > 1e-9 {
		t.Fatalf("the integer-input Huber fit is (%.6f, %.6f), want (2, 1)",
			huber.Coefficients[0], huber.Coefficients[1])
	}
	x := mustFromInts(t, []int64{0, 1, 2, 3, 4, 5, 6}, 7)
	intercept, slope, err := TheilSenRegression(x, y)
	if err != nil {
		t.Fatalf("TheilSenRegression on integer input: %v", err)
	}
	if slope != 1 || intercept != 2 {
		t.Fatalf("the integer-input Theil-Sen fit is (%.17g, %.17g), want exactly (2, 1)", intercept, slope)
	}
}

// TestHuberValidation refuses the malformed inputs, tuning constants
// included: the tuning constant is the estimator's shape, and a zero
// or non-finite one has no Huber loss behind it.
func TestHuberValidation(t *testing.T) {
	design := mustFromFloats(t, []float64{1, 0, 1, 1, 1, 2, 1, 3, 1, 4, 1, 5}, 6, 2)
	y := mustFromFloats(t, []float64{1, 3, 2, 5, 4, 7}, 6)
	if _, err := HuberRegressionTuned(design, y, 0); err == nil {
		t.Fatalf("a zero tuning constant was accepted")
	}
	if _, err := HuberRegressionTuned(design, y, -1); err == nil {
		t.Fatalf("a negative tuning constant was accepted")
	}
	if _, err := HuberRegressionTuned(design, y, math.NaN()); err == nil {
		t.Fatalf("a NaN tuning constant was accepted")
	}
	if _, err := HuberRegressionTuned(design, y, math.Inf(1)); err == nil {
		t.Fatalf("an infinite tuning constant was accepted")
	}
	if _, err := HuberRegressionTuned(design, mustFromFloats(t, []float64{1, 2, 3, math.NaN(), 5, 7}, 6), 1.345); err == nil {
		t.Fatalf("a non-finite response was accepted")
	}
	if _, err := HuberRegressionTuned(mustFromFloats(t, []float64{1, 0, 1, 1}, 2, 2), mustFromFloats(t, []float64{1, 2}, 2), 1.345); err == nil {
		t.Fatalf("n = p was accepted")
	}
	if _, err := HuberRegressionTuned(mustFromFloats(t, []float64{1, 1, 1, 2, 1, 3}, 3, 2), mustFromFloats(t, []float64{1, 2, 3, 4}, 4), 1.345); err == nil {
		t.Fatalf("a length mismatch was accepted")
	}
	complexDesign := core.New(core.Complex, 6, 2)
	if _, err := HuberRegressionTuned(complexDesign, y, 1.345); err == nil {
		t.Fatalf("complex input was accepted")
	}
	// A rank-deficient design is refused by the shared solve.
	singular := mustFromFloats(t, []float64{1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2}, 6, 2)
	if _, err := HuberRegressionTuned(singular, y, 1.345); err == nil {
		t.Fatalf("a rank-deficient design was accepted")
	}
}