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

package integrate

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

import (
	"errors"
	"math"
)

// The variable-order stiff workhorse above IntegrateBDF2: the backward
// differentiation formula of order one through five, with the step size
// and the order both adapted every step in the VODE manner. Each step
// interpolates a polynomial of degree k through the k most recent
// states and the unknown end value and requires its derivative at the
// new time to equal f, the same implicit relation BDF2 solves; the
// Newton iteration, the LU machinery, the Hairer-Nørsett-Wanner initial
// step probe and the step controller are the ones IntegrateBDF2
// already carries.
//
// The coefficients are the variable-step, divided-difference form: the
// Newton form of the interpolating polynomial through (tNext, z) and
// the stored back values, written per component from a small divided-
// difference table over the stored times. The form was chosen over the
// fixed-coefficient one because the package keeps a solution history
// rather than a Nordsieck array, because the divided differences feed
// the order selection (the a-priori error estimate per candidate order
// falls out of the same table) and because the relation leaves the
// Newton contract α·z − h·f(tNext, z) = β of odeNewton untouched. At
// order two with equal steps the assembled α and β agree with
// bdf2Coefficients to rounding, so the shipped BDF2 behaviour is the
// special case the driver degrades to.
//
// The local error estimate is the Milne-type one: the gap between the
// corrector and the degree-k predictor extrapolated from the k+1
// newest states, scaled by the constant that turns the gap into the
// corrector's own error. The variable-step constant generalises the
// 2/11 of bdf2Milne: with α the derivative weight of the new point and
// S the span from tNext to the oldest predictor node, the estimate is
// (z − seed)/(1 + α·S), which for equal steps of order two reproduces
// 2/11 exactly. The order itself is chosen before the solve, from the
// divided differences of the stored states: the (k+1)-th divided
// difference approximates y^(k+1)/(k+1)!, and the candidate whose
// implied optimal step is largest wins, with a margin so the order
// does not flicker between neighbours.
//
// The first step is backward Euler, sized by the shared probe; the
// order ramps up as the history accumulates, one level per step.

// BDFVarStats reports what a variable-order run did: the accepted and
// rejected steps and the highest order the driver reached.
type BDFVarStats struct {
	Steps    int
	Rejected int
	MaxOrder int
}

// BDFVarOptions tunes IntegrateBDFVar. RelTol ≤ 0 means 1e-6, AbsTol ≤ 0
// means 1e-9, MaxSteps ≤ 0 means 100000, the ODEOptions defaults. Stats,
// when not nil, receives the run's counters.
type BDFVarOptions struct {
	RelTol   float64
	AbsTol   float64
	MaxSteps int
	Stats    *BDFVarStats
}

// bdfVarOrderMax is the highest order the driver raises to. bdfVarKeep
// is the number of states held back: order k needs k back values for
// its corrector, k+1 for its predictor and k+2 for the a-priori order
// comparison, so seven states serve order five in every role.
const (
	bdfVarOrderMax = 5
	bdfVarKeep     = bdfVarOrderMax + 2
)

// IntegrateBDFVar integrates y' = f(t, y) from t0 to t1 with the
// variable-step, variable-order BDF scheme of orders one through five
// and returns y(t1). Backward integration works: a t1 < t0 simply
// integrates in the negative direction. An exhausted step budget, a
// collapsed step size, an f that returns a wrongly shaped state, or a
// Newton iteration that cannot converge even as the step shrinks is an
// error, never a silently truncated trajectory.
func IntegrateBDFVar(f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y0 *core.Array, opts BDFVarOptions) (*core.Array, error) {
	end, err := integrateBDFVar("IntegrateBDFVar", f, t0, t1, y0, opts, bdfVarOrderMax, false)
	if err != nil {
		return nil, err
	}
	return arrayFromVector(end), nil
}

// integrateBDFVar drives the variable-order loop. maxOrder caps the
// order adaptation and lockOrder pins the order at maxOrder once the
// history ramp reaches it, which is the fixed-order hook the tests
// drive; the public entry always asks for adaptive order five.
func integrateBDFVar(name string, f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y0 *core.Array, opts BDFVarOptions, maxOrder int, lockOrder bool) ([]float64, error) {
	if maxOrder < 1 || maxOrder > bdfVarOrderMax {
		return nil, base.Errf("%s: maxOrder must be between 1 and %d, got %d", name, bdfVarOrderMax, maxOrder)
	}
	y, err := odeCheck(name, y0, nil)
	if err != nil {
		return nil, err
	}
	relTol, absTol, maxSteps := opts.RelTol, opts.AbsTol, opts.MaxSteps
	if relTol <= 0 {
		relTol = 1e-6
	}
	if absTol <= 0 {
		absTol = 1e-9
	}
	if maxSteps <= 0 {
		maxSteps = 100000
	}
	n := len(y)
	h, err := bdf2InitialStep(name, f, t0, t1, y, &ODEOptions{RelTol: relTol, AbsTol: absTol})
	if err != nil {
		return nil, err
	}
	var stats BDFVarStats
	hist := &bdfVarHistory{}
	hist.push(t0, y)
	// The implicit relation's right side, the predictor seed and the
	// divided-difference workspace live in reused buffers: all are fully
	// rewritten at the top of every step. The Newton result lands in a
	// per-solve scratch buffer that never touches the history window:
	// the order selection and the coefficient assembly reread the whole
	// window on a retry, so a rejected attempt must leave every stored
	// state intact. Only an accepted step copies the state into the
	// ring slot its push then occupies.
	beta := make([]float64, n)
	seed := make([]float64, n)
	zbuf := make([]float64, n)
	dd := make([]float64, bdfVarKeep)
	nodes := make([]float64, bdfVarKeep)
	spans := make([]float64, bdfVarKeep)
	ddTab := make([][]float64, bdfVarKeep)
	for level := range ddTab {
		ddTab[level] = make([]float64, n)
	}
	budget := odeBudget{max: maxSteps}
	w := &odeWork{}
	t := t0
	carried := 1
	for !odeArrived(t, t1) {
		if err := budget.spend(name, t, t1); err != nil {
			return nil, err
		}
		// Never step past t1; t1−t carries the integration direction.
		h = odeClampStep(h, t, t1)
		tNext := t + h
		hN := tNext - t
		var alpha, weight float64
		var order int
		estimated := false
		if hist.n == 1 {
			// The very first step has no history and runs backward
			// Euler, seeded with the semi-implicit prediction: the
			// house starter IntegrateBDF2 begins with.
			alpha, weight, order = 1, hN, 1
			copy(beta, y)
			fy, ferr := odeEval(name, f, tNext, y, n, &w.views)
			if ferr != nil {
				return nil, ferr
			}
			for i := range n {
				seed[i] = y[i] + hN*fy[i]
			}
		} else {
			order = min(carried, maxOrder, hist.n-1)
			switch {
			case lockOrder && hist.n > maxOrder:
				// The fixed-order contract: once the history ramp can
				// feed the requested order, every step runs at it.
				order = maxOrder
			case !lockOrder && hist.n >= 3:
				bdfVarDividedDifferences(hist, n, ddTab, dd, nodes)
				order = bdfVarPickOrder(order, maxOrder, hist.n, tNext, h, hist, y, absTol, relTol, ddTab, spans)
			}
			weight = 1
			alpha = bdfVarCoefficients(order, tNext, hist, beta, seed, dd, nodes, spans)
			estimated = true
		}
		if nerr := odeNewton(name, f, w, tNext, alpha, weight, beta, seed, zbuf, absTol, relTol); nerr != nil {
			if errors.Is(nerr, errNewtonStalled) {
				// The implicit solve struggled: halve the step and
				// retry the same interval, within the step budget.
				h *= 0.5
				continue
			}
			return nil, nerr
		}
		factor := 1.0
		if estimated {
			// Milne-type local error estimate against the mixed
			// absolute and relative tolerance.
			_, tOldest := hist.back(order)
			c := 1 / (1 + alpha*(tNext-tOldest))
			errNorm := 0.0
			for i := range n {
				scale := absTol + relTol*math.Max(math.Abs(y[i]), math.Abs(zbuf[i]))
				ratio := c * (zbuf[i] - seed[i]) / scale
				errNorm += ratio * ratio
			}
			errNorm = math.Sqrt(errNorm/float64(n)) + 1e-10
			if errNorm <= 1 {
				factor = min(2, max(0.2, 0.9*math.Pow(1/errNorm, 1/float64(order+1))))
			} else {
				// Rejected: retry the same interval with a smaller step.
				stats.Rejected++
				h *= max(0.1, min(1, 0.9*math.Pow(1/errNorm, 1/float64(order+1))))
				continue
			}
		}
		// Accepted: the corrector is copied into the ring slot the push
		// fills and becomes the working state, so back(0) is always
		// (t, y) and the buffers flow without copying.
		slot := hist.y[hist.next]
		if slot == nil {
			slot = make([]float64, n)
		}
		copy(slot, zbuf)
		hist.push(tNext, slot)
		y = slot
		carried = order
		if order > stats.MaxOrder {
			stats.MaxOrder = order
		}
		stats.Steps++
		prevT := t
		t = tNext
		h *= factor
		// Collapse is "t did not move": a span below the absolute time
		// scale is integrable, and an accepted step that arrives at the
		// end exactly is not a failure either.
		if t == prevT {
			return nil, base.Errf("%s: the step size shrank below the resolution of t at t=%g", name, prevT)
		}
	}
	if opts.Stats != nil {
		*opts.Stats = stats
	}
	return y, nil
}

// bdfVarHistory holds the last bdfVarKeep accepted states with their
// times in a fixed ring. back(0) is the newest state, back(1) the one
// before it, and so on; slots are recycled only once they are too old
// to serve any order, so the buffers flow without copying.
type bdfVarHistory struct {
	y    [bdfVarKeep][]float64
	t    [bdfVarKeep]float64
	next int
	n    int
}

// push records an accepted state and its time as the new newest entry.
func (h *bdfVarHistory) push(t float64, y []float64) {
	h.y[h.next], h.t[h.next] = y, t
	h.next = (h.next + 1) % bdfVarKeep
	if h.n < bdfVarKeep {
		h.n++
	}
}

// back returns the state i steps behind the newest one.
func (h *bdfVarHistory) back(i int) ([]float64, float64) {
	j := (h.next - 1 - i + bdfVarKeep) % bdfVarKeep
	return h.y[j], h.t[j]
}

// bdfVarCoefficients assembles the variable-step BDF relation of the
// given order for a step to tNext from the newest history state. It
// writes α's companions β and the Newton seed into the caller's
// buffers and returns α: the implicit equation is α·z − f(tNext, z) =
// β, the weight already scaled out. Both buffers are fully overwritten.
// The seed is the degree-order polynomial through the order+1 newest
// states evaluated at tNext, the predictor the error estimate reads.
// All differences are signed, so backward integration needs no
// separate path.
//
// The construction is the divided-difference (Newton) form: with nodes
// x_0 = tNext and x_q = the q-th back time, the interpolating
// polynomial's derivative at tNext is Σ_j c_j·Π_j where c_j are the
// divided differences of the data (z at x_0, the back values after)
// and Π_j the Newton basis products. Splitting c_j into its z part,
// 1/Π_j, and its history part gives α = Σ 1/(tNext − x_m), the Lagrange
// derivative weight of the new point, and β from the history-only
// table, all from one per-component recursion.
func bdfVarCoefficients(order int, tNext float64, hist *bdfVarHistory,
	beta, seed, dd, nodes, spans []float64) float64 {
	nodes[0] = tNext
	// The ring's nodes and value slices are the same for every
	// component: gather both once, on the stack, instead of walking the
	// ring inside the per-element loop.
	var backVals [bdfVarKeep][]float64
	for q := range order + 1 {
		backVals[q], nodes[q+1] = hist.back(q)
	}
	// spans[m] is Π_m, the product of tNext − x_q over q < m: the
	// Newton basis value the level-m coefficients multiply.
	spans[0] = 1
	for m := 1; m <= order; m++ {
		spans[m] = spans[m-1] * (tNext - nodes[m])
	}
	alpha := 0.0
	for m := 1; m <= order; m++ {
		alpha += 1 / (tNext - nodes[m])
	}
	for i := range beta {
		// dd[q] starts as the value at node q: zero at tNext, the back
		// values after. One level of the recursion per Newton term;
		// level order leaves dd[0] holding the order-th divided
		// difference over the new point and dd[1] the one over the
		// stored values, which is the predictor's top coefficient.
		dd[0] = 0
		for q := range order + 1 {
			dd[q+1] = backVals[q][i]
		}
		seed[i] = dd[1]
		betaSum := 0.0
		for level := 1; level <= order; level++ {
			for q := range order + 2 - level {
				dd[q] = (dd[q+1] - dd[q]) / (nodes[q+level] - nodes[q])
			}
			betaSum += dd[0] * spans[level-1]
			seed[i] += dd[1] * spans[level]
		}
		beta[i] = -betaSum
	}
	return alpha
}

// bdfVarDividedDifferences fills tab with the divided differences of
// the stored back values alone: tab[level][i] is the level-th divided
// difference of (y_n, y_{n-1}, …) over their times for component i.
// The (order+1)-th entry approximates y^(order+1)/(order+1)! and is
// what the a-priori order comparison reads.
func bdfVarDividedDifferences(hist *bdfVarHistory, n int, tab [][]float64, dd, times []float64) {
	// The ring's times and value slices do not depend on the component:
	// gather both once, on the stack, instead of walking the ring
	// inside the per-element loops.
	var backVals [bdfVarKeep][]float64
	for q := range hist.n {
		backVals[q], times[q] = hist.back(q)
	}
	for i := range n {
		for q := range hist.n {
			dd[q] = backVals[q][i]
		}
		for level := 1; level < hist.n; level++ {
			for q := range hist.n - level {
				dd[q] = (dd[q+1] - dd[q]) / (times[q+level] - times[q])
			}
			tab[level][i] = dd[0]
		}
	}
}

// bdfVarPickOrder returns the order for the coming step. Every order
// the history supports gets an a-priori optimal step: the local error
// the divided differences predict, raised to the power that would
// bring it to the tolerance. The scan runs from order 1 upward and a
// candidate must beat the running best by a clear margin, so the
// effective pick is the lowest order within 15 percent of the largest
// predicted step: short histories and cheap coefficients win near
// ties, and the order settles instead of flickering between equals.
// The carried order survives the scan only before a second state is
// held; after that some candidate always displaces it.
func bdfVarPickOrder(carried, maxOrder, held int, tNext, h float64, hist *bdfVarHistory, y []float64,
	absTol, relTol float64, tab [][]float64, spans []float64) int {
	best, bestH := carried, 0.0
	for j := 1; j <= min(maxOrder, held-1); j++ {
		hj := math.Abs(h)
		if j <= held-2 {
			e := bdfVarPriorNorm(j, tNext, hist, y, absTol, relTol, tab, spans)
			hj = math.Abs(h) * math.Pow(1/e, 1/float64(j+1))
		}
		if hj > bestH*1.15 {
			best, bestH = j, hj
		}
	}
	return best
}

// bdfVarPriorNorm estimates the RMS error norm a step of size h at the
// given order would produce: the (order+1)-th divided difference of
// the stored states approximates y^(order+1)/(order+1)!, and the
// order's local error scales that by the Newton basis product over α,
// the same estimate the Milne constant formalises a posteriori.
func bdfVarPriorNorm(order int, tNext float64, hist *bdfVarHistory, y []float64,
	absTol, relTol float64, tab [][]float64, spans []float64) float64 {
	spans[0] = 1
	alpha := 0.0
	for q := range order {
		_, tq := hist.back(q)
		spans[q+1] = spans[q] * math.Abs(tNext-tq)
		alpha += 1 / math.Abs(tNext-tq)
	}
	w := spans[order] / alpha
	norm := 0.0
	for i := range y {
		scale := absTol + relTol*math.Abs(y[i])
		ratio := math.Abs(tab[order+1][i]) * w / scale
		norm += ratio * ratio
	}
	return math.Sqrt(norm/float64(len(y))) + 1e-10
}