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

package optim

import (
	"math"
	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
	"sourcedock.dev/petrbalvin/tensor/internal/engine"
	"sourcedock.dev/petrbalvin/tensor/linalg"
)

// FitStatus states how an iterative fit ended.
type FitStatus int

const (
	// FitConverged marks a run that met one of the tolerances: the
	// relative chi2 improvement, the gradient norm GradTol or the
	// step size StepTol.
	FitConverged FitStatus = iota
	// FitStalled marks a run whose step died: the normal equations
	// turned singular or the damping collapsed without the residual
	// meeting the tolerance. The returned point is the best one the
	// run reached, never a partial step past it.
	FitStalled
	// FitBudget marks a run that spent its iteration budget before
	// any tolerance or stall fired. The returned point is the last
	// iterate.
	FitBudget
)

// FitResult carries everything a fit reports: the point it ended on,
// the (weighted) residual sum of squares there, how the run ended
// and, when requested, the parameter covariance at that point.
type FitResult struct {
	Parameters *core.Array
	Chi2       float64
	Status     FitStatus
	Covariance *core.Array
}

// LMOptions tunes LevenbergMarquardt and LevenbergMarquardtFit.
// Lambda ≤ 0 means 1e-3 (the initial damping factor), Tolerance ≤ 0
// means 1e-10 (the relative χ² improvement threshold) and
// MaxIterations ≤ 0 means 200.
type LMOptions struct {
	MaxIterations int
	Tolerance     float64
	Lambda        float64
	Jacobian      func(p *core.Array) (*core.Array, error)
	// AllowBudgetExit makes a run that exhausts MaxIterations report
	// its last point instead of an error. The default is false, so a
	// budget stop is never mistaken for a converged answer; the flag
	// mirrors LBFGSOptions.AllowBudgetExit. LevenbergMarquardtFit
	// needs no flag: it reports the budget stop as FitBudget.
	AllowBudgetExit bool
	// ParallelJacobian lets the central-difference Jacobian sweep its
	// columns on several goroutines. Setting it is the caller's
	// consent that the residual callback may run concurrently from
	// more than one goroutine: the default false keeps every
	// evaluation on the caller's goroutine, and the fit is bit for bit
	// the same either way, because the columns are independent and
	// each one is differenced by the same stencil. The field does
	// nothing while Jacobian supplies the analytic matrix.
	ParallelJacobian bool
	// GradTol converges the fit once the infinity norm of the
	// gradient Jᵀr falls to it, the test that catches the flat
	// optimum where chi2 still falls in slivers while the step
	// directions carry no information. A value ≤ 0 disables the test,
	// which is the default: a caller who sets it picks the scale.
	GradTol float64
	// StepTol converges the fit once an accepted step's infinity norm
	// falls to StepTol·(‖p‖∞ + StepTol), the relative step test that
	// stops a fit whose parameters have stopped moving meaningfully.
	// A value ≤ 0 disables the test, which is the default.
	StepTol float64
	// Sigma weights the residuals by the measurement covariance: a
	// vector holds one positive variance per residual, a square
	// matrix holds the full nR×nR covariance and must be exactly
	// symmetric and positive definite. The fit whitens the residuals
	// and the Jacobian through the factor once, chi2 becomes rᵀC⁻¹r
	// and a requested covariance becomes (JᵀC⁻¹J)⁻¹. Nil, the
	// default, leaves every residual unweighted.
	Sigma *core.Array
	// RequestCovariance fills FitResult.Covariance with (JᵀJ)⁻¹ at
	// the returned point, or (JᵀC⁻¹J)⁻¹ under Sigma. The answer costs
	// one more Jacobian at the final point, which on the
	// difference route is two residual evaluations per parameter. A
	// Jacobian that is rank-deficient at the returned point has no
	// covariance to report and the run fails naming that, so a
	// caller asking for a covariance accepts the trade on a fit it
	// expects to stall.
	RequestCovariance bool
}

// LevenbergMarquardt minimises ‖r(p)‖² by LM damping of the
// Gauss-Newton step, with a central-difference or user-supplied
// analytic Jacobian and the library's LU solver for the normal
// equations.
//
// The historical error contract stays: a run that stalls (a singular
// solve, a collapsed damping) or spends its budget without
// AllowBudgetExit is an error, not a point. LevenbergMarquardtFit
// reports those conditions as a FitResult instead and carries the
// gradient, step, weight and covariance options.
func LevenbergMarquardt(residual func(*core.Array) (*core.Array, error), p0 *core.Array, opts LMOptions) (*core.Array, float64, error) {
	res, err := runLevenbergMarquardt(residual, p0, opts, true)
	if err != nil {
		return nil, 0, err
	}
	return res.Parameters, res.Chi2, nil
}

// LevenbergMarquardtFit is LevenbergMarquardt with the full report:
// the status that says how the run ended and, on request, the
// parameter covariance. Where LevenbergMarquardt keeps its historical
// error contract, this one reports every ended run as a result: a
// stalled step and a spent budget come back as FitStalled and
// FitBudget on the best point reached, never as an error. The errors
// here are the model's own fault (a residual that fails or turns
// non-finite at a state the fit adopts, a malformed Sigma) and, with
// RequestCovariance, a rank-deficient Jacobian at the answer.
func LevenbergMarquardtFit(residual func(*core.Array) (*core.Array, error), p0 *core.Array, opts LMOptions) (*FitResult, error) {
	return runLevenbergMarquardt(residual, p0, opts, false)
}

// denseFloats returns the array's float64 payload when a is a dense
// float64 array and nil otherwise: hot loops branch once on the result
// and sweep the payload directly, falling back to the widening
// accessor for views and other dtypes. The elements are identical
// either way, so a dense sweep computes the same bits as the accessor
// walk it replaces.
func denseFloats(a *core.Array) []float64 {
	if !a.Strided() && a.Dtype() == core.Float {
		return a.RawFloats()
	}
	return nil
}

// whitener carries the Sigma factor: the per-residual divisors of a
// variance vector, or the lower Cholesky factor of a full covariance.
// A nil whitener is the unweighted fit.
type whitener struct {
	diag   []float64
	factor [][]float64
}

// sigmaWhitener validates Sigma against the residual count nR and
// factors it. The matrix form is checked for exact symmetry before
// the factorisation reads one triangle: an asymmetric partner would
// silently weight by a matrix the caller did not pass.
func sigmaWhitener(sigma *core.Array, nR int) (*whitener, error) {
	if sigma == nil {
		return nil, nil
	}
	if err := requireReal("LevenbergMarquardt", "Sigma", sigma); err != nil {
		return nil, err
	}
	if sigma.NDim() == 1 {
		if sigma.Len() != nR {
			return nil, base.Errf("LevenbergMarquardt: Sigma must hold one variance per residual (%d), got %d", nR, sigma.Len())
		}
		w := &whitener{diag: make([]float64, nR)}
		for i := range nR {
			v := sigma.FloatAt(i)
			if math.IsNaN(v) || v <= 0 {
				return nil, base.Errf("LevenbergMarquardt: Sigma must hold positive variances, got %g at %d", v, i)
			}
			w.diag[i] = math.Sqrt(v)
		}
		return w, nil
	}
	if sigma.NDim() != 2 || sigma.Shape()[0] != nR || sigma.Shape()[1] != nR {
		return nil, base.Errf("LevenbergMarquardt: Sigma must be a %d×%d covariance or a vector of %d variances, got shape %s",
			nR, nR, nR, base.ShapeText(sigma.Shape()))
	}
	for i := range nR {
		for j := i + 1; j < nR; j++ {
			up, lo := sigma.FloatAt(i*nR+j), sigma.FloatAt(j*nR+i)
			if up != lo {
				return nil, base.Errf("LevenbergMarquardt: Sigma must be symmetric, got %g and %g at (%d, %d)", up, lo, i, j)
			}
		}
	}
	l, err := linalg.Cholesky(sigma)
	if err != nil {
		return nil, base.Errf("LevenbergMarquardt: Sigma is not positive definite: %w", err)
	}
	w := &whitener{factor: make([][]float64, nR)}
	for i := range nR {
		row := make([]float64, i+1)
		for j := range i + 1 {
			row[j] = l.FloatAt(i*nR + j)
		}
		w.factor[i] = row
	}
	return w, nil
}

// vector whitens a residual in place: y solves L y = r.
func (w *whitener) vector(r []float64) {
	if w == nil {
		return
	}
	if w.diag != nil {
		for i := range r {
			r[i] /= w.diag[i]
		}
		return
	}
	for i := range r {
		s := r[i]
		li := w.factor[i]
		for j := range i {
			s -= li[j] * r[j]
		}
		r[i] = s / li[i]
	}
}

// matrix whitens a Jacobian in place: the rows solve L J' = J, so the
// downstream normal equations accumulate JᵀC⁻¹J without knowing a
// weight exists.
func (w *whitener) matrix(jac [][]float64) {
	if w == nil {
		return
	}
	for i := range jac {
		if w.diag != nil {
			for j := range jac[i] {
				jac[i][j] /= w.diag[i]
			}
			continue
		}
		li := w.factor[i]
		for j := range jac[i] {
			s := jac[i][j]
			for k := range i {
				s -= li[k] * jac[k][j]
			}
			jac[i][j] = s / li[i]
		}
	}
}

// covarianceFromJac inverts the unweighted normal equations of the
// (whitened) Jacobian, which is the parameter covariance. The solve
// runs column by column against the identity and the answer is
// symmetrised explicitly: a pivoted LU on a symmetric matrix may
// leave last-bit asymmetry the covariance must not carry.
func covarianceFromJac(name string, jac [][]float64, nP int) (*core.Array, error) {
	a := make([][]float64, nP)
	for i := range a {
		a[i] = make([]float64, nP)
	}
	for k := range len(jac) {
		row := jac[k]
		for i := range nP {
			xi := row[i]
			ai := a[i]
			for j := range nP {
				ai[j] += xi * row[j]
			}
		}
	}
	rhs := make([][]float64, nP)
	for i := range nP {
		rhs[i] = make([]float64, nP)
		rhs[i][i] = 1
	}
	x, err := base.SolveSystem(name, a, rhs)
	if err != nil {
		return nil, base.Errf("%s: the Jacobian is rank-deficient at the answer, so no covariance exists: %w", name, err)
	}
	out := core.New(core.Float, nP, nP)
	v := out.RawFloats()
	for i := range nP {
		for j := range nP {
			v[i*nP+j] = (x[i][j] + x[j][i]) / 2
		}
	}
	return out, nil
}

// runLevenbergMarquardt carries the fit. The legacy flag restores the
// historical error contract of LevenbergMarquardt: the same stalls
// the FitResult reports come back as errors with the messages the
// package has always published, so existing callers see nothing move.
func runLevenbergMarquardt(residual func(*core.Array) (*core.Array, error), p0 *core.Array, opts LMOptions, legacy bool) (*FitResult, error) {
	if p0.Dtype() == core.Complex {
		return nil, base.Errf("LevenbergMarquardt: complex parameters are not supported")
	}
	nP := p0.Len()
	if nP == 0 {
		return nil, base.Errf("LevenbergMarquardt: the parameter vector must not be empty")
	}
	if opts.MaxIterations <= 0 {
		opts.MaxIterations = 200
	}
	if opts.Tolerance <= 0 {
		opts.Tolerance = 1e-10
	}
	if opts.Lambda <= 0 {
		opts.Lambda = 1e-3
	}

	// cloneDense promotes through FloatAt, so Int and Float32 starting
	// vectors behave exactly like Float64 ones (a RawFloats copy would
	// silently start the fit from zeros for those dtypes).
	p := cloneDense(p0)
	nR := 0
	// whiten carries the Sigma factor; it is still nil for the very
	// first evaluation, and the residual it returns is whitened by
	// hand right after the factor is built.
	var whiten *whitener
	// evalR reads the residual at pp. The finiteness gate is strict
	// for the states the fit adopts (the start point and every
	// accepted iterate): a non-finite residual there poisons chi2 and
	// every comparison against it, and the fit would die later as a
	// bogus "the damping collapsed" diagnosis instead of the model's
	// own fault. Backtracking trials take the lenient variant: a step
	// into a saturating model is a candidate to damp past, not a dead
	// run, the same recovery FindRootSystem's trials make. The Sigma
	// whitening lands here, so every downstream consumer (chi2, the
	// difference stencil, the trial comparison) works on the whitened
	// residual and the weighted fit is the unweighted one on whitened
	// data.
	evalR := func(pp []float64, strict bool, dst []float64) ([]float64, error) {
		a := linalg.ArrayFromFloatsSafe(pp, nP)
		r, err := residual(a)
		if err != nil {
			return nil, err
		}
		if r.NDim() != 1 {
			return nil, base.Errf("LevenbergMarquardt: the residual must be a vector")
		}
		if err := requireReal("LevenbergMarquardt", "residuals", r); err != nil {
			return nil, err
		}
		if nR != 0 && r.Len() != nR {
			return nil, base.Errf("LevenbergMarquardt: the residual length changed from %d to %d mid-fit", nR, r.Len())
		}
		// dst carries a buffer the stencil reuses across columns; the
		// states the fit keeps come back freshly allocated. Every entry
		// of the buffer is written before it is read.
		res := dst
		if cap(res) < r.Len() {
			res = make([]float64, r.Len())
		}
		res = res[:r.Len()]
		// Both branches fill res with the identical elements: the dense
		// sweep reads the payload the accessor walk would widen.
		if fs := denseFloats(r); fs != nil {
			if strict {
				for i, v := range fs {
					if math.IsNaN(v) || math.IsInf(v, 0) {
						return nil, base.Errf("LevenbergMarquardt: the residual returned the non-finite value %g at %d", v, i)
					}
				}
			}
			copy(res, fs)
		} else {
			for i := range r.Len() {
				v := r.FloatAt(i)
				if strict && (math.IsNaN(v) || math.IsInf(v, 0)) {
					return nil, base.Errf("LevenbergMarquardt: the residual returned the non-finite value %g at %d", v, i)
				}
				res[i] = v
			}
		}
		whiten.vector(res)
		return res, nil
	}
	r, rerr := evalR(p, true, nil)
	if rerr != nil {
		return nil, base.Errf("LevenbergMarquardt: %w", rerr)
	}
	nR = len(r)
	if nR < nP {
		return nil, base.Errf("LevenbergMarquardt: underdetermined (%d obs, %d params)", nR, nP)
	}
	whiten, werr := sigmaWhitener(opts.Sigma, nR)
	if werr != nil {
		return nil, werr
	}
	whiten.vector(r)
	chi2 := 0.0
	for i := range nR {
		chi2 += r[i] * r[i]
	}

	// buildJacobian assembles the row-major Jacobian at p, either from
	// the caller's analytic callback or by central differences on the
	// residual, one column per parameter. Its storage is allocated
	// once and refilled per iteration: the sweep writes every entry.
	jac := make([][]float64, nR)
	for i := range nR {
		jac[i] = make([]float64, nP)
	}
	// The difference stencils and the two residual vectors the columns
	// are differenced from, allocated on first use and carried across
	// the whole fit: a fit with an analytic Jacobian pays for none of
	// them. The stencil carries the offset on one parameter at a time,
	// restored as soon as the column is done, so neither a copy of the
	// whole parameter vector nor a residual slice per column is needed.
	var pp, pm []float64
	var resPlus, resMinus []float64
	buildJacobian := func(p []float64) error {
		if opts.Jacobian != nil {
			jm, err := opts.Jacobian(linalg.ArrayFromFloatsSafe(p, nP))
			if err != nil {
				return base.Errf("LevenbergMarquardt: %w", err)
			}
			if jm.NDim() != 2 || jm.Shape()[0] != nR || jm.Shape()[1] != nP {
				return base.Errf("LevenbergMarquardt: the Jacobian must be a %d×%d matrix, got shape %s",
					nR, nP, base.ShapeText(jm.Shape()))
			}
			if err := requireReal("LevenbergMarquardt", "Jacobians", jm); err != nil {
				return err
			}
			if fs := denseFloats(jm); fs != nil {
				for i := range nR {
					copy(jac[i], fs[i*nP:(i+1)*nP])
				}
			} else {
				for i := range nR {
					ji := jac[i]
					for j := range nP {
						ji[j] = jm.FloatAt(i*nP + j)
					}
				}
			}
			whiten.matrix(jac)
			return nil
		}
		// Central differences: evalR copies into the array handed to the
		// callback, so nothing observes later mutation. column walks one
		// parameter's stencil and writes that column of jac, and nothing
		// else, so the bits it produces do not depend on which driver
		// walks the columns.
		column := func(j int, sp, sm, rp, rm []float64) ([]float64, []float64, error) {
			eps := math.Sqrt(base.EpsF) * math.Max(1, math.Abs(p[j]))
			sp[j] += eps
			sm[j] -= eps
			rp, re1 := evalR(sp, true, rp)
			rm, re2 := evalR(sm, true, rm)
			sp[j], sm[j] = p[j], p[j]
			if re1 != nil || re2 != nil {
				return rp, rm, firstError(re1, re2)
			}
			for i := range nR {
				jac[i][j] = (rp[i] - rm[i]) / (2 * eps)
			}
			return rp, rm, nil
		}
		if opts.ParallelJacobian {
			// The consent the option records lets the columns go to the
			// engine's workers: each goroutine owns a disjoint run of
			// columns, writes only into those columns of jac and reads
			// only p, so no two workers write the same address and the
			// sweep needs no locks. A failing column records its error
			// and abandons the rest of its run; the reported one is the
			// lowest failing column, the one the serial walk would hit
			// first. The residual buffers live per worker instead of
			// being carried across columns: the option exists for
			// expensive callbacks, where the carry buys nothing.
			colErrs := make([]error, nP)
			engine.ParallelMin(nP, 1, func(start, end int) {
				sp, sm := make([]float64, nP), make([]float64, nP)
				copy(sp, p)
				copy(sm, p)
				var rp, rm []float64
				for j := start; j < end; j++ {
					var err error
					rp, rm, err = column(j, sp, sm, rp, rm)
					if err != nil {
						colErrs[j] = err
						return
					}
				}
			})
			for _, err := range colErrs {
				if err != nil {
					return err
				}
			}
			return nil
		}
		if cap(pp) < nP {
			pp, pm = make([]float64, nP), make([]float64, nP)
		}
		pp, pm = pp[:nP], pm[:nP]
		copy(pp, p)
		copy(pm, p)
		for j := range nP {
			var err error
			resPlus, resMinus, err = column(j, pp, pm, resPlus, resMinus)
			if err != nil {
				return err
			}
		}
		return nil
	}

	// The normal equations' storage, reused across iterations: the
	// upper triangle of a is refilled by accumulation from an explicit
	// zero and its lower one is mirrored back, bv is cleared likewise,
	// and every other buffer is fully overwritten before it is read.
	a := make([][]float64, nP)
	for i := range nP {
		a[i] = make([]float64, nP)
	}
	bv := make([]float64, nP)
	pNew := make([]float64, nP)
	// One right-hand-side header for the whole fit: the solve writes
	// the step through bv in place, so the wrapper never changes.
	solveRHS := [][]float64{bv}
	lambda := opts.Lambda
	status := FitBudget
	var solveErr error
	var collapseAt float64

	// buildResult packs the fit's answer at the current p. The
	// covariance rebuilds the Jacobian there: the loop's last one
	// belongs to the point the last accepted step left behind, and the
	// covariance must describe the point it is published beside.
	buildResult := func() (*FitResult, error) {
		out := core.New(core.Float, nP)
		copy(out.RawFloats(), p)
		res := &FitResult{Parameters: out, Chi2: chi2, Status: status}
		if opts.RequestCovariance {
			if jerr := buildJacobian(p); jerr != nil {
				return nil, jerr
			}
			cov, cerr := covarianceFromJac("LevenbergMarquardt", jac, nP)
			if cerr != nil {
				return nil, cerr
			}
			res.Covariance = cov
		}
		return res, nil
	}

	if chi2 == 0 {
		// A start whose residual cancels exactly is already the perfect
		// fit: the improvement test below is strict and cannot accept
		// the zero step it produces, so the run would die in the
		// damping collapse for being perfect.
		status = FitConverged
		return buildResult()
	}

	for iter := 0; iter < opts.MaxIterations; iter++ {
		if jerr := buildJacobian(p); jerr != nil {
			return nil, jerr
		}

		// a = JᵀJ + λ·diag(JᵀJ), bv = −Jᵀr. The accumulation walks the
		// rows in ascending order, so each entry sums the same
		// products in the same order the column-wise walk visited.
		// The off-diagonal pair (i, j) and (j, i) sums the same
		// products in the same order, the product commuting bitwise,
		// so the pass runs the upper triangle alone and the mirror
		// below reproduces the lower one exactly.
		for i := range nP {
			ai := a[i]
			for j := i; j < nP; j++ {
				ai[j] = 0
			}
		}
		clear(bv)
		for k := range nR {
			row := jac[k]
			rk := r[k]
			for i := range nP {
				bv[i] -= row[i] * rk
			}
			for i := range nP {
				xi := row[i]
				ai := a[i]
				for j := i; j < nP; j++ {
					ai[j] += xi * row[j]
				}
			}
		}
		for i := range nP {
			ai := a[i]
			for j := i + 1; j < nP; j++ {
				a[j][i] = ai[j]
			}
		}
		for i := range nP {
			a[i][i] *= (1 + lambda)
		}
		if opts.GradTol > 0 {
			// The gradient test runs on bv before the solve: ‖bv‖∞ is
			// ‖Jᵀr‖∞, and a gradient this small says the parameter
			// directions carry nothing the step could spend, which is
			// exactly the flat optimum the chi2 test alone never
			// reaches.
			gInf := 0.0
			for i := range nP {
				gInf = max(gInf, math.Abs(bv[i]))
			}
			if gInf <= opts.GradTol {
				status = FitConverged
				return buildResult()
			}
		}

		delta, derr := base.SolveSystem("LevenbergMarquardt", a, solveRHS)
		if derr != nil {
			// A singular solve leaves the current point standing: it
			// was good enough to build normal equations from, and no
			// step replaced it. The fit reports it and stops.
			status = FitStalled
			solveErr = derr
			break
		}

		for j := range nP {
			pNew[j] = p[j] + delta[0][j]
		}
		// A trial point: a saturating model here is damped past, and a
		// finite chi2New admits only finite components, so an accepted
		// trial never carries poison into the fit state.
		rNew, rerr := evalR(pNew, false, nil)
		if rerr != nil {
			return nil, base.Errf("LevenbergMarquardt: %w", rerr)
		}
		chi2New := 0.0
		for i := range nR {
			chi2New += rNew[i] * rNew[i]
		}

		if chi2New < chi2 {
			// The relative step test runs on the step just accepted,
			// against the scale of the point it left: a step this small
			// says the parameters have stopped moving meaningfully,
			// whatever the residual still promises.
			var stepSmall bool
			if opts.StepTol > 0 {
				stepInf, pInf := 0.0, 0.0
				for j := range nP {
					stepInf = max(stepInf, math.Abs(delta[0][j]))
					pInf = max(pInf, math.Abs(p[j]))
				}
				stepSmall = stepInf <= opts.StepTol*(pInf+opts.StepTol)
			}
			// The tolerance break must accept the better point too:
			// reporting chi2New beside the old p publishes a fit quality
			// the returned parameters do not achieve.
			chi2Old := chi2
			copy(p, pNew)
			r = rNew
			chi2 = chi2New
			if chi2Old-chi2 < opts.Tolerance*(1+chi2Old) {
				status = FitConverged
				return buildResult()
			}
			lambda *= 0.3
			if stepSmall {
				status = FitConverged
				return buildResult()
			}
		} else {
			lambda *= 10
			if lambda > 1e20 {
				// A damping that collapsed has no step left to take:
				// the last point is a stall, never a converged answer.
				status = FitStalled
				collapseAt = lambda
				break
			}
		}
	}

	if legacy {
		switch status {
		case FitStalled:
			if solveErr != nil {
				return nil, base.Errf("LevenbergMarquardt: %w", solveErr)
			}
			return nil, base.Errf("LevenbergMarquardt: the damping collapsed to %g without the residual meeting the tolerance", collapseAt)
		case FitBudget:
			if !opts.AllowBudgetExit {
				return nil, base.Errf("LevenbergMarquardt: the iteration budget of %d ran out without the residual meeting the tolerance", opts.MaxIterations)
			}
		case FitConverged:
		}
	}
	return buildResult()
}