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

package stats

import (
	"math"
	"sync"

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

// Robust regression: the fits that survive the wild observations the
// classical least squares would chase. Two estimators live here. The
// Huber M-estimator bounds the influence of a gross outlier by
// replacing the squared loss with one that grows quadratically only
// inside a band around zero and linearly outside it, fitted by
// iteratively reweighted least squares. Theil-Sen replaces the whole
// least-squares machinery with the median of the pairwise slopes,
// which a minority of broken points cannot move.

// DefaultHuberTuning is the Huber tuning constant the plain
// HuberRegression entry uses. The value 1.345 is the literature's
// standard choice: with the band measured in robust standard
// deviations it puts the estimator at 95 percent asymptotic efficiency
// at the Gaussian while keeping the influence of an outlier bounded at
// 1.345 times what a residual inside the band would have.
const DefaultHuberTuning = 1.345

// huberMaxIterations bounds the reweighting loop, and the tolerance is
// the largest coefficient movement one iteration may leave behind for
// the fit to call itself settled. Both follow the glm.go convention.
const (
	huberMaxIterations = 100
	huberTolerance     = 1e-10
)

// HuberRegressionResult carries a Huber M-estimate of the linear
// model.
type HuberRegressionResult struct {
	// Coefficients are the M-estimates β̂, one per design column, in
	// the design's own order.
	Coefficients []float64
	// StandardErrors are the classical asymptotic standard errors read
	// from σ²·(XᵀWX)⁻¹, W the final weight vector: the same shape of
	// statement WeightedLinearRegression makes, with the robust scale
	// σ in place of the residual standard deviation.
	StandardErrors []float64
	// Weights are the final IRLS weights, one per observation: exactly
	// 1 inside the band |r| ≤ k·σ and tapering as k·σ/|r| outside it.
	// They are the diagnostic a robust fit exists to produce: the
	// contaminated observations are the ones at the bottom of the
	// list.
	Weights []float64
	// Scale is the final robust scale σ, 1.4826 times the median
	// absolute deviation of the residuals, the estimate the band is
	// measured in.
	Scale float64
	// Fitted and Residuals align with the rows of the design.
	Fitted    []float64
	Residuals []float64
	// Iterations counts the reweighting steps taken; Converged reports
	// whether the coefficient updates fell under the tolerance.
	Iterations int
	Converged  bool
}

// HuberRegression fits y = X·β with the Huber M-estimator at the
// default tuning constant. See HuberRegressionTuned for the full
// contract.
func HuberRegression(x, y *core.Array) (*HuberRegressionResult, error) {
	return HuberRegressionTuned(x, y, DefaultHuberTuning)
}

// HuberRegressionTuned fits y = X·β by Huber's M-estimation with the
// tuning constant tuning: the loss is r²/2 inside the band |r| ≤
// tuning·σ and tuning·σ·(|r| − tuning·σ/2) outside it, so a wild
// observation pulls like a linear, not a quadratic, residual. The fit
// runs by iteratively reweighted least squares: ordinary least squares
// to start, then each round re-estimates the robust scale σ as 1.4826
// times the median absolute deviation of the current residuals,
// weights each observation by 1 inside the band and tuning·σ/|r|
// outside it, and solves the weighted normal equations until no
// coefficient moves by more than 1e-10.
//
// The design carries n rows and p columns exactly as
// LinearRegression's, the intercept included by the caller as a
// constant column when wanted, and the same validations apply: n > p,
// a full-rank design, real finite input, and a positive finite tuning
// constant. A robust scale that collapses to zero, more than half the
// residuals landing on one value, stops the iteration as an exact fit
// and is reported rather than divided by.
func HuberRegressionTuned(x, y *core.Array, tuning float64) (*HuberRegressionResult, error) {
	const name = "HuberRegression"
	if x.NDim() != 2 {
		return nil, base.Errf("%s: the design must be rank 2, got shape %s", name, base.ShapeText(x.Shape()))
	}
	if y.NDim() != 1 {
		return nil, base.Errf("%s: the response must be rank 1", name)
	}
	if x.Dtype() == core.Complex || y.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex inputs are not supported", name)
	}
	n, p := x.Shape()[0], x.Shape()[1]
	if y.Len() != n {
		return nil, base.Errf("%s: the design has %d rows but the response %d", name, n, y.Len())
	}
	if n <= p {
		return nil, base.Errf("%s: need n > p, got %d observations and %d columns", name, n, p)
	}
	if p == 0 {
		return nil, base.Errf("%s: the design must carry at least one column", name)
	}
	if err := checkFinite(name, "the design", x); err != nil {
		return nil, err
	}
	if err := checkFinite(name, "the response", y); err != nil {
		return nil, err
	}
	if math.IsNaN(tuning) || math.IsInf(tuning, 0) || tuning <= 0 {
		return nil, base.Errf("%s: the tuning constant must be finite and positive, got %g", name, tuning)
	}
	fx := rawFloats(x)
	fy := rawFloats(y)

	// The start is the ordinary least squares answer: the M-estimator's
	// own optimum is rarely far from it, and the reweighting does the
	// rest. The normal equations go through the shared LU, exactly as
	// LinearRegression solves them.
	beta, err := leastSquaresSolve(name, n, p, x, y, fx, fy, nil)
	if err != nil {
		return nil, err
	}
	residuals := make([]float64, n)
	regressionResiduals(n, p, x, y, fx, fy, beta, residuals)
	scaleScratch := make([]float64, n)
	// One workspace serves every reweighting round: the solve consumes
	// its buffers in place and the next round clears them first.
	ws := newLeastSquaresWorkspace(p)
	out := &HuberRegressionResult{
		Coefficients: beta,
		Weights:      make([]float64, n),
	}
	converged := false
	iterations := huberMaxIterations
	for iter := 1; iter <= huberMaxIterations; iter++ {
		scale := huberScale(residuals, scaleScratch)
		if scale == 0 {
			// The robust scale has collapsed: more than half the
			// residuals sit on one value, the band has nothing to
			// widen, and the fit is as settled as it will ever be. The
			// weights follow the collapsed band, one inside it and
			// zero outside, and the loop stops rather than divide by
			// the zero the tapering would need.
			for i, r := range residuals {
				out.Weights[i] = huberWeight(r, 0)
			}
			out.Scale = 0
			converged = true
			iterations = iter - 1
			break
		}
		band := tuning * scale
		for i, r := range residuals {
			out.Weights[i] = huberWeight(r, band)
		}
		updated, err := leastSquaresSolveWS(name, n, p, x, y, fx, fy, out.Weights, ws)
		if err != nil {
			return nil, err
		}
		worst := 0.0
		for j := range p {
			if d := math.Abs(updated[j] - beta[j]); d > worst {
				worst = d
			}
			beta[j] = updated[j]
		}
		regressionResiduals(n, p, x, y, fx, fy, beta, residuals)
		if worst < huberTolerance {
			out.Scale = huberScale(residuals, scaleScratch)
			// The reported weights follow the answer, not the step
			// that produced it: the collapsed-scale branch refreshes
			// them, and the converged exit must agree with the Scale
			// and Residuals it reports (huberWeight at band 0 is the
			// collapsed rule, one inside and zero outside).
			band := tuning * out.Scale
			for i, r := range residuals {
				out.Weights[i] = huberWeight(r, band)
			}
			converged = true
			iterations = iter
			break
		}
	}
	if !converged {
		return nil, base.Errf("%s: %d iterations did not converge", name, huberMaxIterations)
	}
	out.Iterations = iterations
	out.Converged = converged

	// Fitted and Residuals from the settled coefficients, and the
	// standard errors from σ²·(XᵀWX)⁻¹, its diagonal read from one
	// factorisation of the pristine normal equations against all p unit
	// columns at once, the way LinearRegression reads its own
	// covariance diagonal.
	out.Fitted = make([]float64, n)
	out.Residuals = make([]float64, n)
	copy(out.Residuals, residuals)
	weighted := out.Scale * out.Scale
	wxxPristine := weightedNormal(n, p, x, fx, out.Weights)
	out.StandardErrors = make([]float64, p)
	unit := make([][]float64, p)
	for j := range p {
		unit[j] = make([]float64, p)
		unit[j][j] = 1
	}
	inv, err := base.SolveSystem(name, wxxPristine, unit)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	for j := range p {
		v := weighted * inv[j][j]
		switch {
		case v > 0:
			out.StandardErrors[j] = math.Sqrt(v)
		case v == 0:
			// An exact fit: nothing to estimate, a zero standard error
			// beside the zero residual variance.
			out.StandardErrors[j] = 0
		default:
			return nil, base.Errf("%s: the design is near-collinear: the variance of coefficient %d came out negative (%g)", name, j, v)
		}
	}
	return out, nil
}

// huberWeight is the IRLS weight of one residual for a band: exactly
// one inside the band, the band over the magnitude outside it, which
// is the taper that turns the quadratic loss linear. A collapsed band
// leaves the zero residual at weight one and everything else at zero.
func huberWeight(residual, band float64) float64 {
	a := math.Abs(residual)
	if a <= band {
		return 1
	}
	return band / a
}

// huberScale estimates the robust scale of a residual sample as
// 1.4826 times the median absolute deviation, the constant that makes
// the estimate consistent for the standard deviation of Gaussian
// residuals. It is scale-free against a shift because the median is
// taken twice, once at the centre and once over the deviations.
//
// scratch is a buffer as long as residuals, reused across the
// reweighting rounds: the centre pass sorts it, and the deviation pass
// then writes every slot of it before the second median reads any, so
// the rounds cannot see each other's values.
func huberScale(residuals, scratch []float64) float64 {
	vals := scratch[:len(residuals)]
	copy(vals, residuals)
	centre := medianSlice(vals)
	for i, r := range residuals {
		vals[i] = math.Abs(r - centre)
	}
	return 1.4826 * medianSlice(vals)
}

// medianSlice rearranges vals in place and returns their median,
// averaging the two middle values on even length exactly as Median does:
// the halved-magnitudes sum, immune to the overflow the literal average
// risks. Callers hand over a disposable slice.
//
// The two middle order statistics come from a selection, not from a sort:
// at the Theil-Sen observation cap the slope list holds millions of
// entries and the sort cost more than everything else in the fit put
// together. The selected values are the ones the sort left at the middle
// indices, with one honest difference: a sort is unstable, so among values
// that compare equal but differ in bits (the two zeros, a NaN payload) the
// sort's choice of which one lands at the middle is unspecified, and the
// selection may return the other. Every other input answers the identical
// number.
func medianSlice(vals []float64) float64 {
	n := len(vals)
	if n%2 == 1 {
		selectNth(vals, n/2)
		return vals[n/2]
	}
	selectNth(vals, n/2)
	// Everything before index n/2 is no larger than the upper middle, so
	// the lower middle is the largest of that prefix.
	lo := vals[0]
	for _, v := range vals[:n/2] {
		if v > lo {
			lo = v
		}
	}
	return lo/2 + vals[n/2]/2
}

// selectNth rearranges vals so that the element at index k is the k-th
// smallest, every element before it is no larger, and every element after
// it is no smaller, using Hoare's partition with a median-of-three pivot
// and an insertion sort for short ranges. The pivot choice and the range
// cutoff decide speed only: whichever pivot is taken, the value that ends
// up at k is the k-th order statistic.
func selectNth(vals []float64, k int) {
	lo, hi := 0, len(vals)-1
	for hi-lo > 24 {
		mid := lo + (hi-lo)/2
		if vals[mid] < vals[lo] {
			vals[mid], vals[lo] = vals[lo], vals[mid]
		}
		if vals[hi] < vals[lo] {
			vals[hi], vals[lo] = vals[lo], vals[hi]
		}
		if vals[hi] < vals[mid] {
			vals[hi], vals[mid] = vals[mid], vals[hi]
		}
		pivot := vals[mid]
		i, j := lo, hi
		for i <= j {
			for vals[i] < pivot {
				i++
			}
			for pivot < vals[j] {
				j--
			}
			if i <= j {
				vals[i], vals[j] = vals[j], vals[i]
				i++
				j--
			}
		}
		if k <= j {
			hi = j
			continue
		}
		if k >= i {
			lo = i
			continue
		}
		return
	}
	for i := lo + 1; i <= hi; i++ {
		v := vals[i]
		j := i - 1
		for j >= lo && v < vals[j] {
			vals[j+1] = vals[j]
			j--
		}
		vals[j+1] = v
	}
}

// regressionResiduals recomputes r = y − Xβ into residuals.
func regressionResiduals(n, p int, x, y *core.Array, fx, fy []float64, beta, residuals []float64) {
	for i := range n {
		f := 0.0
		if fx != nil {
			row := fx[i*p : i*p+p]
			for j, xj := range row {
				f += beta[j] * xj
			}
		} else {
			for j := range p {
				f += beta[j] * x.FloatAt(i*p+j)
			}
		}
		var yv float64
		if fy != nil {
			yv = fy[i]
		} else {
			yv = y.FloatAt(i)
		}
		residuals[i] = yv - f
	}
}

// weightedNormal assembles the weighted normal equations matrix
// XᵀWX from a weight vector, or the unweighted XᵀX when w is nil.
// The caller consumes the matrix through the shared solve, which
// factors in place, so every call builds fresh.
func weightedNormal(n, p int, x *core.Array, fx []float64, w []float64) [][]float64 {
	m := make([][]float64, p)
	for i := range p {
		m[i] = make([]float64, p)
	}
	weightedNormalInto(m, n, p, x, fx, w)
	return m
}

// weightedNormalInto assembles XᵀWX (or XᵀX when w is nil) into the
// provided matrix, which the caller has cleared: every entry
// accumulates from zero, so a cleared buffer answers exactly what a
// fresh allocation answered. The accumulation order is the row walk's
// own, unchanged.
func weightedNormalInto(m [][]float64, n, p int, x *core.Array, fx []float64, w []float64) {
	for r := range n {
		wr := 1.0
		if w != nil {
			wr = w[r]
		}
		if fx != nil {
			row := fx[r*p : r*p+p]
			for a, xa := range row {
				ga := wr * xa
				ma := m[a]
				for b, xb := range row {
					ma[b] += ga * xb
				}
			}
		} else {
			for a := range p {
				xa := x.FloatAt(r*p + a)
				ga := wr * xa
				for b := range p {
					m[a][b] += ga * x.FloatAt(r*p+b)
				}
			}
		}
	}
}

// leastSquaresWorkspace carries the normal-equation buffers one
// reweighting loop refills: the shared solve factors its matrix and
// overwrites its right-hand side in place, so every round clears the
// same backing arrays and the accumulation sees the zero state a fresh
// allocation carried.
type leastSquaresWorkspace struct {
	mat [][]float64
	rhs []float64
}

func newLeastSquaresWorkspace(p int) *leastSquaresWorkspace {
	m := make([][]float64, p)
	for i := range p {
		m[i] = make([]float64, p)
	}
	return &leastSquaresWorkspace{mat: m, rhs: make([]float64, p)}
}

// leastSquaresSolve solves the (weighted) normal equations in one
// shot: XᵀWX β = XᵀWy with W the weights, or the plain XᵀX system
// when w is nil. The right-hand side is assembled alongside the
// matrix so both see the same weights. The returned slice is the
// workspace's right-hand side, overwritten in place by the solve: the
// caller must consume it before the workspace is refilled.
func leastSquaresSolve(name string, n, p int, x, y *core.Array, fx, fy []float64, w []float64) ([]float64, error) {
	return leastSquaresSolveWS(name, n, p, x, y, fx, fy, w, newLeastSquaresWorkspace(p))
}

// leastSquaresSolveWS is leastSquaresSolve on a caller-owned
// workspace, for a reweighting loop that solves once per round.
func leastSquaresSolveWS(name string, n, p int, x, y *core.Array, fx, fy []float64, w []float64, ws *leastSquaresWorkspace) ([]float64, error) {
	for i := range p {
		clear(ws.mat[i])
	}
	weightedNormalInto(ws.mat, n, p, x, fx, w)
	rhs := ws.rhs
	clear(rhs)
	for r := range n {
		wr := 1.0
		if w != nil {
			wr = w[r]
		}
		var yv float64
		if fy != nil {
			yv = fy[r]
		} else {
			yv = y.FloatAt(r)
		}
		g := wr * yv
		if fx != nil {
			row := fx[r*p : r*p+p]
			for a, xa := range row {
				rhs[a] += xa * g
			}
		} else {
			for a := range p {
				rhs[a] += x.FloatAt(r*p+a) * g
			}
		}
	}
	solved, err := base.SolveSystem(name, ws.mat, [][]float64{rhs})
	if err != nil {
		return nil, base.Errf("%s: the design is rank deficient (%w)", name, err)
	}
	return solved[0], nil
}

// TheilSenMaxObservations is the exactness contract of
// TheilSenRegression: the median of the pairwise slopes is computed
// over all n(n−1)/2 of them, which at 4096 observations is already
// some eight million slopes and a good fraction of a gigabyte of
// working memory. Beyond the cap the estimator refuses rather than
// silently degrade to a sample of itself; the cost is named in the
// error so the caller can subsample deliberately.
const TheilSenMaxObservations = 4096

// theilSenParallelPairs is the pairwise-slope count from which the walk
// splits across workers: below it a crew costs more to start than the
// walk it would carry, and above it a worker is handed this many
// slopes, so the number of blocks follows the pair count rather than
// the row count.
const theilSenParallelPairs = 1 << 18

// theilSenSlopePoolMax bounds the pair-slope buffer the pool keeps, in
// float64 entries. The cap-sized fit needs 8,386,560 of them, inside
// the bound; a longer list allocates fresh and is dropped on return,
// so one oversized call cannot pin a larger buffer on every
// processor, and sync.Pool forgets what it holds at each garbage
// collection besides.
const theilSenSlopePoolMax = 1 << 23

// theilSenSlopes is the pooled pair-slope buffer. The pointer form
// keeps the Put from boxing a slice header on every return.
type theilSenSlopes struct{ s []float64 }

var theilSenSlopePool = sync.Pool{
	New: func() any { return new(theilSenSlopes) },
}

// TheilSenRegression fits the simple linear model y = a + b·x by the
// Theil-Sen estimator: the slope is the median of the pairwise slopes
// (y_j − y_i)/(x_j − x_i) over all pairs with distinct predictors,
// and the intercept is the median of y_i − b·x_i at that slope. Both
// medians are exact: the slope breaks down only when nearly half the
// points are broken, and one wild observation among hundreds cannot
// drag the answer at all. At least three observations with at least
// two distinct predictors are needed; every input must be finite, and
// samples beyond TheilSenMaxObservations are refused with the cost
// named rather than answered approximately.
func TheilSenRegression(x, y *core.Array) (intercept, slope float64, err error) {
	const name = "TheilSenRegression"
	if x.NDim() != 1 {
		return 0, 0, base.Errf("%s: the predictor must be rank 1, got shape %s", name, base.ShapeText(x.Shape()))
	}
	if y.NDim() != 1 {
		return 0, 0, base.Errf("%s: the response must be rank 1", name)
	}
	if x.Dtype() == core.Complex || y.Dtype() == core.Complex {
		return 0, 0, base.Errf("%s: complex inputs are not supported", name)
	}
	n := x.Len()
	if y.Len() != n {
		return 0, 0, base.Errf("%s: the predictor has %d samples but the response %d", name, n, y.Len())
	}
	if n < 3 {
		return 0, 0, base.Errf("%s: at least three observations are needed, got %d", name, n)
	}
	if n > TheilSenMaxObservations {
		return 0, 0, base.Errf("%s: %d observations would need the exact median over %d pairwise slopes; the exactness contract ends at %d, subsample deliberately instead",
			name, n, n*(n-1)/2, TheilSenMaxObservations)
	}
	if err := checkFinite(name, "the predictor", x); err != nil {
		return 0, 0, err
	}
	if err := checkFinite(name, "the response", y); err != nil {
		return 0, 0, err
	}
	xs := make([]float64, n)
	ys := make([]float64, n)
	if fx := rawFloats(x); fx != nil {
		copy(xs, fx)
	} else {
		for i := range n {
			xs[i] = x.FloatAt(i)
		}
	}
	if fy := rawFloats(y); fy != nil {
		copy(ys, fy)
	} else {
		for i := range n {
			ys[i] = y.FloatAt(i)
		}
	}
	// All pairwise slopes over the pairs whose predictor actually
	// differs: a repeated predictor carries no slope information, and
	// dividing by its zero would poison the median. The slopes land in
	// one buffer in the row walk's own order, so the rows can be filled
	// by a crew and the collected slice is the serial walk's own,
	// element for element; the per-row pair counts are what let the
	// blocks be cut before the walk starts.
	counts := make([]int, n)
	seen := make(map[float64]int, n)
	for i := n - 1; i >= 0; i-- {
		equal := seen[xs[i]]
		seen[xs[i]] = equal + 1
		counts[i] = n - 1 - i - equal
	}
	offsets := make([]int, n+1)
	for i := range n {
		offsets[i+1] = offsets[i] + counts[i]
	}
	kept := offsets[n]
	if kept == 0 {
		return 0, 0, base.Errf("%s: the predictor does not vary, no slope exists", name)
	}
	sb := theilSenSlopePool.Get().(*theilSenSlopes)
	slopes := sb.s
	if cap(slopes) < kept {
		slopes = make([]float64, kept)
	}
	slopes = slopes[:kept]
	fill := func(lo, hi int) {
		for i := lo; i < hi; i++ {
			off := offsets[i]
			xi, yi := xs[i], ys[i]
			for j := i + 1; j < n; j++ {
				if dx := xs[j] - xi; dx != 0 {
					slopes[off] = (ys[j] - yi) / dx
					off++
				}
			}
		}
	}
	// The walk's cost falls with i, so an even split of the rows would
	// leave the first worker with a quarter of the work: the blocks are
	// cut where the pair count crosses an equal share instead.
	parts := min(kept/theilSenParallelPairs, n)
	if parts < 2 {
		fill(0, n)
	} else {
		per := kept / parts
		bounds := make([]int, 1, parts+1)
		for i, cut := 1, per; i < n && len(bounds) < parts; i++ {
			if offsets[i] >= cut {
				bounds = append(bounds, i)
				cut += per
			}
		}
		bounds = append(bounds, n)
		engine.Parallel(len(bounds)-1, func(start, end int) {
			for k := start; k < end; k++ {
				fill(bounds[k], bounds[k+1])
			}
		})
	}
	slope = medianSlice(slopes)
	if cap(slopes) <= theilSenSlopePoolMax {
		sb.s = slopes[:cap(slopes)]
		theilSenSlopePool.Put(sb)
	}
	intercepts := make([]float64, n)
	for i := range n {
		intercepts[i] = ys[i] - slope*xs[i]
	}
	return medianSlice(intercepts), slope, nil
}