// 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 stiff workhorse beside IntegrateBackwardEuler: the variable-step
// second-order backward differentiation formula. Where the explicit
// Dormand-Prince pair must keep h·λ inside its stability region, BDF2
// is A-stable and damps the stiff mode like (2hλ)^(−1/2), so the step
// size follows accuracy alone. Each step solves the implicit relation
// α·z − h·f(t_{n+1}, z) = β by Newton over a numerical Jacobian and
// the library's LU solver, the machinery IntegrateBackwardEuler
// already carries.
//
// The step size is driven by a Milne-type estimate of the one-step
// error: the gap between the corrector and the quadratic predictor
// through the three most recent states, scaled by the constant that
// turns that gap into the BDF2 truncation error (2/11 for equal
// steps). The first step runs backward Euler, whose size a
// Hairer-Nørsett-Wanner style probe picks so the starter's own error
// already sits below the tolerance; the second BDF2 step repeats that
// size untested, safe because its truncation error is an order in h
// below the starter's; from the third step on the estimate controls
// everything.

// IntegrateBDF2 integrates y' = f(t, y) from t0 to t1 with the
// variable-step BDF2 scheme 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 IntegrateBDF2(f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y0 *core.Array, opts ODEOptions) (*core.Array, error) {
	const name = "IntegrateBDF2"
	y, err := odeCheck("IntegrateBDF2", y0, &opts)
	if err != nil {
		return nil, err
	}
	n := len(y)
	h, err := bdf2InitialStep(name, f, t0, t1, y, &opts)
	if err != nil {
		return nil, err
	}
	t := t0
	yn := cloneDenseSlice(y)
	var yNm1, yNm2 []float64
	var tNm1, tNm2 float64
	budget := odeBudget{max: opts.MaxSteps}
	// The implicit relation's right side and the Newton seed live in
	// reused buffers: both are fully rewritten at the top of every step
	// and neither outlives the step's solve. The converged state lands
	// in a four-buffer ring: at every acceptance the live history is
	// the three most recent ring slots, so the next slot aliases
	// nothing the step reads, and a rejected or stalled attempt
	// reuses the slot it already holds.
	beta := make([]float64, n)
	seed := make([]float64, n)
	var ring [4][]float64
	next := 0
	w := &odeWork{}
	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
		// The implicit equation and the Newton seed. The very first
		// step has no history and runs backward Euler, seeded with
		// the semi-implicit prediction; from the second step on the
		// BDF2 weights carry the 1/h scaling themselves, so the
		// derivative enters the implicit equation with weight 1.
		var alpha, weight float64
		estimated := yNm2 != nil
		if yNm1 == nil {
			alpha, weight = 1, hN
			copy(beta, yn)
			fy, ferr := odeEval(name, f, tNext, yn, n, &w.views)
			if ferr != nil {
				return nil, ferr
			}
			for i := range n {
				seed[i] = yn[i] + hN*fy[i]
			}
		} else {
			weight = 1
			alpha = bdf2Coefficients(t, tNext, tNm1, tNm2, yn, yNm1, yNm2, beta, seed)
		}
		dst := ring[next]
		if dst == nil {
			dst = make([]float64, n)
			ring[next] = dst
		}
		if nerr := odeNewton(name, f, w, tNext, alpha, weight, beta, seed, dst, opts.AbsTol, opts.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.
			c := bdf2Milne(t, tNext, tNm1, tNm2)
			errNorm := 0.0
			for i := range n {
				scale := opts.AbsTol + opts.RelTol*math.Max(math.Abs(yn[i]), math.Abs(dst[i]))
				ratio := c * (dst[i] - seed[i]) / scale
				errNorm += ratio * ratio
			}
			errNorm = math.Sqrt(errNorm/float64(n)) + 1e-10
			if errNorm <= 1 {
				factor = math.Min(2, math.Max(0.2, 0.9*math.Pow(1/errNorm, 1.0/3)))
			} else {
				// Rejected: retry the same interval with a smaller step.
				h *= math.Max(0.1, math.Min(1, 0.9*math.Pow(1/errNorm, 1.0/3)))
				continue
			}
		}
		// Accepted: shift the history one step forward. The ring slot
		// becomes the working state and the buffers flow without
		// copying; nothing aliases them afterwards.
		yNm2, tNm2 = yNm1, tNm1
		yNm1, tNm1 = yn, t
		yn = dst
		next = (next + 1) % len(ring)
		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)
		}
	}
	return arrayFromVector(yn), nil
}

// bdf2InitialStep picks the first step by probing f: a trial step h0
// compares the derivative at y against the derivative one h0 further,
// and the result sizes the step so a first-order scheme's local error
// sits a factor hundred below the mixed tolerance. The span and a
// hundredfold h0 bound the answer, and the sign carries the
// integration direction.
func bdf2InitialStep(name string, f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y []float64, opts *ODEOptions) (float64, error) {
	span := math.Abs(t1 - t0)
	if span == 0 {
		return 0, nil
	}
	n := len(y)
	f0, err := odeEval(name, f, t0, y, n, nil)
	if err != nil {
		return 0, err
	}
	scale := make([]float64, n)
	d0, d1 := 0.0, 0.0
	for i := range n {
		scale[i] = opts.AbsTol + opts.RelTol*math.Abs(y[i])
		d0 = math.Max(d0, math.Abs(y[i])/scale[i])
		d1 = math.Max(d1, math.Abs(f0[i])/scale[i])
	}
	h0 := 1e-6
	if d0 > 1e-5 && d1 > 1e-5 {
		h0 = 0.01 * d0 / d1
	}
	h0 = math.Min(h0, span)
	// The probe steps in the integration direction: a backward span
	// samples t0−h0 with the derivative subtracted, or the difference
	// f1−f0 measures the wrong side of the dynamics.
	dir := 1.0
	if t1 < t0 {
		dir = -1
	}
	probe := make([]float64, n)
	for i := range n {
		probe[i] = y[i] + dir*h0*f0[i]
	}
	f1, err := odeEval(name, f, t0+dir*h0, probe, n, nil)
	if err != nil {
		return 0, err
	}
	d2 := 0.0
	for i := range n {
		d2 = math.Max(d2, math.Abs(f1[i]-f0[i])/(scale[i]*h0))
	}
	h1 := span
	if d := math.Max(d1, d2); d > 1e-15 {
		h1 = math.Sqrt(0.01 / d)
	}
	h1 = math.Min(h1, math.Min(100*h0, span))
	if t1 < t0 {
		h1 = -h1
	}
	return h1, nil
}

// bdf2Coefficients assembles the variable-step BDF2 relation for a
// step from t, whose previous point sits at tNm1 (and the one before
// that at tNm2 when known), to tNext. It writes α's companions β and
// the Newton seed into the caller's buffers and returns α: the
// implicit equation is α·z − h·f(tNext, z) = β. Both buffers are fully
// overwritten. The seed is the quadratic predictor through the last
// three states when yNm2 is given, otherwise the linear ramp over the
// last two. All differences are signed, so backward integration needs
// no separate path.
func bdf2Coefficients(t, tNext, tNm1, tNm2 float64,
	yn, yNm1, yNm2, beta, seed []float64) float64 {
	n := len(yn)
	hN := tNext - t
	hP := t - tNm1
	alpha := (hP + 2*hN) / ((hP + hN) * hN)
	w1 := (hP + hN) / (hP * hN)
	w0 := hN / (hP * (hP + hN))
	for i := range n {
		beta[i] = w1*yn[i] - w0*yNm1[i]
	}
	if yNm2 != nil {
		l2 := (tNext - tNm1) * (tNext - t) / ((tNm2 - tNm1) * (tNm2 - t))
		l1 := (tNext - tNm2) * (tNext - t) / ((tNm1 - tNm2) * (tNm1 - t))
		l0 := (tNext - tNm2) * (tNext - tNm1) / ((t - tNm2) * (t - tNm1))
		for i := range n {
			seed[i] = l2*yNm2[i] + l1*yNm1[i] + l0*yn[i]
		}
	} else {
		ramp := hN / hP
		for i := range n {
			seed[i] = yn[i] + ramp*(yn[i]-yNm1[i])
		}
	}
	return alpha
}

// bdf2Milne returns the constant that turns the gap between the BDF2
// corrector and the quadratic predictor through the three previous
// states into an estimate of the corrector's one-step error: 2/11 for
// equal steps, from the leading error terms h³y”' of the predictor
// and (2/9)h³y”' of the corrector.
func bdf2Milne(t, tNext, tNm1, tNm2 float64) float64 {
	hN := tNext - t
	hP := t - tNm1
	hPp := tNm1 - tNm2
	return hN * (hN + hP) / ((hP+2*hN)*(hPp+hP+hN) + hN*(hN+hP))
}