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

package optim

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

import "math"

// Limited-memory BFGS optimisation. The plain `Minimise` (Nelder-Mead)
// is derivative-free and robust but needs O(n²) memory for the simplex
// and many objective evaluations per step. L-BFGS is the standard
// quasi-Newton method for smooth objectives on moderate-to-large
// problems: it approximates the inverse Hessian from the last `memory`
// position/gradient pairs by the two-loop recursion, which costs
// O(memory·n) per step, far cheaper than the O(n²) simplex or the
// O(n³) dense Hessian, and far faster to converge than gradient
// descent on ill-conditioned objectives.
//
// The gradient may be supplied by the caller or computed by central
// finite differences, which costs one extra objective evaluation per
// coordinate per step.

// LBFGSOptions tunes the optimiser. MaxIterations ≤ 0 means 10000,
// Tolerance ≤ 0 means 1e-8 (L∞ norm of the projected gradient, see
// below), Memory ≤ 0 means 10.
//
// The tolerance is absolute in the gradient's own scale: an objective
// whose gradient sits many orders of magnitude below one is reported
// converged at its start point. Rescale the objective to O(1) before
// calling MinimiseLBFGS when the natural units are not of that size.
//
// Lower and Upper, when non-nil, bound the search coordinate-wise:
// an infinite entry leaves that side open and a crossed or NaN pair of
// walls is an error. The iterate is then kept feasible by projecting
// the line search onto the box, and a coordinate sitting on a wall the
// gradient pushes against is pinned: the two-loop recursion and the
// curvature pairs run in the free subspace, and the convergence
// measure is the projected gradient, which is the KKT residual of the
// box problem. A gradient by finite differences turns one-sided at a
// wall, so the objective is never evaluated outside the box. x0 is
// projected onto the box rather than refused.
type LBFGSOptions struct {
	MaxIterations int
	Tolerance     float64
	Memory        int
	Lower         []float64
	Upper         []float64
	// AllowBudgetExit makes a run that exhausts MaxIterations report
	// the best point it reached with a nil error instead of refusing
	// it. MinimiseConstrained sets it for its inner solves, whose
	// accuracy the outer loop's feasibility check judges; a direct
	// caller leaves it false, so a budget stop is never mistaken for a
	// converged answer. A run that converges normally is unaffected.
	AllowBudgetExit bool
}

// MinimiseLBFGS returns the point and value of a local minimum of f
// near x0 by the limited-memory BFGS method with backtracking Armijo
// line search. The gradient grad may be nil, in which case central
// finite differences are used (costing two extra objective evaluations
// per coordinate per step).
func MinimiseLBFGS(f func(*core.Array) (float64, error), grad func(*core.Array) (*core.Array, error),
	x0 *core.Array, opts LBFGSOptions) (*core.Array, float64, error) {
	n := x0.Len()
	if n == 0 {
		return nil, 0, base.Errf("MinimiseLBFGS: the starting point must have at least one element")
	}
	if x0.Dtype() == core.Complex {
		return nil, 0, base.Errf("MinimiseLBFGS: complex starting points are not supported")
	}
	if opts.MaxIterations <= 0 {
		opts.MaxIterations = 10000
	}
	if opts.Tolerance <= 0 {
		opts.Tolerance = 1e-8
	}
	if opts.Memory <= 0 {
		opts.Memory = 10
	}
	mem := min(opts.Memory, n)

	lower, upper, err := boundsOf(opts, n)
	if err != nil {
		return nil, 0, base.Errf("MinimiseLBFGS: %w", err)
	}

	x := make([]float64, n)
	for i := range n {
		x[i] = min(max(x0.FloatAt(i), lower[i]), upper[i])
	}
	fx, err := f(linalg.ArrayFromFloatsSafe(x, n))
	if err != nil {
		return nil, 0, base.Errf("MinimiseLBFGS: %w", err)
	}
	if math.IsNaN(fx) || math.IsInf(fx, 0) {
		return nil, 0, base.Errf("MinimiseLBFGS: the objective returned the non-finite value %g at the starting point", fx)
	}
	g := make([]float64, n)
	xp := make([]float64, n)
	xm := make([]float64, n)
	if err := evalGrad(f, grad, x, g, lower, upper, xp, xm); err != nil {
		return nil, 0, base.Errf("MinimiseLBFGS: %w", err)
	}
	pinnedMask := make([]bool, n)
	projected := make([]float64, n)
	project := func() float64 {
		// The projected gradient zeroes a coordinate that sits on a
		// wall the gradient pushes against: its norm is the KKT
		// residual of the box problem, its support the free subspace
		// the two-loop recursion runs in. The true gradient stays
		// untouched in g.
		norm := 0.0
		for i := range n {
			pinnedMask[i] = pinnedAt(lower, upper, x, g, i)
			if pinnedMask[i] {
				projected[i] = 0
			} else {
				projected[i] = g[i]
			}
			if v := math.Abs(projected[i]); v > norm {
				norm = v
			}
		}
		return norm
	}
	gNorm := project()
	if gNorm <= opts.Tolerance {
		out, fv := packResult(x, fx)
		return out, fv, nil
	}

	// Limited-memory correction pairs: (s_i, y_i) with ρ_i = 1/(y_iᵀs_i).
	// The window is a ring of pairs, each pair's vectors allocated with
	// the first pair that needs them: count is how many of
	// its slots hold a pair, head is the oldest of them, and the pair at
	// position d of the window sits at (head+d) % mem, position 0 being
	// the oldest. A step that arrives once the window is full overwrites
	// the slot it drops, so an accepted step allocates nothing.
	history := make([]lmPair, mem)
	head, count := 0, 0

	// Scratch reused across iterations: the two-loop direction, its
	// per-pair store, the accepted-step buffers and the finite-difference
	// step and its two stencils. Each is fully overwritten before it is
	// read, and nothing escapes the iteration by reference.
	q := make([]float64, n)
	store := make([]float64, mem)
	dir := q
	trial := make([]float64, n)
	xNew := make([]float64, n)
	gNew := make([]float64, n)
	sVec := make([]float64, n)
	yVec := make([]float64, n)

	converged := false
	for iter := 0; iter < opts.MaxIterations; iter++ {
		gNorm = project()
		if gNorm <= opts.Tolerance {
			converged = true
			break
		}
		// Two-loop recursion for the search direction, driven by the
		// projected gradient and confined to the free subspace: the
		// dots skip pinned coordinates and q is re-zeroed on them
		// after every update, because history pairs predate the
		// current active set and their stale components would drag a
		// pinned coordinate off its wall or corrupt the slope test.
		twoLoopRecursion(projected, history, head, count, mem, pinnedMask, q, store)
		// The search direction is −q (the negative approximate
		// inverse Hessian times the gradient); q is dead past this
		// point, so the negation happens in place.
		for j := range n {
			dir[j] = -q[j]
		}
		dirNorm := maxAbs(dir)
		if dirNorm == 0 {
			return nil, 0, base.Errf("MinimiseLBFGS: the search direction vanished with a projected gradient of %g", gNorm)
		}

		// Backtracking Armijo line search along the free subspace,
		// projected onto the box so every trial point stays feasible.
		// The trial step starts at one unit, the standard choice, and
		// halves until the objective's sufficient decrease is met: the
		// step that actually reduces a quadratic is ≈ 1/L for a
		// curvature L, so a stiff objective (penalty weights, mixed
		// units) needs a long halving chain before any trial is
		// acceptable. A search that cannot find an acceptable point
		// within the budget is a failure, never a silent success: the
		// caller would otherwise read the start point as the answer.
		step := 1.0
		slope := 0.0
		for j := range n {
			slope += g[j] * dir[j]
		}
		if slope >= 0 {
			return nil, 0, base.Errf("MinimiseLBFGS: the search direction is not a descent direction (slope %g) at a projected gradient of %g",
				slope, gNorm)
		}
		const (
			armijo    = 1e-4
			backtrack = 60
		)
		accepted := false
		var ft float64
		var ferr error
		for range backtrack {
			for j := range n {
				trial[j] = min(max(x[j]+step*dir[j], lower[j]), upper[j])
			}
			ft, ferr = f(linalg.ArrayFromFloatsSafe(trial, n))
			if ferr != nil {
				step *= 0.5
				continue
			}
			if ft <= fx+armijo*step*slope {
				accepted = true
				break
			}
			step *= 0.5
		}
		if !accepted {
			if ferr != nil {
				// The last trial failed by error, not by value: the
				// caller sees the cause instead of a bare stall.
				return nil, 0, base.Errf("MinimiseLBFGS: the line search stalled at a step of %g with a projected gradient of %g: %w",
					step, gNorm, ferr)
			}
			return nil, 0, base.Errf("MinimiseLBFGS: the line search stalled at a step of %g with a projected gradient of %g",
				step, gNorm)
		}
		if math.IsNaN(ft) || math.IsInf(ft, 0) {
			return nil, 0, base.Errf("MinimiseLBFGS: the objective returned the non-finite value %g during the line search", ft)
		}
		// The accepted point is the projected trial, not the unprojected
		// step: the displacement s must be what actually moved. Its
		// value is the line search's own: the trial already sat at
		// xNew, so a second evaluation would ask the objective for the
		// number it just returned.
		copy(xNew, trial)
		fxNew := ft
		if err := evalGrad(f, grad, xNew, gNew, lower, upper, xp, xm); err != nil {
			return nil, 0, base.Errf("MinimiseLBFGS: %w", err)
		}
		ys := 0.0
		for j := range n {
			sVec[j] = xNew[j] - x[j]
			yVec[j] = gNew[j] - g[j]
			// A coordinate the projection held still carries no
			// curvature information; its gradient change is noise the
			// second loop would otherwise mix into the direction.
			if sVec[j] == 0 {
				yVec[j] = 0
			}
			ys += yVec[j] * sVec[j]
		}
		if ys > 1e-16 {
			if count < mem {
				p := &history[(head+count)%mem]
				if p.s == nil {
					// The vectors arrive with the first pair that needs
					// them, so a solve that takes few steps does not pay
					// for the whole window up front.
					p.s = make([]float64, n)
					p.y = make([]float64, n)
				}
				copy(p.s, sVec)
				copy(p.y, yVec)
				p.rho = 1 / ys
				count++
			} else {
				p := &history[head]
				copy(p.s, sVec)
				copy(p.y, yVec)
				p.rho = 1 / ys
				head = (head + 1) % mem
			}
		}
		copy(x, xNew)
		copy(g, gNew)
		fx = fxNew
	}
	// The last accepted step updated x, g and fx after the loop-top
	// test, so a run whose tolerance was met exactly on the final
	// iteration must re-test before the budget refusal reports it.
	gNorm = project()
	if gNorm <= opts.Tolerance {
		converged = true
	}
	// Falling out of the loop means the budget ran out, not that a
	// minimum was found: the projected gradient is still above the
	// tolerance, and reporting the point as a converged answer is the
	// silent-wrongness the other exits refuse. The best point is not
	// returned, exactly as the stall and direction exits do not return
	// theirs.
	if !converged {
		if !opts.AllowBudgetExit {
			return nil, 0, base.Errf("MinimiseLBFGS: the iteration budget of %d ran out with a projected gradient of %g, above the tolerance %g",
				opts.MaxIterations, gNorm, opts.Tolerance)
		}
		out, fv := packResult(x, fx)
		return out, fv, nil
	}
	out, fv := packResult(x, fx)
	return out, fv, nil
}

// lmPair is one limited-memory correction pair: the position
// difference s, the gradient difference y and ρ = 1/(yᵀs).
type lmPair struct {
	s, y []float64
	rho  float64
}

// twoLoopRecursion drives the search direction through the stored
// pairs, driven by the projected gradient q starts from. When no
// coordinate is pinned the four sweeps run mask-free, which is the
// common case for an interior iterate; with pins the masked walk
// skips them in the dots and re-zeroes q on them after every update.
// Both walks produce the same bits for the same state: the mask-free
// path is the masked one with an always-false branch removed.
func twoLoopRecursion(projected []float64, history []lmPair, head, count, mem int, pinned []bool, q, store []float64) {
	n := len(projected)
	copy(q, projected)
	store = store[:count]
	free := 0
	for j := range n {
		if !pinned[j] {
			free++
		}
	}
	if free == n {
		for i := count - 1; i >= 0; i-- {
			p := &history[(head+i)%mem]
			dot := 0.0
			for j := range n {
				dot += p.s[j] * q[j]
			}
			store[i] = dot * p.rho
			for j := range n {
				q[j] -= store[i] * p.y[j]
			}
		}
	} else {
		for i := count - 1; i >= 0; i-- {
			p := &history[(head+i)%mem]
			dot := 0.0
			for j := range n {
				if !pinned[j] {
					dot += p.s[j] * q[j]
				}
			}
			store[i] = dot * p.rho
			for j := range n {
				q[j] -= store[i] * p.y[j]
				if pinned[j] {
					q[j] = 0
				}
			}
		}
	}
	gamma := 1.0
	if count > 0 {
		last := &history[(head+count-1)%mem]
		num, den := 0.0, 0.0
		for j := range n {
			num += last.s[j] * last.y[j]
			den += last.y[j] * last.y[j]
		}
		if den != 0 {
			gamma = num / den
		}
	}
	for j := range n {
		q[j] *= gamma
	}
	if free == n {
		for i := range count {
			p := &history[(head+i)%mem]
			dot := 0.0
			for j := range n {
				dot += p.y[j] * q[j]
			}
			coeff := store[i] - p.rho*dot
			for j := range n {
				q[j] += coeff * p.s[j]
			}
		}
		return
	}
	for i := range count {
		p := &history[(head+i)%mem]
		dot := 0.0
		for j := range n {
			if !pinned[j] {
				dot += p.y[j] * q[j]
			}
		}
		coeff := store[i] - p.rho*dot
		for j := range n {
			q[j] += coeff * p.s[j]
			if pinned[j] {
				q[j] = 0
			}
		}
	}
}

// boundsOf resolves the optional box walls: a nil slice opens every
// side, an infinite entry opens one side, and NaN walls or a crossed
// pair are errors.
func boundsOf(opts LBFGSOptions, n int) (lower, upper []float64, err error) {
	lower = make([]float64, n)
	upper = make([]float64, n)
	for i := range n {
		lower[i], upper[i] = math.Inf(-1), math.Inf(1)
	}
	if opts.Lower != nil {
		if len(opts.Lower) != n {
			return nil, nil, base.Errf("the lower bounds hold %d entries for %d variables", len(opts.Lower), n)
		}
		copy(lower, opts.Lower)
	}
	if opts.Upper != nil {
		if len(opts.Upper) != n {
			return nil, nil, base.Errf("the upper bounds hold %d entries for %d variables", len(opts.Upper), n)
		}
		copy(upper, opts.Upper)
	}
	for i := range n {
		if math.IsNaN(lower[i]) || math.IsNaN(upper[i]) || lower[i] > upper[i] {
			return nil, nil, base.Errf("variable %d has bounds [%g, %g]", i, lower[i], upper[i])
		}
	}
	return lower, upper, nil
}

// pinnedAt reports whether coordinate i sits on a wall the gradient
// pushes against: at a lower wall with a positive gradient the descent
// direction -g points out of the box, so the KKT condition holds and
// the coordinate is pinned. The wall test carries a relative tolerance
// wide enough to cover the drift of accepted steps that landed beside
// the wall without crossing it; release needs only the gradient to
// turn favourable, so a generous tolerance cannot trap a coordinate.
func pinnedAt(lower, upper, x, g []float64, i int) bool {
	const wall = 1e-7
	// An equality-bounded coordinate can never move, whatever the
	// gradient says: pinning it keeps the two-loop direction and the
	// projected-gradient norm consistent with the frozen wall.
	if !math.IsInf(lower[i], 0) && lower[i] == upper[i] {
		return true
	}
	if g[i] > 0 && !math.IsInf(lower[i], 0) && x[i]-lower[i] <= wall*math.Max(1, math.Abs(lower[i])) {
		return true
	}
	if g[i] < 0 && !math.IsInf(upper[i], 0) && upper[i]-x[i] <= wall*math.Max(1, math.Abs(upper[i])) {
		return true
	}
	return false
}

// evalGrad fills g with the gradient of f at x: the caller-supplied
// gradient if available, else finite differences with step
// √ε·max(1,|xᵢ|). A coordinate whose central stencil would leave the
// box falls back to the one-sided stencil that stays inside it, so a
// bounded objective is never evaluated outside its domain. xp and xm
// are scratch of length len(x): the perturbed copy carries the offset
// on one coordinate only, which is restored once the coordinate is
// done, so neither a copy of the whole point nor a fresh slice per
// coordinate is needed. linalg.ArrayFromFloatsSafe copies into the
// array handed to f, so the objective never observes later mutation.
func evalGrad(f func(*core.Array) (float64, error), grad func(*core.Array) (*core.Array, error),
	x, g, lower, upper, xp, xm []float64) error {
	if grad != nil {
		ga, err := grad(linalg.ArrayFromFloatsSafe(x, len(x)))
		if err != nil {
			return err
		}
		if err := requireReal("MinimiseLBFGS", "gradients", ga); err != nil {
			return err
		}
		if ga.Len() != len(x) {
			return base.Errf("MinimiseLBFGS: the gradient callback returned %d elements for %d variables",
				ga.Len(), len(x))
		}
		for i := range x {
			g[i] = ga.FloatAt(i)
			// A non-finite gradient coordinate must be an error, not a
			// NaN that skips the projected-gradient max and reads as
			// converged.
			if math.IsNaN(g[i]) || math.IsInf(g[i], 0) {
				return base.Errf("MinimiseLBFGS: the gradient callback returned the non-finite value %g at coordinate %d", g[i], i)
			}
		}
		return nil
	}
	copy(xp, x)
	copy(xm, x)
	for i := range x {
		eps := math.Sqrt(base.EpsF) * math.Max(1, math.Abs(x[i]))
		xp[i] = min(x[i]+eps, upper[i])
		xm[i] = max(x[i]-eps, lower[i])
		if xp[i] == xm[i] {
			// An equality-bounded coordinate clamps both stencil
			// points to the same wall; the central difference would
			// divide 0/0 into a NaN gradient.
			g[i] = 0
			xp[i], xm[i] = x[i], x[i]
			continue
		}
		fp, err := f(linalg.ArrayFromFloatsSafe(xp, len(x)))
		if err != nil {
			return err
		}
		fm, err := f(linalg.ArrayFromFloatsSafe(xm, len(x)))
		if err != nil {
			return err
		}
		g[i] = (fp - fm) / (xp[i] - xm[i])
		xp[i], xm[i] = x[i], x[i]
		// The finite-difference path validates what the analytic one
		// does: a NaN the stencil produces is invisible to the
		// projected-gradient max and would read as converged.
		if math.IsNaN(g[i]) || math.IsInf(g[i], 0) {
			return base.Errf("MinimiseLBFGS: the objective returned the non-finite values %g and %g around coordinate %d", fp, fm, i)
		}
	}
	return nil
}

// maxAbs returns the maximum absolute value of a slice.
func maxAbs(v []float64) float64 {
	m := 0.0
	for _, a := range v {
		if x := math.Abs(a); x > m {
			m = x
		}
	}
	return m
}

// packResult wraps a flat vector as a rank-1 core.Array.
func packResult(x []float64, fx float64) (*core.Array, float64) {
	out := core.New(core.Float, len(x))
	copy(out.RawFloats(), x)
	return out, fx
}