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

// stiffCosine returns f for y' = −k(y − cos t), the canonical stiff
// problem: a slow forcing with a transient decaying at rate k.
func stiffCosine(k float64) func(t float64, y *core.Array) (*core.Array, error) {
	return func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{-k * (y.FloatAt(0) - math.Cos(t))}, 1)
	}
}

// TestIntegrateBDF2Stiff is the demonstration the stiff solver exists
// for: y' = −10^5(y − cos t) carries a transient of width 10^−5 under
// a slow forcing, and BDF2 crosses it and follows the forcing to t=1
// inside a 2000-step budget, landing on the exact solution
// y(1) = (k²·cos 1 + k·sin 1)/(k² + 1).
func TestIntegrateBDF2Stiff(t *testing.T) {
	const k = 1e5
	end, err := IntegrateBDF2(stiffCosine(k), 0, 1, mustFloats(t, []float64{0}),
		ODEOptions{MaxSteps: 2000})
	if err != nil {
		t.Fatalf("IntegrateBDF2: %v", err)
	}
	want := (k*k*math.Cos(1) + k*math.Sin(1)) / (k*k + 1)
	if math.Abs(end.FloatAt(0)-want) > 1e-6 {
		t.Fatalf("y(1) = %.14g, want %.14g", end.FloatAt(0), want)
	}
}

// TestIntegrateBDF2StiffBeatsExplicit shows the same problem is out
// of reach for the explicit pair: stability pins DOPRI to steps of
// order 1/k, so a 5000-step budget dies a fifth of the way in.
func TestIntegrateBDF2StiffBeatsExplicit(t *testing.T) {
	_, err := IntegrateODE(stiffCosine(1e5), 0, 1, mustFloats(t, []float64{0}),
		ODEOptions{MaxSteps: 5000})
	if err == nil {
		t.Fatal("explicit DOPRI was expected to exhaust its step budget on the stiff problem")
	}
	if !strings.Contains(err.Error(), "MaxSteps=5000") {
		t.Fatalf("want a step-budget error, got %v", err)
	}
}

// TestBDF2FixedStepOrder verifies the second order of the underlying
// formula directly: with exact history on y' = −y and uniform steps,
// halving h must quarter the global error. The Milne constant for
// equal steps is pinned to 2/11 along the way.
func TestBDF2FixedStepOrder(t *testing.T) {
	if got := bdf2Milne(1, 2, 0, -1); math.Abs(got-2.0/11) > 1e-12 {
		t.Fatalf("bdf2Milne for equal steps = %.14g, want 2/11", got)
	}
	errAt := func(steps int) float64 {
		h := 1.0 / float64(steps)
		alpha := 1.5 / h
		now := 0.0
		yn := []float64{1}
		yNm1 := []float64{math.Exp(h)} // exact history at t−h
		w := &odeWork{}
		for range steps {
			tNext := now + h
			beta := []float64{2*yn[0]/h - yNm1[0]/(2*h)}
			z := make([]float64, 1)
			err := odeNewton("TestBDF2FixedStepOrder", decay, w, tNext, alpha, 1,
				beta, []float64{math.Exp(-tNext)}, z, 1e-13, 1e-13)
			if err != nil {
				t.Fatalf("odeNewton: %v", err)
			}
			yNm1 = yn
			yn = z
			now = tNext
		}
		return math.Abs(yn[0] - math.Exp(-1))
	}
	e20, e40 := errAt(20), errAt(40)
	if e20 < 1e-12 {
		t.Skipf("error already at round-off (%v)", e20)
	}
	ratio := e20 / e40
	if ratio < 3 || ratio > 5.2 {
		t.Fatalf("error ratio over a halved step = %.2g, want ≈ 4 for a second-order scheme", ratio)
	}
}

// TestIntegrateBDF2Accuracy checks the adaptive driver on a smooth
// problem against the analytic decay at a tolerance far below the
// default, and over a full oscillator period with a two-dimensional
// state, exercising the vector Newton path.
func TestIntegrateBDF2Accuracy(t *testing.T) {
	end, err := IntegrateBDF2(decay, 0, 1, mustFloats(t, []float64{1}),
		ODEOptions{RelTol: 1e-8, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateBDF2: %v", err)
	}
	if math.Abs(end.FloatAt(0)-math.Exp(-1)) > 1e-5 {
		t.Fatalf("y(1) = %.14g, want %.14g ± 1e-5", 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 := IntegrateBDF2(oscillator, 0, 2*math.Pi, mustFloats(t, []float64{1, 0}),
		ODEOptions{RelTol: 1e-8, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateBDF2 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))
	}
}

// TestIntegrateBDF2Backward integrates the decay backwards from t=1
// to t=0; the signed-step formulation must return the start value.
func TestIntegrateBDF2Backward(t *testing.T) {
	end, err := IntegrateBDF2(decay, 1, 0, mustFloats(t, []float64{math.Exp(-1)}),
		ODEOptions{RelTol: 1e-8, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateBDF2 backward: %v", err)
	}
	if math.Abs(end.FloatAt(0)-1) > 1e-5 {
		t.Fatalf("backward y(0) = %.14g, want 1 ± 1e-5", end.FloatAt(0))
	}
}

// TestIntegrateBDF2Errors pins the error contract: a degenerate span
// returns the initial state unchanged, a wrong-shaped f, a rank-2
// state, an empty state and an exhausted step budget are errors.
func TestIntegrateBDF2Errors(t *testing.T) {
	y0 := mustFloats(t, []float64{1})
	same, err := IntegrateBDF2(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 := IntegrateBDF2(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 := IntegrateBDF2(decay, 0, 1, matrixState, ODEOptions{}); err == nil {
		t.Fatal("expected an error for a rank-2 state")
	}
	if _, err := IntegrateBDF2(decay, 0, 1, mustFloats(t, nil), ODEOptions{}); err == nil {
		t.Fatal("expected an error for an empty state")
	}
	if _, err := IntegrateBDF2(decay, 0, 1, y0, ODEOptions{MaxSteps: 2}); err == nil {
		t.Fatal("expected an error for an exhausted step budget")
	}
	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 := IntegrateBDF2(boom, 0, 1, y0, ODEOptions{}); err == nil {
		t.Fatal("expected the operator error to propagate")
	}
}

// TestBDF2InitialStepBackwardProbe pins the probe direction: on a
// backward span the initial-step probe must sample the dynamics at
// t0 − h0, not extrapolate forward, so every evaluation time either
// sits before t0 or the probe is wrong.
func TestBDF2InitialStepBackwardProbe(t *testing.T) {
	var calls []float64
	f := func(tt float64, y *core.Array) (*core.Array, error) {
		calls = append(calls, tt)
		return mustFloats(t, []float64{0}), nil
	}
	h, err := bdf2InitialStep("TestBDF2InitialStep", f, 5, 0, []float64{1}, &ODEOptions{})
	if err != nil {
		t.Fatalf("bdf2InitialStep: %v", err)
	}
	if h >= 0 {
		t.Fatalf("backward span must yield a negative first step, got %g", h)
	}
	backward := false
	for _, c := range calls {
		if c < 5 {
			backward = true
		}
	}
	if !backward {
		t.Fatalf("the probe never stepped backward from t0 = 5, evaluated at %v", calls)
	}
}