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

// Ordinary differential equation solvers for initial value problems
// y' = f(t, y). The state y is a rank-1 vector of length n; a system
// of higher rank flattens to its leading-axis vector first.
//
// Three schemes cover the standard regimes. IntegrateODE is the
// workhorse: an adaptive embedded Runge-Kutta pair (Dormand-Prince
// 4(5)) that controls the local error against a mixed absolute and
// relative tolerance. IntegrateRK4 is the classical fixed-step
// fourth-order scheme, useful when a uniform step or simple
// reproducibility per step matters. IntegrateBackwardEuler is the
// entry-level stiff scheme: fully implicit, with each step's
// nonlinear equation solved by Newton over a numerical Jacobian and
// the library's LU solver.

// ODEOptions tunes the adaptive integrator. RelTol ≤ 0 means 1e-6,
// AbsTol ≤ 0 means 1e-9, MaxSteps ≤ 0 means 100000.
type ODEOptions struct {
	RelTol   float64
	AbsTol   float64
	MaxSteps int
}

// Dormand-Prince 4(5): node offsets, stage coefficients, and the
// 5th- and 4th-order solution weights. Stage 7 shares the 5th-order
// weights (the FSAL property), which is why it needs no separate row.
var (
	odeC = [7]float64{0, 1.0 / 5, 3.0 / 10, 4.0 / 5, 8.0 / 9, 1, 1}
	odeA = [][]float64{
		{},
		{1.0 / 5},
		{3.0 / 40, 9.0 / 40},
		{44.0 / 45, -56.0 / 15, 32.0 / 9},
		{19372.0 / 6561, -25360.0 / 2187, 64448.0 / 6561, -212.0 / 729},
		{9017.0 / 3168, -355.0 / 33, 46732.0 / 5247, 49.0 / 176, -5103.0 / 18656},
		{35.0 / 384, 0, 500.0 / 1113, 125.0 / 192, -2187.0 / 6784, 11.0 / 84},
	}
	odeB5 = [7]float64{35.0 / 384, 0, 500.0 / 1113, 125.0 / 192, -2187.0 / 6784, 11.0 / 84, 0}
	odeB4 = [7]float64{5179.0 / 57600, 0, 7571.0 / 16695, 393.0 / 640, -92097.0 / 339200, 187.0 / 2100, 1.0 / 40}
)

// IntegrateODE integrates y' = f(t, y) from t0 to t1 with the adaptive
// Dormand-Prince 4(5) pair and returns y(t1). Backward integration
// works: a t1 < t0 simply integrates in the negative direction. An
// exhausted step budget, a collapsed step size or an f that returns a
// wrongly shaped state is an error, never a silently truncated
// trajectory.
func IntegrateODE(f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y0 *core.Array, opts ODEOptions) (*core.Array, error) {
	return odeRun(f, t0, t1, y0, opts, nil)
}

// readVector copies a's elements into dst, sweeping the raw float64
// payload when a is a dense float64 array and falling back to the
// widening accessor for views and other dtypes. The values written are
// identical either way.
func readVector(dst []float64, a *core.Array) {
	if !a.Strided() && a.Dtype() == core.Float {
		copy(dst, a.RawFloats())
		return
	}
	for i := range dst {
		dst[i] = a.FloatAt(i)
	}
}

// denseFloats returns a's elements as a plain float64 slice, sharing
// the payload when a is a dense float64 array and copying the widened
// values otherwise. The values read are the ones the accessor
// returned; a shared slice is read-only, and only an array the caller
// owns may be written through it. A caller sweeping the elements of a
// solver's result uses this instead of one accessor call per element.
func denseFloats(a *core.Array) []float64 {
	if !a.Strided() && a.Dtype() == core.Float {
		return a.RawFloats()
	}
	out := make([]float64, a.Len())
	for i := range out {
		out[i] = a.FloatAt(i)
	}
	return out
}

// odeRun drives the adaptive Dormand-Prince loop over the whole span.
// When watch is not nil it is called after every accepted step with
// the interval just integrated and clones of the states at both ends;
// a true return stops the integration there, and the watch's error
// aborts it. Everything else behaves exactly like IntegrateODE.
func odeRun(f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y0 *core.Array, opts ODEOptions,
	watch func(tPrev, tNow float64, yPrev, yNow []float64) (bool, error)) (*core.Array, error) {
	const name = "IntegrateODE"
	y, err := odeCheck(name, y0, &opts)
	if err != nil {
		return nil, err
	}
	n := len(y)
	w := &odeWork{}
	w.useStage(n)
	k := make([][]float64, 8) // k[1..7] are the stages; k[0] unused
	for i := 1; i <= 7; i++ {
		k[i] = make([]float64, n)
	}
	// One scratch accumulator serves every stage of every step: it is
	// rebuilt from y at the top of each stage call and read only by
	// that call's f evaluation, the same transient view the package's
	// fixed-step solvers hand out.
	acc := w.stage
	stage := func(i int, t float64, h float64) error {
		copy(acc, y)
		row := odeA[i-1]
		for j := 1; j < i; j++ {
			if row[j-1] == 0 {
				continue
			}
			// The step-scaled weight is one product, the same (h·a)·k
			// grouping the plain component loop evaluated.
			hj := h * row[j-1]
			kj := k[j]
			for m := range n {
				acc[m] += hj * kj[m]
			}
		}
		out, err := odeCall(name, f, t+odeC[i-1]*h, acc, n, &w.views)
		if err != nil {
			return err
		}
		readVector(k[i], out)
		return nil
	}

	t := t0
	h := odeInitialStep(t0, t1, y)
	budget := odeBudget{max: opts.MaxSteps}
	yEnd := w.yEnd
	// The solution weights scaled by the step size: one product each,
	// the same (h·b)·k grouping the component loop evaluated.
	var hb5, hb4 [7]float64
	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)
		for j := 1; j <= 7; j++ {
			hb5[j-1] = h * odeB5[j-1]
			hb4[j-1] = h * odeB4[j-1]
		}
		for i := 1; i <= 7; i++ {
			if err := stage(i, t, h); err != nil {
				return nil, err
			}
		}
		// The embedded pair: the 5th-order solution advances, the gap
		// to the 4th-order one estimates the local error.
		errNorm := 0.0
		for m := range n {
			y5, y4 := y[m], y[m]
			for j := 1; j <= 7; j++ {
				y5 += hb5[j-1] * k[j][m]
				y4 += hb4[j-1] * k[j][m]
			}
			yEnd[m] = y5
			scale := opts.AbsTol + opts.RelTol*math.Max(math.Abs(y[m]), math.Abs(y5))
			ratio := (y5 - y4) / scale
			errNorm += ratio * ratio
		}
		errNorm = math.Sqrt(errNorm/float64(n)) + 1e-10

		factor := math.Min(5, math.Max(0.2, 0.9*math.Pow(1/errNorm, 1.0/5)))
		if errNorm <= 1 {
			if watch != nil {
				stop, werr := watch(t, t+h, cloneDenseSlice(y), cloneDenseSlice(yEnd))
				if werr != nil {
					return nil, werr
				}
				if stop {
					return arrayFromVector(yEnd), nil
				}
			}
			copy(y, yEnd)
			prevT := t
			t += h
			h *= factor
			// Collapse is "t did not move", not "h is small": a span
			// far below the absolute time scale is perfectly
			// integrable, and the old absolute floor refused it.
			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 *= math.Max(0.2, factor)
		}
	}
	return arrayFromVector(y), nil
}

// IntegrateODEPath integrates y' = f(t, y) from t0 to t1 and returns
// the trajectory sampled at nSamples evenly spaced points, endpoints
// included: times[i] is the sample time and states[i] the state there,
// so states[0] is the initial state and states[nSamples−1] the answer
// IntegrateODE would return. Every interval between neighbouring
// samples is integrated on its own, so the adaptive step control never
// has to align with the sampling grid. Backward integration
// (t1 < t0) works, and the error contract of IntegrateODE applies
// per interval.
func IntegrateODEPath(f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y0 *core.Array, nSamples int, opts ODEOptions) ([]float64, []*core.Array, error) {
	if nSamples < 2 {
		return nil, nil, base.Errf("IntegrateODEPath: nSamples must be ≥ 2, got %d", nSamples)
	}
	y, err := odeCheck("IntegrateODEPath", y0, &opts)
	if err != nil {
		return nil, nil, base.Errf("IntegrateODEPath: %w", err)
	}
	times := make([]float64, nSamples)
	states := make([]*core.Array, nSamples)
	times[0] = t0
	states[0] = wrapVector(y)
	// The last sample is pinned to t1 exactly; the intermediate ones
	// are the evenly spaced grid.
	for i := 1; i < nSamples; i++ {
		times[i] = t0 + float64(i)*(t1-t0)/float64(nSamples-1)
	}
	times[nSamples-1] = t1
	for i := 1; i < nSamples; i++ {
		states[i], err = IntegrateODE(f, times[i-1], times[i], states[i-1], opts)
		if err != nil {
			return nil, nil, base.Errf("IntegrateODEPath: %w", err)
		}
	}
	return times, states, nil
}

// IntegrateODESteps integrates y' = f(t, y) from t0 to t1 and returns
// the trajectory as recorded at every accepted solver step: times[i]
// carries states[i] = y(times[i]), starting with (t0, y0) and ending
// with (t1, y(t1)). The accepted steps are where the adaptive control
// judged the local error within tolerance, so they are the natural
// interpolation nodes for post-processing, sensitivity analysis and
// adjoint passes. Backward integration records descending times; the
// error contract of IntegrateODE applies.
func IntegrateODESteps(f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y0 *core.Array, opts ODEOptions) ([]float64, []*core.Array, error) {
	y, err := odeCheck("IntegrateODESteps", y0, &opts)
	if err != nil {
		return nil, nil, base.Errf("IntegrateODESteps: %w", err)
	}
	times := []float64{t0}
	states := []*core.Array{wrapVector(y)}
	watch := func(tPrev, tNow float64, yPrev, yNow []float64) (bool, error) {
		times = append(times, tNow)
		// yNow is the run's own per-call clone; nothing aliases it.
		states = append(states, wrapVector(yNow))
		return false, nil
	}
	if _, err := odeRun(f, t0, t1, y0, opts, watch); err != nil {
		return nil, nil, base.Errf("IntegrateODESteps: %w", err)
	}
	// The run's last accepted boundary is t+h with h = t1−t, which
	// rounds a few ulps off t1 whenever the magnitudes demand it; the
	// documented endpoint is t1 exactly, and the recorded state there
	// is already the run's answer y(t1).
	times[len(times)-1] = t1
	return times, states, nil
}

// IntegrateRK4 integrates y' = f(t, y) with the classical fixed-step
// fourth-order Runge-Kutta scheme over the given number of equal
// steps, returning y(t1).
func IntegrateRK4(f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y0 *core.Array, steps int) (*core.Array, error) {
	if steps <= 0 {
		return nil, base.Errf("IntegrateRK4: steps must be ≥ 1, got %d", steps)
	}
	y, err := odeCheck("IntegrateRK4", y0, nil)
	if err != nil {
		return nil, err
	}
	n := len(y)
	h := (t1 - t0) / float64(steps)
	k1 := make([]float64, n)
	k2 := make([]float64, n)
	k3 := make([]float64, n)
	k4 := make([]float64, n)
	tmp := make([]float64, n)
	views := &odeViews{}
	call := func(t float64, v []float64, out []float64) error {
		o, err := odeCall("IntegrateRK4", f, t, v, n, views)
		if err != nil {
			return err
		}
		readVector(out, o)
		// A non-finite stage flows straight into the state with no
		// rejection mechanism to catch it, and the fixed-step run
		// would publish NaN with a nil error; the adaptive drivers
		// reject it, this one has to refuse it.
		for i := range n {
			if math.IsNaN(out[i]) || math.IsInf(out[i], 0) {
				return base.Errf("IntegrateRK4: f returned the non-finite value %g at coordinate %d, t=%g", out[i], i, t)
			}
		}
		return nil
	}
	// The stage times come from the exact grid t0 + i·h, never from an
	// accumulated t += h: the addition's rounding walks over a long run
	// (measured on y' = cos t from t0 = 1e6 the walk contributes an
	// error of 1.1e-8 that no step count refines away), while each grid
	// point carries a single rounding that stays put.
	for i := range steps {
		t := t0 + float64(i)*h
		if err := call(t, y, k1); err != nil {
			return nil, err
		}
		for i := range n {
			tmp[i] = y[i] + h*k1[i]/2
		}
		if err := call(t+h/2, tmp, k2); err != nil {
			return nil, err
		}
		for i := range n {
			tmp[i] = y[i] + h*k2[i]/2
		}
		if err := call(t+h/2, tmp, k3); err != nil {
			return nil, err
		}
		for i := range n {
			tmp[i] = y[i] + h*k3[i]
		}
		if err := call(t+h, tmp, k4); err != nil {
			return nil, err
		}
		for i := range n {
			y[i] += h * (k1[i] + 2*k2[i] + 2*k3[i] + k4[i]) / 6
		}
	}
	return arrayFromVector(y), nil
}

// IntegrateBackwardEuler integrates y' = f(t, y) with the fully
// implicit Euler scheme y_{n+1} = y_n + h·f(t_{n+1}, y_{n+1}), solving
// each step by Newton over a numerical Jacobian and the library's LU
// solver. The extra work per step is what buys stability on stiff
// systems, where the explicit schemes need step sizes far below what
// accuracy alone would ask for.
func IntegrateBackwardEuler(f func(t float64, y *core.Array) (*core.Array, error),
	t0, t1 float64, y0 *core.Array, steps int, opts ODEOptions) (*core.Array, error) {
	if steps <= 0 {
		return nil, base.Errf("IntegrateBackwardEuler: steps must be ≥ 1, got %d", steps)
	}
	y, err := odeCheck("IntegrateBackwardEuler", y0, &opts)
	if err != nil {
		return nil, err
	}
	n := len(y)
	h := (t1 - t0) / float64(steps)
	yn := cloneDenseSlice(y)
	seed := make([]float64, n)
	fy := make([]float64, n)
	// One Newton result buffer serves every step: it aliases neither the
	// state nor the seed, and each step overwrites it fully.
	zbuf := make([]float64, n)
	w := &odeWork{}
	// The step times come from the exact grid t0 + i·h, never from an
	// accumulated t += h: the addition's rounding walks over a long
	// run, while each grid point carries a single rounding that stays
	// put.
	for i := range steps {
		tNext := t0 + float64(i+1)*h
		// Newton on G(z) = z − y_n − h·f(t_{n+1}, z) = 0, seeded with
		// the semi-implicit Euler prediction.
		out, ferr := odeCall("IntegrateBackwardEuler", f, tNext, yn, n, &w.views)
		if ferr != nil {
			return nil, ferr
		}
		readVector(fy, out)
		for i := range n {
			seed[i] = yn[i] + h*fy[i]
		}
		if nerr := odeNewton("IntegrateBackwardEuler", f, w, tNext, 1, h, yn, seed,
			zbuf, opts.AbsTol, opts.RelTol); nerr != nil {
			return nil, nerr
		}
		copy(yn, zbuf)
	}
	return arrayFromVector(yn), nil
}

// errNewtonStalled marks an implicit solve whose Newton iteration ran
// out of budget or hit a singular matrix without converging. A driver
// that can shrink the step retries on it; a failed f evaluation is a
// different, fatal error.
var errNewtonStalled = errors.New("the Newton iteration did not converge")

// odeCall evaluates f at (t, v) and validates that the result is a
// vector of the expected length n, returning it unchanged. views,
// when not nil, caches the read-only wrapper handed to f, which a
// driver calling f repeatedly wants; a cold call site passes nil.
// Callers that need plain floats follow up with odeEval; callers that
// want to place the values themselves read them straight off the
// array.
func odeCall(name string, f func(t float64, y *core.Array) (*core.Array, error),
	t float64, v []float64, n int, views *odeViews) (*core.Array, error) {
	out, err := f(t, views.of(v))
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	if out.NDim() != 1 || out.Len() != n {
		return nil, base.Errf("%s: f returned shape %s, want a vector of length %d",
			name, base.ShapeText(out.Shape()), n)
	}
	return out, nil
}

// odeEval calls f at (t, v) and returns the derivative as a plain
// float64 slice.
func odeEval(name string, f func(t float64, y *core.Array) (*core.Array, error),
	t float64, v []float64, n int, views *odeViews) ([]float64, error) {
	out, err := odeCall(name, f, t, v, n, views)
	if err != nil {
		return nil, err
	}
	r := make([]float64, n)
	readVector(r, out)
	return r, nil
}

// odeBudget counts attempted solver steps against the MaxSteps option:
// the budget is spent before the step's stages are evaluated, so a
// rejected step consumes it like an accepted one.
type odeBudget struct {
	used int
	max  int
}

// spend spends one step of the budget, failing once it is exhausted.
func (b *odeBudget) spend(name string, t, t1 float64) error {
	if b.used >= b.max {
		return base.Errf("%s: reached MaxSteps=%d at t=%g before t1=%g", name, b.max, t, t1)
	}
	b.used++
	return nil
}

// odeClampStep caps h so a step never overshoots t1; t1−t carries the
// integration direction.
func odeClampStep(h, t, t1 float64) float64 {
	if math.Abs(h) > math.Abs(t1-t) {
		return t1 - t
	}
	return h
}

// odeArrived reports whether t sits within a few ulps of t1. The
// accumulating t += h can miss the exact endpoint by rounding once
// t and t1 differ in magnitude beyond Sterbenz territory, and the
// residual distance is indistinguishable from zero at working
// precision, so the solvers treat it as arrived rather than report a
// collapsed step over it.
func odeArrived(t, t1 float64) bool {
	if t == t1 {
		return true
	}
	return math.Abs(t1-t) <= 8*base.EpsF*math.Max(math.Abs(t), math.Abs(t1))
}

// odeWork holds the scratch the implicit schemes reuse across the
// steps of one run: the Newton vectors, one flat numerical Jacobian
// with its difference stencils, the factored matrix's rows, the ROS4
// stage buffers and the read-only views handed to f. Every buffer is
// overwritten before it is read, so a run builds the workspace once
// and no step allocates scratch of its own.
type odeWork struct {
	// The Newton iteration's residual, step and derivative buffers,
	// plus the DAE mass-matrix product M·z.
	g, col, fzs []float64
	mz          []float64
	// jac is the numerical Jacobian as one flat n×n buffer, row-major:
	// jac[i*n+j] is ∂f_i/∂z_j.
	jac []float64
	// The Jacobian's central-difference stencils: the two perturbed
	// states and their two results.
	zp, zm, fp, fm []float64
	// mat is the implicit relation's matrix in row-major rows, ROS4's
	// (1/(γh))I − J or Newton's α·I − h·J depending on the caller, with
	// perm the row permutation its factorisation produced. Factor
	// shuffles the rows in place, and every rebuild rewrites the lot.
	mat  [][]float64
	perm []int
	// The ROS4 stage buffers: the four divided differences, the stage
	// value, its right side, f's result and the candidate end state.
	ks             [4][]float64
	stage, rhs, fy []float64
	yEnd           []float64
	// The Newton iterate the implicit schemes iterate in, and the visit
	// bitmap the stage permutation walks: both reused across every step
	// and every attempt of one run.
	zwork   []float64
	visited []bool
	// views caches the wrapper handed to f per scratch slice.
	views odeViews
}

// use returns the workspace's buffers at the state length n, growing
// them on first use. The content is left as the previous step wrote
// it: every consumer overwrites its buffer before reading it.
func (w *odeWork) use(n int) {
	w.g = sizedBuf(w.g, n)
	w.col = sizedBuf(w.col, n)
	w.fzs = sizedBuf(w.fzs, n)
	w.mz = sizedBuf(w.mz, n)
	w.jac = sizedBuf(w.jac, n*n)
	w.zp = sizedBuf(w.zp, n)
	w.zm = sizedBuf(w.zm, n)
	w.fp = sizedBuf(w.fp, n)
	w.fm = sizedBuf(w.fm, n)
	w.mat = sizedRows(w.mat, n)
	for s := range w.ks {
		w.ks[s] = sizedBuf(w.ks[s], n)
	}
	w.stage = sizedBuf(w.stage, n)
	w.rhs = sizedBuf(w.rhs, n)
	w.fy = sizedBuf(w.fy, n)
	w.yEnd = sizedBuf(w.yEnd, n)
	w.zwork = sizedBuf(w.zwork, n)
	w.visited = sizedBools(w.visited, n)
}

// useStage returns the explicit step loop's buffers at the state
// length n, growing them on first use: the stage accumulator and the
// candidate end state are all an explicit pair needs, and sizing the
// implicit buffers here would allocate the Jacobian and the factored
// matrix for a loop that never takes a derivative.
func (w *odeWork) useStage(n int) {
	w.stage = sizedBuf(w.stage, n)
	w.yEnd = sizedBuf(w.yEnd, n)
}

// sizedBuf returns b cut to length n, reusing its storage when it is
// large enough.
func sizedBuf(b []float64, n int) []float64 {
	if cap(b) < n {
		return make([]float64, n)
	}
	return b[:n]
}

// sizedBools returns b cut to length n, reusing its storage when it is
// large enough.
func sizedBools(b []bool, n int) []bool {
	if cap(b) < n {
		return make([]bool, n)
	}
	return b[:n]
}

// odePermuteColumn reorders col in place so that col[i] takes the
// value that sat at perm[i], the permutation the workspace's LU
// factorisation produced. The walk carries each displaced value around
// its cycle exactly as the library's PermuteColumn does, plain
// assignments moving each value once, so the column ends up bit-
// identical; the difference is that the visit bitmap is the caller's
// reused scratch rather than a fresh allocation per call. The bitmap
// is cleared on entry, so a dirty buffer behaves exactly like a fresh
// one.
func odePermuteColumn(col []float64, perm []int, visited []bool) {
	vis := visited[:len(col)]
	clear(vis)
	for i := range col {
		if vis[i] || perm[i] == i {
			vis[i] = true
			continue
		}
		// Carry the displaced value around the cycle.
		tmp := col[i]
		j := i
		for {
			vis[j] = true
			k := perm[j]
			if k == i {
				break
			}
			col[j] = col[k]
			j = k
		}
		col[j] = tmp
	}
}

// arrayFromVector copies a float64 slice into a fresh rank-1 float64
// array: the trajectory endpoint contract, without the intermediate
// wrapper a cloneArray(wrapVector(...)) pair built. The result never
// aliases the input.
func arrayFromVector(v []float64) *core.Array {
	a := core.New(core.Float, len(v))
	copy(a.RawFloats(), v)
	return a
}

// sizedRows returns m as n row slices of length n, reusing the rows it
// already holds. The rows' content is the caller's to overwrite.
func sizedRows(m [][]float64, n int) [][]float64 {
	if cap(m) < n {
		m = make([][]float64, n)
	}
	m = m[:n]
	for i := range m {
		m[i] = sizedBuf(m[i], n)
	}
	return m
}

// odeViews caches the read-only wrapper handed to f for one scratch
// slice, so a driver that calls f thousands of times builds the
// wrapper once per slice instead of once per call. The values behind
// the wrapper are the driver's own scratch and keep changing exactly
// as they did; only the Array header is reused. The slices a driver
// hands in are few, so a linear scan beats a map; a driver that
// presents a fresh slice every call cannot grow the cache without
// bound, because the oldest entry makes way.
type odeViews struct {
	entries []odeView
}

type odeView struct {
	vals []float64
	arr  *core.Array
}

// odeViewSlots bounds the cache. A driver holds a handful of scratch
// slices at once, and each slot pins one sliced buffer, so the bound
// keeps both the scan and the retention small.
const odeViewSlots = 8

// of returns a read-only view of s, reusing the one already built for
// that slice. A nil cache builds a fresh view, which is what a cold
// call site wants.
func (v *odeViews) of(s []float64) *core.Array {
	if v == nil || len(s) == 0 {
		return wrapVector(s)
	}
	for i := range v.entries {
		e := &v.entries[i]
		if len(e.vals) == len(s) && &e.vals[0] == &s[0] {
			return e.arr
		}
	}
	arr := wrapVector(s)
	if len(v.entries) == odeViewSlots {
		copy(v.entries, v.entries[1:])
		v.entries = v.entries[:odeViewSlots-1]
	}
	v.entries = append(v.entries, odeView{vals: s, arr: arr})
	return arr
}

// odeJacobian fills the workspace's flat Jacobian with the central
// differences of f at (t, z), one column per state component:
// entry i*n+j is ∂f_i/∂z_j. The returned slice is the workspace's, so
// it stays valid until the next Jacobian. The two perturbed stencils
// and their two result buffers are reused across columns: each round
// rebuilds the stencils from z and overwrites both results before
// reading them.
func odeJacobian(name string, f func(t float64, y *core.Array) (*core.Array, error),
	t float64, z []float64, w *odeWork) ([]float64, error) {
	n := len(z)
	w.use(n)
	jac := w.jac
	for j := range n {
		eps := math.Sqrt(base.EpsF) * math.Max(1, math.Abs(z[j]))
		copy(w.zp, z)
		copy(w.zm, z)
		w.zp[j] += eps
		w.zm[j] -= eps
		out, e1 := odeCall(name, f, t, w.zp, n, &w.views)
		if e1 != nil {
			return nil, e1
		}
		readVector(w.fp, out)
		out, e2 := odeCall(name, f, t, w.zm, n, &w.views)
		if e2 != nil {
			return nil, e2
		}
		readVector(w.fm, out)
		for i := range n {
			jac[i*n+j] = (w.fp[i] - w.fm[i]) / (2 * eps)
		}
	}
	return jac, nil
}

// odeNewton solves the implicit step equation α·z − h·f(tNext, z) = β
// for z by Newton over a numerical Jacobian and the library's LU
// solver, writing the converged state into dst and returning an error
// otherwise. The Jacobian is frozen from the seed and rebuilt twice
// when convergence drags, so a converged step costs one Jacobian and a
// handful of f evaluations. The iteration runs in the workspace's own
// buffer and the column permutation walks the workspace's bitmap, so
// the only allocation a converged step costs is the caller's dst: a
// driver whose dst recycles through a ring allocates its result
// buffers once per solve, not once per step, and a rejected or stalled
// attempt allocates nothing. dst must not alias seed or beta; the
// workspace overwrites it fully at convergence. Convergence is
// measured on the residual against the error scale the caller
// integrates to, an order of magnitude below it, but never below the
// floating-point floor of the residual's own terms, which would
// otherwise be unreachable at the tiny steps a stiff start begins
// with. An iteration that outlives twenty rounds, or a singular Newton
// matrix, surfaces as errNewtonStalled so a stepping driver can retry
// with a smaller step; an f that fails is the fatal error it is.
func odeNewton(name string, f func(t float64, y *core.Array) (*core.Array, error),
	w *odeWork, tNext float64, alpha, h float64, beta, seed, dst []float64,
	absTol, relTol float64) error {
	n := len(seed)
	w.use(n)
	// The iterate starts as the seed and stays the workspace's buffer:
	// the views cache then hands f one stable wrapper for every Newton
	// call of the run.
	z := sizedBuf(w.zwork, n)
	copy(z, seed)
	for iteration := range 20 {
		out, err := odeCall(name, f, tNext, z, n, &w.views)
		if err != nil {
			return err
		}
		readVector(w.fzs, out)
		worst, terms := 0.0, 0.0
		for i := range n {
			w.g[i] = alpha*z[i] - h*w.fzs[i] - beta[i]
			worst = math.Max(worst, math.Abs(w.g[i]))
			terms = math.Max(terms, math.Abs(alpha*z[i])+math.Abs(h*w.fzs[i])+math.Abs(beta[i]))
		}
		limit := math.Max(0.1*(absTol+relTol*normInfOfStep(z)), 8*base.EpsF*terms)
		if worst <= limit {
			copy(dst, z)
			return nil
		}
		if iteration == 0 || iteration == 4 || iteration == 10 {
			jac, jerr := odeJacobian(name, f, tNext, z, w)
			if jerr != nil {
				return jerr
			}
			// Newton matrix α·I − h·J, a fresh LU for the frozen
			// Jacobian; the iterations that follow only substitute.
			// Every row is rebuilt entry by entry before the
			// factorisation reads it.
			for i := range n {
				row := w.mat[i]
				for j := range n {
					row[j] = -h * jac[i*n+j]
				}
				row[i] += alpha
			}
			w.perm, _ = base.Factor(w.mat)
			if err := base.CheckSingular(name, w.mat); err != nil {
				return base.Errf("%s: %w, singular Newton matrix at t=%g",
					name, errNewtonStalled, tNext)
			}
		}
		for i := range n {
			w.col[i] = -w.g[i]
		}
		odePermuteColumn(w.col, w.perm, w.visited)
		base.SolveColumn(w.mat, w.col)
		for i := range n {
			z[i] += w.col[i]
		}
	}
	return base.Errf("%s: %w at t=%g", name, errNewtonStalled, tNext)
}

// odeCheck validates the initial state, applies option defaults and
// returns the flat float64 working state.
func odeCheck(name string, y0 *core.Array, opts *ODEOptions) ([]float64, error) {
	if y0.NDim() != 1 {
		return nil, base.Errf("%s: the state must be a vector, got shape %s", name, base.ShapeText(y0.Shape()))
	}
	if y0.Len() == 0 {
		return nil, base.Errf("%s: the state must not be empty", name)
	}
	if y0.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex states are not supported", name)
	}
	if opts != nil {
		if opts.RelTol <= 0 {
			opts.RelTol = 1e-6
		}
		if opts.AbsTol <= 0 {
			opts.AbsTol = 1e-9
		}
		if opts.MaxSteps <= 0 {
			opts.MaxSteps = 100000
		}
	}
	return cloneDense(y0), nil
}

// odeInitialStep guesses the first step size as a small fraction of
// the integration span, carrying the direction in its sign.
func odeInitialStep(t0, t1 float64, y []float64) float64 {
	h := 0.01 * math.Abs(t1-t0)
	if h == 0 {
		h = 1e-6
	}
	// The step sign carries the integration direction: a t1 < t0
	// integrates backwards.
	if t1 < t0 {
		h = -h
	}
	return h
}

// cloneArray copies an array element by element, so the integration
// steps' results never alias a buffer already handed out.
func cloneArray(a *core.Array) *core.Array {
	out := core.New(a.Dtype(), a.Shape()...)
	switch a.Dtype() {
	case core.Float:
		copy(out.RawFloats(), a.RawFloats())
	case core.Float32:
		copy(out.RawFloat32s(), a.RawFloat32s())
	case core.Int:
		copy(out.RawInts(), a.RawInts())
	default:
		copy(out.RawComplexes(), a.RawComplexes())
	}
	return out
}

// 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
}

// wrapVector views a float64 slice as a rank-1 Array without copying.
// The caller must treat the result as read-only.
func wrapVector(v []float64) *core.Array {
	a, _ := core.FloatsFromArray(v, len(v))
	return a
}

// cloneDense copies an array's elements into a plain float64 slice,
// widening int and float32 elements exactly.
func cloneDense(y0 *core.Array) []float64 {
	vals := make([]float64, y0.Len())
	for i := range vals {
		vals[i] = y0.FloatAt(i)
	}
	return vals
}

// cloneDenseSlice copies a float64 slice.
func cloneDenseSlice(v []float64) []float64 {
	out := make([]float64, len(v))
	copy(out, v)
	return out
}