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

// Quantile regression: the linear fit that asks not for the mean of
// the response but for its tau-th conditional quantile, by minimising
// the check loss
//
//	ρ_τ(r) = r·τ when r ≥ 0, r·(τ − 1) when r < 0,
//
// the asymmetric absolute loss that rewards a fitted line for putting
// the right fraction of the data beneath it. The implementation is the
// Frisch-Newton interior-point method on the dual, the form
// Portnoy and Koenker put the method in: the check loss problem is
// the linear program
//
//	min Σᵢ τ·uᵢ + (1−τ)·vᵢ  subject to  u − v = y − Xβ, u, v ≥ 0,
//
// and its dual asks for
//
//	max wᵀy  subject to  Xᵀw = (1−τ)·Xᵀ1,  w ∈ [0, 1]ⁿ.
//
// At the optimum an observation with a positive residual carries w = 1,
// one with a negative residual w = 0, and the observations the fit
// reproduces exactly carry w strictly inside the box. A logarithmic
// barrier is laid on the box, every iteration solves the barrier's
// Newton system exactly in the p×p form XᵀD⁻¹X through the shared LU
// solve, and the barrier parameter falls geometrically. The start is
// exactly feasible, w = 1−τ for every observation, and the primal fit
// is read back off the interior set at every iteration, so the run
// records a monotone descent of the true check loss.

// The documented schedule of the interior-point loop: the barrier
// parameter starts at the mean absolute residual of the ordinary
// least squares start, shrinks by this factor every iteration, and the
// run is settled when no coefficient of the dual moves by more than
// the tolerance, or when the parameter has fallen sixteen orders of
// magnitude, whichever comes first.
const (
	quantileMaxIterations = 100
	quantileMuShrink      = 0.25
	quantileMuFloorRatio  = 1e-16
	quantileStepFraction  = 0.995
	quantileTolerance     = 1e-14
)

// QuantileRegressionResult carries a quantile regression fit.
type QuantileRegressionResult struct {
	// Coefficients are the quantile estimates β̂, one per design
	// column, in the design's own order. The intercept, supplied by
	// the caller as a constant column, is estimated like any other
	// coefficient: the check loss pulls it to the response's tau-th
	// quantile at x = 0.
	Coefficients []float64
	// Fitted and Residuals align with the rows of the design.
	Fitted    []float64
	Residuals []float64
	// Tau is the quantile the fit minimises the check loss for.
	Tau float64
	// CheckLoss is the minimised check loss Σᵢ ρ_τ(rᵢ) at the fit.
	CheckLoss float64
	// Objective records the best check loss seen after every
	// interior-point iteration, starting from the ordinary least
	// squares start. It is the instrument that shows the optimisation
	// descending: monotone non-increasing by construction, because an
	// iterate only enters the record by improving on every iterate
	// before it.
	Objective []float64
	// Iterations counts the interior-point iterations taken;
	// Converged reports whether the run settled by its own stopping
	// rules.
	Iterations int
	Converged  bool
}

// QuantileRegression fits y = X·β for the tau-th conditional quantile
// by the Frisch-Newton interior-point method on the dual of the check
// loss program. 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. tau must lie strictly inside
// (0, 1): the closed ends have no regression answer, tau = 0 and
// tau = 1 being the envelope of the data rather than a fit.
//
// The run starts from the ordinary least squares fit and follows the
// barrier path in the dual, where the equality constraint is satisfied
// exactly from the first iterate to the last. The primal fit is read
// back off the dual's interior set, the observations the optimal fit
// reproduces exactly, by least squares over that set; a degenerate
// optimum, whose interior set is underdetermined, is resolved to the
// smallest-norm member of the optimal face by a ridge of the order of
// the solve's own rounding.
func QuantileRegression(x, y *core.Array, tau float64) (*QuantileRegressionResult, error) {
	const name = "QuantileRegression"
	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)
	}
	// NaN compares false against both bounds, so this refuses it too.
	if !(tau > 0 && tau < 1) {
		return nil, base.Errf("%s: tau must lie strictly inside (0, 1), got %g", name, tau)
	}
	if err := checkFinite(name, "the design", x); err != nil {
		return nil, err
	}
	if err := checkFinite(name, "the response", y); err != nil {
		return nil, err
	}
	fx := rawFloats(x)
	fy := rawFloats(y)
	if fy == nil {
		// The response is a view or a narrower dtype: widen it once so
		// the barrier loops below sweep a plain slice, the same bits
		// the widening accessor returns.
		fy = make([]float64, n)
		for i := range n {
			fy[i] = y.FloatAt(i)
		}
	}

	// The start is the ordinary least squares answer. It is not the
	// optimum of any asymmetric loss, but it is inside the basin, its
	// residual scale sets the barrier's starting parameter, and a
	// design the shared LU cannot solve is refused here, rank
	// deficiency and all.
	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)
	meanAbs := 0.0
	for _, r := range residuals {
		meanAbs += math.Abs(r)
	}
	meanAbs /= float64(n)
	lossAt := func(b []float64) float64 {
		total := 0.0
		for i := range n {
			r := fy[i]
			for j := range p {
				var xj float64
				if fx != nil {
					xj = fx[i*p+j]
				} else {
					xj = x.FloatAt(i*p + j)
				}
				r -= b[j] * xj
			}
			if r >= 0 {
				total += r * tau
			} else {
				// r and tau − 1 are both negative, so the product is
				// the positive |r|·(1 − tau) the check loss asks for.
				total += r * (tau - 1)
			}
		}
		return total
	}
	loss := lossAt(beta)
	out := &QuantileRegressionResult{
		Coefficients: beta,
		Tau:          tau,
		CheckLoss:    loss,
		Objective:    []float64{loss},
	}
	if meanAbs == 0 {
		// The ordinary least squares fit reproduces the response
		// exactly: the check loss is zero and no asymmetric loss can
		// beat zero. The run is settled before it begins.
		out.Fitted = make([]float64, n)
		out.Residuals = make([]float64, n)
		copy(out.Residuals, residuals)
		for i := range n {
			out.Fitted[i] = fy[i]
		}
		out.Converged = true
		return out, nil
	}

	// The dual barrier state. w = 1−τ for every observation satisfies
	// the equality Xᵀw = (1−τ)·Xᵀ1 identically and sits exactly in the
	// middle of the box, the interior point the method asks for.
	mu := meanAbs
	muFloor := mu * quantileMuFloorRatio
	w := make([]float64, n)
	dw := make([]float64, n)
	gradient := make([]float64, n)
	curvature := make([]float64, n)
	cols := make([]float64, n)
	colSums := make([]float64, p)
	// The backtracking buffer holds only complete candidate vectors: every
	// iteration rewrites all n slots before the objective reads any, so one
	// buffer serves the whole run.
	trial := make([]float64, n)
	for i := range n {
		w[i] = 1 - tau
	}
	for i := range n {
		for j := range p {
			var xj float64
			if fx != nil {
				xj = fx[i*p+j]
			} else {
				xj = x.FloatAt(i*p + j)
			}
			colSums[j] += xj
		}
	}
	bestBeta := append([]float64(nil), beta...)
	bestLoss := loss
	converged := false
	iterations := quantileMaxIterations
	// The barrier system's storage rides the whole run: the upper
	// triangle is refilled by accumulation from an explicit zero and
	// its mirror copies the lower one, and the right-hand wrapper is
	// fixed because the solve writes through rhs in place.
	normal := make([][]float64, p)
	for j := range p {
		normal[j] = make([]float64, p)
	}
	rhs := make([]float64, p)
	solveRHS := [][]float64{rhs}
	// The barrier objective Φ = wᵀy + μΣ(log w + log(1−w)) rises
	// along the Newton direction of this concave problem by
	// ∇ΦᵀΔw = ΔwᵀDΔw ≥ 0 identically, so a backtracking half-step
	// finds an ascent all the way to the barrier maximum.
	objective := func(wv []float64, muv float64) float64 {
		total := 0.0
		for i := range n {
			total += wv[i]*fy[i] + muv*(math.Log(wv[i])+math.Log(1-wv[i]))
		}
		return total
	}
	for iter := 1; iter <= quantileMaxIterations; iter++ {
		// The Newton system of the barrier problem, eliminated to the
		// p×p form XᵀD⁻¹X Δλ = XᵀD⁻¹z: z the barrier gradient
		// y + μ(1/w − 1/(1−w)), D⁻¹ the reciprocal of the barrier's
		// curvature μ((1−w)² + w²)/(w²(1−w)²). The solve is exact, the
		// equality direction XᵀΔw = 0 comes out of it, and the step
		// keeps the equality exactly because it started exact.
		for j := range p {
			nj := normal[j]
			for k := j; k < p; k++ {
				nj[k] = 0
			}
		}
		clear(rhs)
		for i := range n {
			if fx != nil {
				for j := range p {
					cols[j] = fx[i*p+j]
				}
			} else {
				for j := range p {
					cols[j] = x.FloatAt(i*p + j)
				}
			}
			oneMinus := 1 - w[i]
			gradient[i] = fy[i] + mu*(1/w[i]-1/oneMinus)
			curvature[i] = w[i] * w[i] * oneMinus * oneMinus /
				(mu * (oneMinus*oneMinus + w[i]*w[i]))
			scale := gradient[i] * curvature[i]
			for j := range p {
				rhs[j] += cols[j] * scale
				for k := j; k < p; k++ {
					normal[j][k] += cols[j] * cols[k] * curvature[i]
				}
			}
		}
		for j := range p {
			for k := range j {
				normal[j][k] = normal[k][j]
			}
		}
		solved, err := base.SolveSystem(name, normal, solveRHS)
		if err != nil {
			return nil, base.Errf("%s: the barrier system is singular (%w)", name, err)
		}
		deltaLambda := solved[0]
		mostMove := 0.0
		for i := range n {
			pred := 0.0
			if fx != nil {
				for j := range p {
					pred += fx[i*p+j] * deltaLambda[j]
				}
			} else {
				for j := range p {
					pred += x.FloatAt(i*p+j) * deltaLambda[j]
				}
			}
			dw[i] = curvature[i] * (gradient[i] - pred)
			if m := math.Abs(dw[i]); m > mostMove {
				mostMove = m
			}
		}
		// Fraction to the boundary: the largest step that keeps every
		// w strictly inside the box, taken at a documented fraction of
		// it and never more than the full Newton step.
		t := 1.0
		for i := range n {
			if dw[i] < 0 {
				t = min(t, quantileStepFraction*(-w[i]/dw[i]))
			}
			if dw[i] > 0 {
				t = min(t, quantileStepFraction*((1-w[i])/dw[i]))
			}
		}
		// The barrier objective rises along the Newton direction of
		// this concave problem (the note above the closure), so a
		// backtracking half-step finds an ascent all the way to the
		// barrier maximum.
		current := objective(w, mu)
		for {
			for i := range n {
				trial[i] = w[i] + t*dw[i]
			}
			if objective(trial, mu) >= current || t < quantileTolerance {
				break
			}
			t /= 2
		}
		if t < quantileTolerance || mostMove < quantileTolerance {
			// The barrier problem is stationary at this μ: further
			// pressure buys nothing until μ falls, and μ is about to.
			// The run settles when the barrier itself is exhausted.
			converged = true
			iterations = iter
			break
		}
		copy(w, trial)
		// The primal fit read off this iterate's interior set, and the
		// monotone record it feeds.
		candidate, err := quantileRecover(name, n, p, x, y, fx, fy, w)
		if err != nil {
			return nil, err
		}
		if l := lossAt(candidate); l < bestLoss {
			bestLoss = l
			copy(bestBeta, candidate)
		}
		out.Objective = append(out.Objective, bestLoss)
		mu *= quantileMuShrink
		if mu < muFloor {
			converged = true
			iterations = iter
			break
		}
	}
	if !converged {
		return nil, base.Errf("%s: %d iterations did not converge", name, quantileMaxIterations)
	}
	out.Iterations = iterations
	out.Converged = true
	out.CheckLoss = bestLoss
	copy(out.Coefficients, bestBeta)
	out.Fitted = make([]float64, n)
	out.Residuals = make([]float64, n)
	regressionResiduals(n, p, x, y, fx, fy, out.Coefficients, out.Residuals)
	for i := range n {
		out.Fitted[i] = fy[i] - out.Residuals[i]
	}
	return out, nil
}

// quantileRecover reads a primal fit off a dual barrier point. The
// observations whose w sits strictly inside the box are the ones the
// optimal fit reproduces exactly, their residual being zero, so the
// least squares over that interior set is the fit; the degenerate
// case, an interior set that underdetermines the coefficients, is
// resolved to the smallest-norm member of the optimal face by a ridge
// a trillionth of the normal equations' own scale, far below the
// solve's meaningful digits. An empty interior set falls back to the
// plain least squares fit, the same reading the first iterate's
// all-interior box gives.
func quantileRecover(name string, n, p int, x, y *core.Array, fx, fy []float64, w []float64) ([]float64, error) {
	const band = 1e-5
	normal := make([][]float64, p)
	for j := range p {
		normal[j] = make([]float64, p)
	}
	rhs := make([]float64, p)
	trace := 0.0
	count := 0
	for i := range n {
		if !(w[i] > band && w[i] < 1-band) {
			continue
		}
		count++
		var yv float64
		if fy != nil {
			yv = fy[i]
		} else {
			yv = y.FloatAt(i)
		}
		for a := range p {
			var xa float64
			if fx != nil {
				xa = fx[i*p+a]
			} else {
				xa = x.FloatAt(i*p + a)
			}
			rhs[a] += xa * yv
			for b := range p {
				var xb float64
				if fx != nil {
					xb = fx[i*p+b]
				} else {
					xb = x.FloatAt(i*p + b)
				}
				normal[a][b] += xa * xb
			}
		}
	}
	if count < p {
		// An interior set too small to identify the coefficients: the
		// whole box was effectively at its bounds, and the plain least
		// squares fit is the honest reading. The caller's
		// best-by-check-loss record keeps this from ever worsening the
		// answer.
		return leastSquaresSolve(name, n, p, x, y, fx, fy, nil)
	}
	for a := range p {
		trace += normal[a][a]
	}
	for a := range p {
		normal[a][a] += 1e-12 * trace / float64(p)
	}
	solved, err := base.SolveSystem(name, normal, [][]float64{rhs})
	if err != nil {
		return nil, base.Errf("%s: the interior set is degenerate (%w)", name, err)
	}
	return solved[0], nil
}