// 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 Rosenbrock-Wanner workhorse: the four-stage L-stable scheme ROS4
// of Hairer and Wanner, fourth order accurate with an embedded
// third-order solution driving the step control. Where BDF solves each
// step by Newton, a W method puts the Jacobian into the formula
// itself: every stage is one linear solve against the frozen matrix
// (1/(γh))I − J, so a step costs one numerical Jacobian, one LU
// factorisation and four back-substitutions, and no Newton iteration
// ever stalls. The scheme is L-stable: its stability function vanishes
// at infinity, so a step far beyond the transient's time constant
// damps the stiff mode instead of amplifying it.
//
// The stage form is the divided one the standard implementations use.
// With A = (1/(γh))I − J factored once per step, stage one solves
// A·k₁ = f(t, y) and stage i solves A·kᵢ = f(t + cᵢh, yᵢ) +
// (Σⱼ<ᵢ cᵢⱼkⱼ)/h with the stage value yᵢ = y + Σⱼ<ᵢ aᵢⱼkⱼ; the step
// advances by y + Σ bᵢkᵢ and the embedded estimate is Σ êᵢkᵢ. Stage
// four shares its node and its stage value with stage three, so its f
// evaluation carries over and a step costs three f calls. The digits
// are the published ones, cross-checked against the standard
// implementations of the scheme.
//
// Two honest limits of the tableau. The Jacobian enters at (t, y) and
// the tableau's time-derivative weights are left out, because the
// package's f contract carries no partial derivative in t: the fourth
// order therefore holds for autonomous systems (every problem in this
// package's stiff tests is one), while a genuinely time-dependent f
// loses the second-order local terms those weights carry and can
// degrade to first order globally; measured on y' = −y + t with
// uniform steps the global ratios come out near 2, not near 16. And
// the scheme is L-stable but not stiffly accurate: the fourth stage's
// weights differ from the solution weights, the stiff damping coming
// from the stability limit rather than from stage-end agreement.

var (
	// rowGamma is the abscissa of the stage solves and the source of
	// the L-stability: the stability function's denominator is
	// (1 − γz)⁴ and its numerator vanishes at infinity.
	rowGamma = 0.57282

	// rowNodes are the stage times as multiples of h. The second node
	// is the published tableau's c2 = 0.114564 (Hairer and Wanner's
	// ROS4); its doubled form 2·gamma appears in some reprints, and the
	// digit never enters the arithmetic of an autonomous problem, where
	// stage times cancel, so either form integrates identically here.
	rowNodes = [4]float64{0, 0.114564, 0.65521686381559, 0.65521686381559}

	// rowA holds the stage-value weights a[i][j], the coefficient of
	// k_j inside stage i's value.
	rowA = [4][3]float64{
		{},
		{2},
		{1.867943637803922, 0.2344449711399156},
		{1.867943637803922, 0.2344449711399156, 0},
	}

	// rowC holds the Jacobian-coupling weights c[i][j], the coefficient
	// of k_j inside stage i's right side, divided by h.
	rowC = [4][3]float64{
		{},
		{-7.13761503641231},
		{2.580708087951457, 0.6515950076447975},
		{-2.137148994382534, -0.3214669691237626, -0.6949742501781779},
	}

	// rowB advances the solution; rowE carries the embedded third-order
	// estimate (the difference between the fourth-order and the
	// embedded third-order weights).
	rowB = [4]float64{2.255570073418735, 0.2870493262186792, 0.435317943184018, 1.093502252409163}
	rowE = [4]float64{-0.2815431932141155, -0.0727619912493892, -0.1082196201495311, -1.093502252409163}
)

// IntegrateROS4 integrates y' = f(t, y) from t0 to t1 with the
// four-stage L-stable Rosenbrock-Wanner scheme ROS4 and returns y(t1).
// The fourth order holds for autonomous systems; a genuinely
// time-dependent f loses the second-order local terms the tableau's
// time-derivative weights would carry (the package's f contract has no
// partial derivative in t), and can degrade to first order globally:
// measured on y' = −y + t with uniform steps the ratios come out near
// 2, not near 16. The adaptive controller still holds its tolerance
// there, at the cost of more steps. The step size follows the embedded
// error under the classic accept-or-shrink control, with the numerical
// Jacobian taken once per step through the same central-difference
// helper the Newton path uses; there is no user Jacobian parameter.
// 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 stage that leaves the
// finite range even as the step shrinks is an error, never a silently
// truncated trajectory.
func IntegrateROS4(f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y0 *core.Array, opts ODEOptions) (*core.Array, error) {
	const name = "IntegrateROS4"
	y, err := odeCheck(name, y0, &opts)
	if err != nil {
		return nil, err
	}
	h, err := bdf2InitialStep(name, f, t0, t1, y, &opts)
	if err != nil {
		return nil, err
	}
	budget := odeBudget{max: opts.MaxSteps}
	w := &odeWork{}
	t := t0
	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)
		yEnd, errNorm, nerr := ros4Step(name, f, w, t, y, h, opts.AbsTol, opts.RelTol)
		if nerr != nil {
			if errors.Is(nerr, errNewtonStalled) {
				// The stage matrix closed onto singularity: halve the
				// step and retry the same interval, within the budget.
				h *= 0.5
				continue
			}
			return nil, nerr
		}
		if math.IsNaN(errNorm) || math.IsInf(errNorm, 0) {
			// A stage left the finite range: shrink hard and retry.
			h *= 0.25
			continue
		}
		factor := min(5, max(0.2, 0.9*math.Pow(1/errNorm, 1.0/4)))
		if errNorm <= 1 {
			copy(y, yEnd)
			prevT := t
			t += h
			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)
			}
		} else {
			// Rejected: retry the same interval with the smaller step.
			h *= max(0.2, factor)
		}
	}
	return arrayFromVector(y), nil
}

// ros4Step attempts one ROS4 step of the given size from (t, y) and
// returns the candidate end state, the embedded error norm against the
// mixed absolute and relative tolerance, and a nil error when the
// stages all solved. Every buffer belongs to the workspace, so a step
// allocates nothing of its own: the caller's y is left untouched, and
// the returned slice stays valid until the next step.
func ros4Step(name string, f func(t float64, y *core.Array) (*core.Array, error),
	w *odeWork, t float64, y []float64, h, absTol, relTol float64) ([]float64, float64, error) {
	n := len(y)
	w.use(n)
	// One scratch buffer per role: the stage value, the stage right
	// side, f's result and the candidate end state are rebuilt at the
	// top of their use and read only by it.
	ks, stage, rhs, fy, yEnd := w.ks, w.stage, w.rhs, w.fy, w.yEnd
	// One numerical Jacobian per step, the house central-difference
	// helper the Newton path shares, and one factorisation of I − hγJ,
	// which is (1/(γh))I − J scaled by γh: the scale folds into the
	// stage right sides instead of the matrix.
	jac, jerr := odeJacobian(name, f, t, y, w)
	if jerr != nil {
		return nil, 0, jerr
	}
	lu := w.mat
	// The step-scaled matrix weight and the stage weight are one
	// product each, the same (−h·γ) and (h·γ) groupings the element
	// loops evaluated.
	hgr := -h * rowGamma
	hg := h * rowGamma
	for i := range n {
		row := lu[i]
		for j := range n {
			row[j] = hgr * jac[i*n+j]
		}
		row[i]++
	}
	w.perm, _ = base.Factor(lu)
	if err := base.CheckSingular(name, lu); err != nil {
		return nil, 0, base.Errf("%s: %w, singular stage matrix at t=%g", name, errNewtonStalled, t)
	}
	// Stage one.
	out, ferr := odeCall(name, f, t, y, n, &w.views)
	if ferr != nil {
		return nil, 0, ferr
	}
	readVector(fy, out)
	for i := range n {
		rhs[i] = hg * fy[i]
	}
	odePermuteColumn(rhs, w.perm, w.visited)
	base.SolveColumn(lu, rhs)
	copy(ks[0], rhs)
	// Stages two through four. Stage four repeats stage three's node
	// and stage value, so its derivative carries over.
	for s := 1; s < 4; s++ {
		if s < 3 {
			copy(stage, y)
			for j := range s {
				aj := rowA[s][j]
				if aj == 0 {
					continue
				}
				kjs := ks[j]
				for i := range n {
					stage[i] += aj * kjs[i]
				}
			}
			out, ferr = odeCall(name, f, t+rowNodes[s]*h, stage, n, &w.views)
			if ferr != nil {
				return nil, 0, ferr
			}
			readVector(fy, out)
		}
		for i := range n {
			rhs[i] = hg * fy[i]
		}
		for j := range s {
			cj := rowC[s][j]
			if cj == 0 {
				continue
			}
			gcj := rowGamma * cj
			kjs := ks[j]
			for i := range n {
				rhs[i] += gcj * kjs[i]
			}
		}
		odePermuteColumn(rhs, w.perm, w.visited)
		base.SolveColumn(lu, rhs)
		copy(ks[s], rhs)
	}
	// The embedded pair: the b weights advance, the gap to the
	// embedded third-order solution estimates the local error.
	errNorm := 0.0
	for i := range n {
		e, advance := 0.0, 0.0
		for s := range 4 {
			e += rowE[s] * ks[s][i]
			advance += rowB[s] * ks[s][i]
		}
		yEnd[i] = y[i] + advance
		scale := absTol + relTol*math.Max(math.Abs(y[i]), math.Abs(yEnd[i]))
		ratio := e / scale
		errNorm += ratio * ratio
	}
	return yEnd, math.Sqrt(errNorm/float64(n)) + 1e-10, nil
}