// 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 (
	"math"
	"strings"
	"testing"
)

// TestROS4FixedStepOrder pins the fourth order of the scheme: on the
// oscillator, whose exact rotation is known, uniform steps must shrink
// the global error by roughly sixteen per halving. The convergence is
// measured against the driven step, the way the order is defined.
func TestROS4FixedStepOrder(t *testing.T) {
	oscillator := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{y.FloatAt(1), -y.FloatAt(0)}, 2)
	}
	errAt := func(steps int) float64 {
		h := 1.0 / float64(steps)
		y := []float64{1, 0}
		now := 0.0
		w := &odeWork{}
		for range steps {
			yEnd, _, err := ros4Step("TestROS4FixedStepOrder", oscillator, w, now, y, h, 1e-300, 1e-300)
			if err != nil {
				t.Fatalf("ros4Step: %v", err)
			}
			copy(y, yEnd)
			now += h
		}
		return math.Max(math.Abs(y[0]-math.Cos(1)), math.Abs(y[1]+math.Sin(1)))
	}
	e8, e16, e32 := errAt(8), errAt(16), errAt(32)
	if e8 < 1e-13 {
		t.Skipf("error already at round-off (%v)", e8)
	}
	for _, r := range []float64{e8 / e16, e16 / e32} {
		if r < 12 || r > 20 {
			t.Fatalf("error ratio over a halved step = %.2g, want ≈ 16 for a fourth-order scheme", r)
		}
	}
}

// TestROS4NonAutonomousDegradation pins the documented limit the
// autonomous fourth order carries with it: on y' = −y + t, whose exact
// answer y = t − 1 + e^{−t} is known, the tableau's missing
// time-derivative weights cost the second-order local terms and the
// uniform-step ratios sit near 2, first order, not near 16. A change
// that lifts this must move the doc comment with it.
func TestROS4NonAutonomousDegradation(t *testing.T) {
	forced := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{-y.FloatAt(0) + t}, 1)
	}
	errAt := func(steps int) float64 {
		h := 1.0 / float64(steps)
		y := []float64{0}
		now := 0.0
		w := &odeWork{}
		for range steps {
			yEnd, _, err := ros4Step("TestROS4NonAutonomousDegradation", forced, w, now, y, h, 1e-300, 1e-300)
			if err != nil {
				t.Fatalf("ros4Step: %v", err)
			}
			copy(y, yEnd)
			now += h
		}
		return math.Abs(y[0] - (1 - 1/math.E))
	}
	e8, e16, e32 := errAt(8), errAt(16), errAt(32)
	for _, r := range []float64{e8 / e16, e16 / e32} {
		if r >= 4 {
			t.Fatalf("error ratio over a halved step = %.2g, the forced system runs at first order (the doc comment names this limit)", r)
		}
	}
}

// TestIntegrateROS4Quadrature pins exactness on y' = 1: the constant
// right side is reproduced to rounding whatever the accepted steps do.
func TestIntegrateROS4Quadrature(t *testing.T) {
	one := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{1}, 1)
	}
	end, err := IntegrateROS4(one, 0, 1, mustFloats(t, []float64{0}), ODEOptions{})
	if err != nil {
		t.Fatalf("IntegrateROS4: %v", err)
	}
	if math.Abs(end.FloatAt(0)-1) > 1e-12 {
		t.Fatalf("y(1) = %.16g, want 1 to rounding", end.FloatAt(0))
	}
}

// TestIntegrateROS4Accuracy checks the adaptive driver on the analytic
// decay and over a full oscillator period with a two-dimensional state.
func TestIntegrateROS4Accuracy(t *testing.T) {
	end, err := IntegrateROS4(decay, 0, 1, mustFloats(t, []float64{1}),
		ODEOptions{RelTol: 1e-8, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateROS4: %v", err)
	}
	if math.Abs(end.FloatAt(0)-math.Exp(-1)) > 1e-7 {
		t.Fatalf("y(1) = %.14g, want %.14g ± 1e-7", end.FloatAt(0), math.Exp(-1))
	}
	oscillator := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{y.FloatAt(1), -y.FloatAt(0)}, 2)
	}
	full, err := IntegrateROS4(oscillator, 0, 2*math.Pi, mustFloats(t, []float64{1, 0}),
		ODEOptions{RelTol: 1e-8, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateROS4 oscillator: %v", err)
	}
	if math.Abs(full.FloatAt(0)-1) > 1e-4 || math.Abs(full.FloatAt(1)) > 1e-4 {
		t.Fatalf("full period = (%.10g, %.10g), want (1, 0)",
			full.FloatAt(0), full.FloatAt(1))
	}
	back, err := IntegrateROS4(decay, 1, 0, mustFloats(t, []float64{math.Exp(-1)}),
		ODEOptions{RelTol: 1e-8, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateROS4 backward: %v", err)
	}
	if math.Abs(back.FloatAt(0)-1) > 1e-7 {
		t.Fatalf("backward y(0) = %.14g, want 1 ± 1e-7", back.FloatAt(0))
	}
}

// TestIntegrateROS4LStability pins the L-stable damping: on y' =
// −10^8(y − 1) and y' = −10^8 y the steps are far beyond the transient
// and a non-L-stable scheme blows up, while the W scheme lands on the
// forcing, respectively on zero, with a bounded step count.
func TestIntegrateROS4LStability(t *testing.T) {
	rise := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{-1e8 * (y.FloatAt(0) - 1)}, 1)
	}
	end, err := IntegrateROS4(rise, 0, 1, mustFloats(t, []float64{0}), ODEOptions{MaxSteps: 1000})
	if err != nil {
		t.Fatalf("IntegrateROS4 stiff rise: %v", err)
	}
	if v := end.FloatAt(0); math.IsNaN(v) || math.Abs(v-1) > 1e-9 {
		t.Fatalf("stiff rise y(1) = %.14g, want 1", v)
	}
	decayStiff := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{-1e8 * y.FloatAt(0)}, 1)
	}
	end, err = IntegrateROS4(decayStiff, 0, 1, mustFloats(t, []float64{1}), ODEOptions{MaxSteps: 1000})
	if err != nil {
		t.Fatalf("IntegrateROS4 stiff decay: %v", err)
	}
	if v := end.FloatAt(0); math.IsNaN(v) || math.Abs(v) > 1e-9 {
		t.Fatalf("stiff decay y(1) = %.14g, want 0", v)
	}
	// The one-step amplification at a stiff eigenvalue must die out:
	// a giant step on the decay damps the state by orders of magnitude
	// instead of amplifying it.
	yEnd, _, err := ros4Step("TestIntegrateROS4LStability", decay, &odeWork{}, 0, []float64{1}, 1e6, 1e-3, 1e-3)
	if err != nil {
		t.Fatalf("ros4Step at z = 1e6: %v", err)
	}
	if math.Abs(yEnd[0]) > 1e-4 {
		t.Fatalf("one step at h·λ = 1e6 multiplied the state by %.3g, want heavy damping", yEnd[0])
	}
}

// TestIntegrateROS4VanDerPol integrates the Van der Pol oscillator in
// the stiff relaxation regime: μ = 1000 carries a transient of width
// 10^−3 under a slow motion, and the W scheme must cross it and follow
// the slow branch inside the step budget.
func TestIntegrateROS4VanDerPol(t *testing.T) {
	const mu = 1000.0
	vdp := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{
			y.FloatAt(1),
			mu*(1-y.FloatAt(0)*y.FloatAt(0))*y.FloatAt(1) - y.FloatAt(0),
		}, 2)
	}
	end, err := IntegrateROS4(vdp, 0, 2, mustFloats(t, []float64{2, 0}),
		ODEOptions{MaxSteps: 100000})
	if err != nil {
		t.Fatalf("IntegrateROS4 Van der Pol: %v", err)
	}
	for i := range 2 {
		if math.IsNaN(end.FloatAt(i)) || math.IsInf(end.FloatAt(i), 0) {
			t.Fatalf("Van der Pol state[%d] = %g left the finite range", i, end.FloatAt(i))
		}
	}
	// The trajectory returns onto the slow branch near x = 2 with a
	// small velocity; anything else means the jump was not resolved.
	if math.Abs(end.FloatAt(0)-2) > 0.01 || math.Abs(end.FloatAt(1)) > 0.01 {
		t.Fatalf("Van der Pol end = (%.10g, %.10g), want the slow branch near (2, 0)",
			end.FloatAt(0), end.FloatAt(1))
	}
}

// TestIntegrateROS4Errors pins the error contract: a degenerate span
// returns the initial state unchanged, a wrong-shaped f, a rank-2
// state, an empty state, an exhausted step budget and an f that blows
// up mid-span are errors.
func TestIntegrateROS4Errors(t *testing.T) {
	y0 := mustFloats(t, []float64{1})
	same, err := IntegrateROS4(decay, 1, 1, y0, ODEOptions{})
	if err != nil {
		t.Fatalf("zero span: %v", err)
	}
	if math.Abs(same.FloatAt(0)-1) > 0 {
		t.Fatalf("zero span moved the state to %v", same.FloatAt(0))
	}
	wrongShape := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{1, 1}, 2)
	}
	if _, err := IntegrateROS4(wrongShape, 0, 1, y0, ODEOptions{}); err == nil {
		t.Fatal("expected an error when f returns the wrong shape")
	}
	matrixState := mustFloats(t, []float64{1, 1}, 1, 2)
	if _, err := IntegrateROS4(decay, 0, 1, matrixState, ODEOptions{}); err == nil {
		t.Fatal("expected an error for a rank-2 state")
	}
	if _, err := IntegrateROS4(decay, 0, 1, mustFloats(t, nil), ODEOptions{}); err == nil {
		t.Fatal("expected an error for an empty state")
	}
	if _, err := IntegrateROS4(decay, 0, 1, y0, ODEOptions{MaxSteps: 2}); err == nil {
		t.Fatal("expected an error for an exhausted step budget")
	} else if !strings.Contains(err.Error(), "MaxSteps=2") {
		t.Fatalf("want a step-budget error, got %v", err)
	}
	boom := func(t float64, y *core.Array) (*core.Array, error) {
		if t > 0.5 {
			return nil, base.Errf("detector tripped")
		}
		return core.MulF(y, -1), nil
	}
	if _, err := IntegrateROS4(boom, 0, 1, y0, ODEOptions{}); err == nil {
		t.Fatal("expected the operator error to propagate")
	}
	nan := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{math.NaN()}, 1)
	}
	if _, err := IntegrateROS4(nan, 0, 1, y0, ODEOptions{MaxSteps: 200}); err == nil {
		t.Fatal("expected an error when f returns NaN throughout")
	}
}

// TestROS4StepErrors pins the error paths of a single attempted step:
// an f failing inside the numerical Jacobian and an f failing at the
// later stage times are both fatal to the step.
func TestROS4StepErrors(t *testing.T) {
	always := func(t float64, y *core.Array) (*core.Array, error) {
		return nil, base.Errf("detector tripped")
	}
	if _, _, err := ros4Step("TestROS4StepErrors", always, &odeWork{}, 0, []float64{1}, 0.1, 1e-6, 1e-9); err == nil {
		t.Fatal("expected the Jacobian's f error to propagate")
	}
	gated := func(t float64, y *core.Array) (*core.Array, error) {
		if t > 0 {
			return nil, base.Errf("detector tripped")
		}
		return core.FromFloats([]float64{0}, 1)
	}
	if _, _, err := ros4Step("TestROS4StepErrors", gated, &odeWork{}, 0, []float64{1}, 0.1, 1e-6, 1e-9); err == nil {
		t.Fatal("expected the stage f error to propagate")
	}
}