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

// Broyden's quasi-Newton maintenance for FindRootSystem, enabled by
// RootSystemOptions.UseBroyden. The central-difference Jacobian is
// built once at the start and the iteration walks on its maintained
// inverse: after every accepted step the rank-one update corrects the
// inverse so that it satisfies the secant equation H·y = s for the
// step just taken.
//
// The update maintained here is the bad Broyden form in its inverse
// shape, H + (s − H·y)·yᵀ/(yᵀy) with H the maintained inverse: a
// rank-one correction along the residual-difference direction y that
// satisfies the secant equation exactly for the step just taken, and
// turns every later iteration into a single matrix-vector product,
// δ = −H·r. The good form, whose correction runs along sᵀH instead,
// satisfies the same equation with a different matrix; maintaining it
// needs the product sᵀH·y on top of the step, which spends back the
// per-iteration saving the option exists for.

// broydenInvert fills h with the inverse of the central-difference
// Jacobian jac by solving jac·h = I column-wise through the library's
// LU solver: one factorisation, n right-hand columns, the same cost
// class as the Newton path's single solve. jacWork is consumed in
// place, exactly as the Newton path consumes it, so jac stays
// pristine for the steepest-descent fallback. A singular Jacobian is
// returned as the solver's own error for the caller to fall back on.
func broydenInvert(jac, jacWork, h [][]float64) error {
	n := len(jac)
	rhs := make([][]float64, n)
	for i := range n {
		rhs[i] = make([]float64, n)
		rhs[i][i] = 1
	}
	for i := range n {
		copy(jacWork[i], jac[i])
	}
	sol, err := base.SolveSystem("FindRootSystem", jacWork, rhs)
	if err != nil {
		return err
	}
	// sol[k] is jac⁻¹·eₖ, the kth column of the inverse.
	for i := range n {
		for k := range n {
			h[i][k] = sol[k][i]
		}
	}
	return nil
}

// broydenMaintain applies the rank-one update
//
//	H ← H + (s − H·y)·yᵀ/(yᵀy)
//
// to the maintained inverse, where s is the accepted step and y the
// residual change it produced, and reports whether the inverse is fit
// to carry forward, together with the running count of consecutive
// steps that failed to lower the residual infinity norm. A false
// report leaves h untouched and orders FindRootSystem to rebuild the
// Jacobian numerically before it steps again, the restart the option
// documents. Two observations order the restart, both a degradation of
// the rank-one model:
//
//   - the update denominator yᵀy is zero, non-finite, or at rounding
//     level against the residual's own scale (‖y‖∞ ≤ ε·max(1, ‖r‖∞)):
//     the division would amplify cancellation noise into H, and a y
//     that small carries no curvature information at all;
//   - two consecutive accepted steps each failed to lower the
//     residual infinity norm. The damping guarantees the residual sum
//     of squares falls on every accepted step, so a flat infinity
//     norm twice in a row means the maintained inverse has stopped
//     predicting the landscape and a fresh Jacobian is cheaper than
//     more crawling.
func broydenMaintain(h [][]float64, s, y, r []float64, res, resPrev float64, stalled int) (bool, int) {
	den := 0.0
	ynorm := 0.0
	for i := range y {
		den += y[i] * y[i]
		if v := math.Abs(y[i]); v > ynorm {
			ynorm = v
		}
	}
	if den == 0 || math.IsNaN(den) || math.IsInf(den, 0) || ynorm <= base.EpsF*math.Max(1, normInfOfStep(r)) {
		return false, 0
	}
	if res < resPrev {
		stalled = 0
	} else {
		stalled++
		if stalled >= 2 {
			return false, 0
		}
	}
	for i := range h {
		hy := 0.0
		for k := range y {
			hy += h[i][k] * y[k]
		}
		w := (s[i] - hy) / den
		for k := range y {
			h[i][k] += w * y[k]
		}
	}
	return true, stalled
}

// broydenStep writes the quasi-Newton step δ = −H·r into step: the
// whole per-iteration linear algebra the maintained inverse leaves,
// a single matrix-vector product.
func broydenStep(h [][]float64, r, step []float64) {
	for i := range step {
		s := 0.0
		for k := range r {
			s += h[i][k] * r[k]
		}
		step[i] = -s
	}
}