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

package optim

import (
	"fmt"
	"math"

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

// Root finding for systems of nonlinear equations r(x) = 0 with as
// many equations as unknowns. FindRoot and FindRootNewton cover the
// one-dimensional case; implicit solvers, equilibrium conditions and
// closure relations live in many dimensions, where a residual vector
// replaces the single equation.
//
// The method is the damped Newton iteration: each step solves the
// linearised system J·δ = −r with a central-difference Jacobian and
// the library's LU solver, then backtracks along δ until the residual
// norm actually falls, which keeps the iteration from leaping out of
// the basin of attraction when the linear model overshoots.
//
// RootSystemOptions.UseBroyden switches how the Jacobian is
// maintained: it is built by central differences once at the start
// and carried between steps by the rank-one Broyden update on
// its inverse, with a numerical rebuild whenever the update degrades
// (broydenMaintain documents the triggers). The default false keeps
// the per-step Jacobian and the iteration exactly as described here.

// RootSystemOptions tunes FindRootSystem. Tolerance ≤ 0 means 1e-10
// (an infinity-norm threshold on both the residual and the scaled
// step), MaxIterations ≤ 0 means 100.
type RootSystemOptions struct {
	Tolerance     float64
	MaxIterations int
	// UseBroyden maintains the Jacobian across steps instead of
	// rebuilding it every round: one central-difference Jacobian at
	// the start, then the rank-one Broyden update on its inverse
	// after every accepted step, with a numerical rebuild whenever the
	// update degrades (the triggers are documented on
	// broydenMaintain). The default false leaves the per-step
	// Jacobian and the iteration's results exactly as they are
	// without the option.
	UseBroyden bool
	// ParallelJacobian lets the central-difference Jacobian sweep its
	// columns on several goroutines, the build at the start and every
	// Broyden restart included. 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 run is bit for bit the same either
	// way, because the columns are independent and each one is
	// differenced by the same stencil.
	ParallelJacobian bool
}

// FindRootSystem solves r(x) = 0 for a vector function r of an n-vector,
// returning the solution and the residual infinity norm at it. The
// function must return a vector of the same length as x0. A singular
// Jacobian, a step that collapses before the tolerance is met or an
// exhausted iteration budget are errors, never silent answers.
//
// The iteration is a local solver: it follows the residual norm
// downhill, so a start whose basin contains no root, or one sitting
// in a parasitic minimum of the residual norm, reports the unfinished
// residual instead of pretending to converge. When the Newton system
// is singular the step falls back to the steepest-descent direction
// of the residual norm, which lets the iteration slide off the
// singular locus rather than dying there.
func FindRootSystem(f func(x *core.Array) (*core.Array, error), x0 *core.Array,
	opts RootSystemOptions) (*core.Array, float64, error) {
	if opts.Tolerance <= 0 {
		opts.Tolerance = 1e-10
	}
	if opts.MaxIterations <= 0 {
		opts.MaxIterations = 100
	}
	if x0.Dtype() == core.Complex {
		return nil, 0, base.Errf("FindRootSystem: complex unknowns are not supported")
	}
	n := x0.Len()
	if n == 0 {
		return nil, 0, base.Errf("FindRootSystem: the initial guess must not be empty")
	}
	// eval reads the residual at x. The finiteness gate is strict for
	// the states the iteration adopts (the start point and accepted
	// iterates): a non-finite residual there is invisible to the
	// norm's strict comparisons and would read as a converged root,
	// the same refusal LevenbergMarquardt makes. Backtracking trials
	// take the lenient variant: a trial that stepped into overflow is
	// a rejected candidate to halve past, not a dead run, exactly the
	// recovery the damping exists for.
	eval := func(x []float64, strict bool, dst []float64) ([]float64, error) {
		// The callback receives a copy, not an aliasing view of the
		// reused iterate or stencil buffers: arrays are contractually
		// immutable, and a callback that retains its argument must not
		// observe the later writes. Minimise, L-BFGS and Levenberg-
		// Marquardt copy for the same reason.
		xa, aerr := core.FromFloats(x, len(x))
		if aerr != nil {
			return nil, aerr
		}
		out, err := f(xa)
		if err != nil {
			return nil, err
		}
		if out.NDim() != 1 || out.Len() != n {
			return nil, base.Errf("residual has shape %s, want a vector of length %d",
				base.ShapeText(out.Shape()), n)
		}
		if err := requireReal("FindRootSystem", "residuals", out); err != nil {
			return nil, err
		}
		// dst carries a buffer the Jacobian's stencil reuses across
		// columns; the states the iteration keeps come back freshly
		// allocated. Every entry of the buffer is written before it is
		// read.
		res := dst
		if cap(res) < n {
			res = make([]float64, n)
		}
		r := res[:n]
		for i := range n {
			r[i] = out.FloatAt(i)
			if strict && (math.IsNaN(r[i]) || math.IsInf(r[i], 0)) {
				// Bare, as Minimise's eval explains: every call site
				// wraps once with the entry-point name, so the surfaced
				// message carries exactly one prefix.
				return nil, fmt.Errorf("the residual returned the non-finite value %g at coordinate %d", r[i], i)
			}
		}
		return r, nil
	}

	// cloneDense promotes through FloatAt, so Int and Float32 starting
	// vectors behave exactly like Float64 ones (a RawFloats fast path
	// would hand the iteration a nil slice for those dtypes).
	x := cloneDense(x0)
	r, err := eval(x, true, nil)
	if err != nil {
		return nil, 0, base.Errf("FindRootSystem: %w", err)
	}
	res := normInfOfStep(r)
	jac := make([][]float64, n)
	for i := range n {
		jac[i] = make([]float64, n)
	}
	// The Newton solve runs on a working copy of the Jacobian, refilled
	// from jac every round: base.Factor consumes its argument in place
	// before it can fail, and the fallback below needs the pristine
	// central-difference matrix.
	jacWork := make([][]float64, n)
	for i := range n {
		jacWork[i] = make([]float64, n)
	}
	// Buffers reused across iterations: the two perturbed stencils, the
	// backtracking candidate, the Newton column and the step. Each is
	// fully rewritten before it is read, and only transiently wrapped
	// views of them reach the callback. The two residual vectors of the
	// Jacobian's stencil are carried the same way: a column is
	// differenced from them and they are dead once it is written, so
	// one pair per round replaces one pair per column.
	xp := make([]float64, n)
	xm := make([]float64, n)
	candidate := make([]float64, n)
	col := make([]float64, n)
	colWrap := [][]float64{nil}
	step := make([]float64, n)
	var resPlus, resMinus []float64
	// Broyden state, allocated only under UseBroyden: the maintained
	// inverse of the Jacobian and the flag saying it is fit to use,
	// false until the first Jacobian is inverted and again after any
	// restart, which sends the next round down the same
	// build-and-invert path as the first.
	useBroyden := opts.UseBroyden
	var invJac [][]float64
	var sVec, yVec []float64
	haveInv := false
	stalled := 0
	if useBroyden {
		invJac = make([][]float64, n)
		for i := range n {
			invJac[i] = make([]float64, n)
		}
		sVec = make([]float64, n)
		yVec = make([]float64, n)
	}
	for range opts.MaxIterations {
		if res <= opts.Tolerance {
			return linalg.ArrayFromFloatsSafe(x, len(x)), res, nil
		}
		var derr error
		if useBroyden && haveInv {
			// The maintained inverse turns the linearised solve into a
			// matrix-vector product, δ = −H·r: no Jacobian build and
			// no factorisation this round.
			broydenStep(invJac, r, step)
		} else {
			// Central-difference Jacobian, one column per unknown. The
			// stencil carries the offset on one unknown at a time,
			// restored as soon as the column is differenced, so the
			// whole point is copied once per round rather than once per
			// column. column walks one unknown'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(x[j]))
				sp[j] += eps
				sm[j] -= eps
				rp, e1 := eval(sp, true, rp)
				rm, e2 := eval(sm, true, rm)
				sp[j], sm[j] = x[j], x[j]
				if e1 != nil || e2 != nil {
					return rp, rm, firstError(e1, e2)
				}
				for i := range n {
					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 x, so no two workers write the same
				// address and the sweep needs no locks. The build runs
				// once at the start and again at every Broyden restart,
				// both through this sweep. 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.
				colErrs := make([]error, n)
				engine.ParallelMin(n, 1, func(start, end int) {
					sp, sm := make([]float64, n), make([]float64, n)
					copy(sp, x)
					copy(sm, x)
					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 nil, 0, base.Errf("FindRootSystem: %w", err)
					}
				}
			} else {
				copy(xp, x)
				copy(xm, x)
				for j := range n {
					var err error
					resPlus, resMinus, err = column(j, xp, xm, resPlus, resMinus)
					if err != nil {
						return nil, 0, base.Errf("FindRootSystem: %w", err)
					}
				}
			}
			for i := range n {
				col[i] = -r[i]
			}
			// Factor consumes the matrix in place: the pivot search swaps
			// the rows and the elimination overwrites the subdiagonal with
			// the multipliers, and the mutation lands in base.go before
			// CheckSingular can reject the matrix. The solve therefore runs
			// on the copy, so the steepest-descent fallback reads the
			// pristine Jacobian: a direction formed from the factored
			// matrix applies the permuted LU factors to r, which is not
			// −Jᵀr and does not descend.
			for i := range n {
				copy(jacWork[i], jac[i])
			}
			colWrap[0] = col
			if useBroyden {
				// First round or restart after a degraded update:
				// rebuild the inverse and step with it at once, so a
				// fresh Jacobian is never paid for without moving. A
				// singular Jacobian takes the same steepest-descent
				// escape the Newton path uses, with the inverse left
				// unfit so the next round rebuilds.
				if ierr := broydenInvert(jac, jacWork, invJac); ierr != nil {
					derr = ierr
					steepestDescentStep(jac, r, step)
				} else {
					haveInv = true
					broydenStep(invJac, r, step)
				}
			} else {
				sol, serr := base.SolveSystem("FindRootSystem", jacWork, colWrap)
				derr = serr
				if serr != nil {
					// A singular Jacobian has no Newton direction, but the
					// steepest-descent direction of the merit function ‖r‖²
					// always exists when r ≠ 0 and always decreases it, so
					// the iteration escapes the singular locus instead of
					// dying there.
					steepestDescentStep(jac, r, step)
				} else {
					for i := range n {
						step[i] = sol[0][i]
					}
				}
			}
		}
		// Backtrack along the step until the residual norm falls, at
		// the Armijo fraction of the predicted linear decrease.
		resSq := 0.0
		for i := range n {
			resSq += r[i] * r[i]
		}
		// Scale the descent fallback down to Newton's magnitude so the
		// first trial is comparable.
		if derr != nil {
			stepNorm := normInfOfStep(step)
			if stepNorm > 0 {
				scale := normInfOfStep(col) / stepNorm
				if scale > 0 && !math.IsInf(scale, 0) {
					for i := range step {
						step[i] *= scale
					}
				}
			}
		}
		alpha := 1.0
		improved := false
		for range 60 {
			for j := range n {
				candidate[j] = x[j] + alpha*step[j]
			}
			rc, cerr := eval(candidate, false, nil)
			if cerr != nil {
				return nil, 0, base.Errf("FindRootSystem: %w", cerr)
			}
			candSq := 0.0
			for i := range n {
				candSq += rc[i] * rc[i]
			}
			// A trial whose residual overflowed is a rejected candidate:
			// its norm is Inf or NaN, compares false below, and the
			// halving continues, which is what the damping exists for.
			// A finite sum of squares admits only finite components, so
			// whatever passes both comparisons is adoptable.
			if candSq <= (1-1e-4*alpha)*resSq || candSq < opts.Tolerance*opts.Tolerance {
				// The candidate buffer is reused by the next round's
				// backtracking, so the accepted point is copied into
				// the iteration's own state.
				if useBroyden {
					// The update's step and residual differences come
					// from the state the accepted step leaves, so they
					// are read before the copy overwrites it.
					for j := range n {
						sVec[j] = candidate[j] - x[j]
						yVec[j] = rc[j] - r[j]
					}
				}
				copy(x, candidate)
				r = rc
				resNew := normInfOfStep(r)
				// The update corrects a maintained inverse, so it only
				// runs when one is fit to be corrected: after a
				// singular round the rebuild next loop takes over.
				if useBroyden && haveInv {
					// A false report leaves the inverse untouched and
					// sends the next round back to a fresh Jacobian.
					haveInv, stalled = broydenMaintain(invJac, sVec, yVec, r, resNew, res, stalled)
				}
				res = resNew
				improved = true
				break
			}
			alpha /= 2
		}
		if !improved {
			return nil, 0, base.Errf("FindRootSystem: the residual cannot be reduced below %g by damping", res)
		}
		// The documented success criterion is both a small step and a
		// residual under the tolerance: a vanished step alone can come
		// from a singular Jacobian whose steepest-descent fallback is
		// numerically zero, and that point is a stall to keep working
		// on, never a root to publish.
		if res <= opts.Tolerance && normInfOfStep(step) <= opts.Tolerance*(1+normInfOfStep(x)) {
			return linalg.ArrayFromFloatsSafe(x, len(x)), res, nil
		}
	}
	// The last accepted step updated r and res after the loop-top test,
	// so a run that met the tolerance exactly on the final iteration
	// must re-test before the budget refusal reports it.
	if res <= opts.Tolerance {
		return linalg.ArrayFromFloatsSafe(x, len(x)), res, nil
	}
	return nil, 0, base.Errf("FindRootSystem: reached MaxIterations=%d with residual %g",
		opts.MaxIterations, res)
}

// normInfOfStep returns the infinity norm of a step vector.
func normInfOfStep(step []float64) float64 {
	worst := 0.0
	for _, v := range step {
		if a := math.Abs(v); a > worst {
			worst = a
		}
	}
	return worst
}

// steepestDescentStep writes the negative gradient of the merit
// function ‖r‖², −Jᵀr, into step: entry i sums column i of the
// Jacobian against the residual, so entry j of the Jacobian's row k
// carries ∂r_k/∂x_j. jac must be the pristine central-difference
// matrix, never a factored one: Factor overwrites the subdiagonal
// with the multipliers and swaps the rows in place, so a factored
// matrix yields a direction that is no descent direction at all.
func steepestDescentStep(jac [][]float64, r, step []float64) {
	for i := range step {
		s := 0.0
		for k := range r {
			s += jac[k][i] * r[k]
		}
		step[i] = -s
	}
}

// firstError returns the first non-nil of two errors.
func firstError(e1, e2 error) error {
	if e1 != nil {
		return e1
	}
	return e2
}