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

package stats

import (
	"math"

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

// Regularised linear regression by coordinate descent: the lasso and
// its elastic net generalisation. The fit runs on the standardised
// design (every column centred and scaled to unit population standard
// deviation, the response centred), the penalty falls on the slopes
// and never on the intercept, and the standardisation is recorded on
// the result and inverted before the coefficients leave, so the
// numbers a caller receives apply to the design as supplied. The path
// over a documented lambda grid is warm-started: every fit after the
// first begins at the previous lambda's coefficients, which is why
// the path costs a fraction of the cold fits it replaces.

// ElasticNetResult carries one regularised fit at a single lambda.
type ElasticNetResult struct {
	// Intercept and Coefficients are the fit on the original scale of
	// the design as supplied: the standardisation the solver worked on
	// has been inverted. Coefficients holds one slope per design
	// column, in order.
	Intercept    float64
	Coefficients []float64
	// Fitted and Residuals align with the rows of the design.
	Fitted    []float64
	Residuals []float64
	// ColumnMeans and ColumnScales record the standardisation the fit
	// ran under: the solver saw (x − ColumnMeans)/ColumnScales, and
	// the lambda is measured against the centred response in the
	// response's own units.
	ColumnMeans  []float64
	ColumnScales []float64
	// Lambda and Alpha are the penalty and the L2 mixing actually used.
	Lambda float64
	Alpha  float64
	// Iterations counts full coordinate-descent cycles and Converged
	// reports whether no slope moved more than the tolerance in the
	// last cycle. An exhausted budget returns the fit found so far
	// rather than an error: coordinate descent on this convex
	// objective cannot diverge, so the iterate stays a usable answer
	// and the flag says how much to trust it.
	Iterations int
	Converged  bool
}

// LassoPathResult carries a warm-started regularisation path.
type LassoPathResult struct {
	// Alpha is the mixing the path ran with.
	Alpha float64
	// Lambdas are the penalty values of the path in descending order.
	// The documented grid is lassoGridSteps values log-spaced from
	// lambdaMax down to lambdaMax/1000, where lambdaMax is the
	// smallest penalty at which the lasso keeps every slope at zero,
	// max_j |Σ_i x_ij (y_i − ȳ)|/n divided by alpha. For alpha = 0
	// the pure-lasso lambdaMax (alpha taken as 1) defines the same
	// grid, so ridge walks exactly the path the lasso would.
	Lambdas []float64
	// Intercepts and Coefficients hold one fit per lambda, on the
	// original scale, in the order of Lambdas.
	Intercepts   []float64
	Coefficients [][]float64
	// Iterations counts the coordinate-descent cycles each lambda
	// needed. The warm start is what keeps them low along the path:
	// each fit but the first begins at the previous lambda's
	// coefficients, far closer to the answer than the zero start a
	// cold fit uses, and the counts fall accordingly.
	Iterations []int
	// Converged reports whether every fit on the path converged
	// within the iteration budget.
	Converged bool
}

// The documented shape of the lambda grid and the documented stopping
// rule of the coordinate descent: lassoGridSteps log-spaced values
// covering three decades down from lambdaMax, and a fit is settled
// once no slope moves more than lassoTolerance in a full cycle, with
// lassoMaxIterations cycles the most any single fit may spend.
const (
	lassoGridSteps = 100
	lassoGridRatio = 1e-3
	lassoTolerance = 1e-8
	lassoMaxIter   = 10000
)

// Lasso fits y by linear regression under the least absolute
// shrinkage penalty at a single lambda: the pure lasso, alpha = 1 of
// ElasticNet. The slopes and only the slopes are penalised; the
// intercept is never. See ElasticNet for the contract the two share.
func Lasso(x, y *core.Array, lambda float64) (*ElasticNetResult, error) {
	return ElasticNet(x, y, lambda, 1)
}

// ElasticNet fits y = X·β under the elastic net penalty
//
//	(1/2n)·Σᵢ(yᵢ − β₀ − xᵢ·β)² + λ·(α·Σ|βⱼ| + (1−α)/2·Σβⱼ²),
//
// the convex combination of the lasso (alpha 1) and ridge (alpha 0)
// penalties, by coordinate descent on the standardised design. The
// design carries n rows and p columns exactly as LinearRegression's,
// but a rank-deficient design is not an error here: the penalty keeps
// the problem well posed, and more columns than rows is a legitimate
// use. n must be at least 2, the design must carry at least one
// non-constant column, alpha must lie in [0, 1] and lambda must not
// be negative; complex or non-finite input is an error.
//
// Two exact anchors hold. On an orthonormal design the lasso answer
// is the soft-thresholded projection, coefficient by coefficient, and
// the fit reproduces that closed form. At alpha = 0 the answer is the
// Tikhonov ridge system (ZᵀZ/n + λI)β = Zᵀ(y − ȳ)/n over the
// standardised design, and the fit agrees with the answer the shared
// LU solve delivers for that system.
func ElasticNet(x, y *core.Array, lambda, alpha float64) (*ElasticNetResult, error) {
	const name = "ElasticNet"
	if math.IsNaN(lambda) || lambda < 0 {
		return nil, base.Errf("%s: lambda must not be negative, got %g", name, lambda)
	}
	if math.IsNaN(alpha) || alpha < 0 || alpha > 1 {
		return nil, base.Errf("%s: alpha must lie in [0, 1], got %g", name, alpha)
	}
	n, p, fx, fy, err := lassoInputs(name, x, y)
	if err != nil {
		return nil, err
	}
	std, err := lassoStandardise(name, n, p, x, y, fx, fy)
	if err != nil {
		return nil, err
	}
	beta, iterations, converged := elasticNetFit(std, lambda, alpha, nil, &lassoWorkspace{})
	return lassoResult(name, n, p, x, y, fx, fy, std, beta, lambda, alpha, iterations, converged)
}

// LassoPath fits the whole regularisation path over the documented
// lambda grid, warm-started: the first fit, at the largest lambda,
// starts from zero coefficients (which are the exact answer there for
// the lasso), and every later fit starts at the previous lambda's
// answer. alpha is the elastic net mixing as in ElasticNet, and the
// alpha = 0 ridge walks the same grid the pure lasso defines. The
// result records how many coordinate cycles each lambda spent, which
// is the evidence the warm start earns its keep.
func LassoPath(x, y *core.Array, alpha float64) (*LassoPathResult, error) {
	const name = "LassoPath"
	if math.IsNaN(alpha) || alpha < 0 || alpha > 1 {
		return nil, base.Errf("%s: alpha must lie in [0, 1], got %g", name, alpha)
	}
	n, p, fx, fy, err := lassoInputs(name, x, y)
	if err != nil {
		return nil, err
	}
	std, err := lassoStandardise(name, n, p, x, y, fx, fy)
	if err != nil {
		return nil, err
	}
	// lambdaMax: the smallest penalty the lasso answers with an
	// all-zero slope vector, divided by alpha because the elastic net
	// subgradient carries only the alpha fraction as the L1 term. The
	// ridge (alpha = 0) has no sparsity threshold to define a largest
	// lambda, so it borrows the pure-lasso grid and walks the same
	// path. A response orthogonal to every column gives lambdaMax 0:
	// no penalty does anything, and the grid value is immaterial, so a
	// unit grid keeps the documented shape.
	gridAlpha := alpha
	if gridAlpha == 0 {
		gridAlpha = 1
	}
	// The projection norms are computed with the same multiply by the
	// reciprocal the coordinate sweep uses, so at the top of the grid
	// the soft threshold compares identical bits against identical
	// bits and the all-zero answer is exact, not nearly zero.
	inv := 1 / float64(n)
	lambdaMax := 0.0
	for j := range p {
		col := std.xs[j]
		s := 0.0
		for i := range n {
			s += col[i] * std.yc[i]
		}
		if a := math.Abs(s*inv) / gridAlpha; a > lambdaMax {
			lambdaMax = a
		}
	}
	if lambdaMax == 0 {
		lambdaMax = 1
	}
	out := &LassoPathResult{
		Alpha:        alpha,
		Lambdas:      make([]float64, lassoGridSteps),
		Intercepts:   make([]float64, lassoGridSteps),
		Coefficients: make([][]float64, lassoGridSteps),
		Iterations:   make([]int, lassoGridSteps),
		Converged:    true,
	}
	warm := make([]float64, p)
	ws := &lassoWorkspace{beta: make([]float64, p), r: make([]float64, n)}
	for k := range lassoGridSteps {
		lambda := lambdaMax * math.Pow(lassoGridRatio, float64(k)/float64(lassoGridSteps-1))
		beta, iterations, converged := elasticNetFit(std, lambda, alpha, warm, ws)
		copy(warm, beta)
		// The path records the fit itself and has no field for a
		// residual: lifting the standardised slopes back to the design
		// as supplied is the whole reading, so the per-lambda pass over
		// the original design that Fitted and Residuals would need is
		// not run.
		intercept, coefficients := lassoCoefficients(std, beta)
		out.Lambdas[k] = lambda
		out.Intercepts[k] = intercept
		out.Coefficients[k] = coefficients
		out.Iterations[k] = iterations
		out.Converged = out.Converged && converged
	}
	return out, nil
}

// lassoStandardisation holds the standardised system the coordinate
// descent runs on: standardised columns stored column-wise for cache
// friendly sweeps, their squared norms, and the centred response.
type lassoStandardisation struct {
	xs     [][]float64 // the standardised columns, each of length n
	zs     []float64   // z_j = Σᵢ x²ᵢⱼ/n, the standardised column energies
	yc     []float64   // the centred response
	yMean  float64     // the response mean the intercept reabsorbs
	means  []float64   // the original column means
	scales []float64   // the original column population standard deviations
}

// lassoWorkspace holds the coordinate descent's running buffers: the
// slope vector and the residual the sweeps keep current. One
// workspace serves a whole path, so a hundred warm-started fits
// allocate none of them again.
type lassoWorkspace struct {
	beta []float64
	r    []float64
}

// lassoInputs validates a design and response exactly as the rest of
// the estimation entries do, and returns the row and column counts
// with the dense raw payloads when they exist.
func lassoInputs(name string, x, y *core.Array) (n, p int, fx, fy []float64, err error) {
	if x.NDim() != 2 {
		return 0, 0, nil, nil, base.Errf("%s: the design must be rank 2, got shape %s", name, base.ShapeText(x.Shape()))
	}
	if y.NDim() != 1 {
		return 0, 0, nil, nil, base.Errf("%s: the response must be rank 1", name)
	}
	if x.Dtype() == core.Complex || y.Dtype() == core.Complex {
		return 0, 0, nil, nil, base.Errf("%s: complex inputs are not supported", name)
	}
	n, p = x.Shape()[0], x.Shape()[1]
	if y.Len() != n {
		return 0, 0, nil, nil, base.Errf("%s: the design has %d rows but the response %d", name, n, y.Len())
	}
	if n < 2 {
		return 0, 0, nil, nil, base.Errf("%s: at least two observations are needed, got %d", name, n)
	}
	if p == 0 {
		return 0, 0, nil, nil, base.Errf("%s: the design must carry at least one column", name)
	}
	if err := checkFinite(name, "the design", x); err != nil {
		return 0, 0, nil, nil, err
	}
	if err := checkFinite(name, "the response", y); err != nil {
		return 0, 0, nil, nil, err
	}
	if rawFloats(x) != nil {
		fx = rawFloats(x)
	}
	if rawFloats(y) != nil {
		fy = rawFloats(y)
	}
	return n, p, fx, fy, nil
}

// lassoStandardise builds the standardised system: every column
// centred and divided by its population standard deviation, the
// response centred. A constant column has no standard deviation to
// divide by and is refused: its slope is unidentified at any penalty.
func lassoStandardise(name string, n, p int, x, y *core.Array, fx, fy []float64) (*lassoStandardisation, error) {
	std := &lassoStandardisation{
		xs:     make([][]float64, p),
		zs:     make([]float64, p),
		yc:     make([]float64, n),
		means:  make([]float64, p),
		scales: make([]float64, p),
	}
	inv := 1 / float64(n)
	for j := range p {
		mean := 0.0
		if fx != nil {
			for i := range n {
				mean += fx[i*p+j]
			}
		} else {
			for i := range n {
				mean += x.FloatAt(i*p + j)
			}
		}
		mean *= inv
		variance := 0.0
		if fx != nil {
			for i := range n {
				d := fx[i*p+j] - mean
				variance += d * d
			}
		} else {
			for i := range n {
				d := x.FloatAt(i*p+j) - mean
				variance += d * d
			}
		}
		scale := math.Sqrt(variance * inv)
		if scale == 0 {
			return nil, base.Errf("%s: design column %d is constant and cannot be standardised", name, j)
		}
		col := make([]float64, n)
		if fx != nil {
			for i := range n {
				col[i] = (fx[i*p+j] - mean) / scale
			}
		} else {
			for i := range n {
				col[i] = (x.FloatAt(i*p+j) - mean) / scale
			}
		}
		std.xs[j] = col
		std.means[j] = mean
		std.scales[j] = scale
	}
	yMean := 0.0
	// The walk is bounded by the response's own element count: a rebased
	// view's payload may run past its visible elements.
	if fy != nil {
		for _, v := range fy[:n] {
			yMean += v
		}
	} else {
		for i := range n {
			yMean += y.FloatAt(i)
		}
	}
	yMean *= inv
	std.yMean = yMean
	for i := range n {
		var yv float64
		if fy != nil {
			yv = fy[i]
		} else {
			yv = y.FloatAt(i)
		}
		std.yc[i] = yv - yMean
	}
	for j := range p {
		s := 0.0
		for i := range n {
			s += std.xs[j][i] * std.xs[j][i]
		}
		std.zs[j] = s * inv
	}
	return std, nil
}

// elasticNetFit runs the coordinate descent on the standardised
// system. start, when given, is the warm start the path hands down;
// a nil start begins at zero, the exact answer at the top of the
// grid. The running buffers live in ws and are reused across calls:
// the path fits one lambda after another into the same workspace, and
// a single fit simply owns one it just allocated. Each cycle sweeps
// the coordinates in order, moving each slope to the exact minimiser
// of the objective along its own coordinate and folding the movement
// into the running residual, until a full cycle moves nothing by more
// than the tolerance. The per-coordinate minimiser is the
// soft-thresholded least-squares coordinate over the denominator the
// elastic net puts there: pure lasso thresholding when alpha is 1,
// plain ridge division when alpha is 0.
func elasticNetFit(std *lassoStandardisation, lambda, alpha float64, start []float64, ws *lassoWorkspace) (beta []float64, iterations int, converged bool) {
	n := len(std.yc)
	p := len(std.xs)
	inv := 1 / float64(n)
	if start == nil {
		beta = append(ws.beta[:0], make([]float64, p)...)
	} else {
		beta = append(ws.beta[:0], start...)
	}
	ws.beta = beta
	// The running residual r = yc − Xβ, kept current by folding each
	// coordinate's movement in as it happens: the sweep then reads the
	// partial residual every coordinate needs without a full product
	// per coordinate.
	r := append(ws.r[:0], std.yc...)
	ws.r = r
	if start != nil {
		for j := range p {
			bj := beta[j]
			if bj == 0 {
				continue
			}
			col := std.xs[j]
			for i := range n {
				r[i] -= col[i] * bj
			}
		}
	}
	for iter := 1; iter <= lassoMaxIter; iter++ {
		worst := 0.0
		for j := range p {
			col := std.xs[j]
			zj := std.zs[j]
			s := 0.0
			for i := range n {
				s += col[i] * r[i]
			}
			rho := s*inv + zj*beta[j]
			newBeta := softThreshold(rho, lambda*alpha) / (zj + lambda*(1-alpha))
			if newBeta != beta[j] {
				d := newBeta - beta[j]
				for i := range n {
					r[i] -= col[i] * d
				}
				beta[j] = newBeta
				if a := math.Abs(d); a > worst {
					worst = a
				}
			}
		}
		if worst < lassoTolerance {
			return beta, iter, true
		}
	}
	return beta, lassoMaxIter, false
}

// softThreshold is the proximal operator of the absolute value: the
// identity past the threshold and zero inside it. It is what makes
// the lasso exact on an orthonormal design, where the coordinates
// decouple and every slope is this function of its own projection.
func softThreshold(v, t float64) float64 {
	if v > t {
		return v - t
	}
	if v < -t {
		return v + t
	}
	return 0
}

// lassoCoefficients lifts the standardised slopes back to the design as
// supplied: the slope of a standardised column scales by the column's
// standard deviation, and the intercept reabsorbs the column means, in
// the columns' own order, so the accumulation is the one the fit's
// reporting pass runs.
func lassoCoefficients(std *lassoStandardisation, betaStd []float64) (intercept float64, coefficients []float64) {
	p := len(betaStd)
	coefficients = make([]float64, p)
	intercept = std.yMean
	for j := range p {
		coefficients[j] = betaStd[j] / std.scales[j]
		intercept -= coefficients[j] * std.means[j]
	}
	return intercept, coefficients
}

// lassoResult lifts the standardised slopes back to the design as
// supplied: the slope of a standardised column scales by the column's
// standard deviation, and the intercept reabsorbs the column means.
// Fitted and Residuals are then computed on the original design, the
// same final sweep WeightedLinearRegression runs.
func lassoResult(name string, n, p int, x, y *core.Array, fx, fy []float64, std *lassoStandardisation, betaStd []float64, lambda, alpha float64, iterations int, converged bool) (*ElasticNetResult, error) {
	intercept, coefficients := lassoCoefficients(std, betaStd)
	out := &ElasticNetResult{
		Intercept:    intercept,
		Coefficients: coefficients,
		ColumnMeans:  std.means,
		ColumnScales: std.scales,
		Lambda:       lambda,
		Alpha:        alpha,
		Iterations:   iterations,
		Converged:    converged,
	}
	out.Fitted = make([]float64, n)
	out.Residuals = make([]float64, n)
	for i := range n {
		f := out.Intercept
		if fx != nil {
			row := fx[i*p : i*p+p]
			for j, xj := range row {
				f += coefficients[j] * xj
			}
		} else {
			for j := range p {
				f += coefficients[j] * x.FloatAt(i*p+j)
			}
		}
		var yv float64
		if fy != nil {
			yv = fy[i]
		} else {
			yv = y.FloatAt(i)
		}
		out.Fitted[i] = f
		out.Residuals[i] = yv - f
	}
	return out, nil
}