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

// TestIntegrateBDFVarStiff is the demonstration pin: on y' =
// −10^5(y − cos t) the variable-order driver lands on the exact y(1) =
// (k²·cos 1 + k·sin 1)/(k² + 1) inside half the step budget BDF2
// needed, having raised to order five on the smooth tail.
func TestIntegrateBDFVarStiff(t *testing.T) {
	const k = 1e5
	var stats BDFVarStats
	end, err := IntegrateBDFVar(stiffCosine(k), 0, 1, mustFloats(t, []float64{0}),
		BDFVarOptions{MaxSteps: 2000, Stats: &stats})
	if err != nil {
		t.Fatalf("IntegrateBDFVar: %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)
	}
	t.Logf("stiff run: %d steps, %d rejected, max order %d", stats.Steps, stats.Rejected, stats.MaxOrder)
	if stats.MaxOrder != 5 {
		t.Fatalf("max order reached = %d, want 5 on the smooth tail", stats.MaxOrder)
	}
	if stats.Steps > 1000 {
		t.Fatalf("the run took %d steps, want well inside the 2000-step budget BDF2 needed", stats.Steps)
	}
}

// TestIntegrateBDFVarOrderAdapts instruments the order counter: the
// first accepted steps run at order one (nothing else has history), so
// a run that ends with order five must have climbed the ladder, and on
// the same stiff problem it must spend far fewer steps than an
// order-one-locked run, which is what step and order adaptation buy.
func TestIntegrateBDFVarOrderAdapts(t *testing.T) {
	const k = 1e5
	var adaptive, locked BDFVarStats
	if _, err := IntegrateBDFVar(stiffCosine(k), 0, 1, mustFloats(t, []float64{0}),
		BDFVarOptions{MaxSteps: 50000, Stats: &adaptive}); err != nil {
		t.Fatalf("IntegrateBDFVar adaptive: %v", err)
	}
	if _, err := integrateBDFVar("TestIntegrateBDFVarOrderAdapts", stiffCosine(k), 0, 1,
		mustFloats(t, []float64{0}), BDFVarOptions{MaxSteps: 50000, Stats: &locked}, 1, true); err != nil {
		t.Fatalf("IntegrateBDFVar order-one locked: %v", err)
	}
	t.Logf("adaptive run: %d steps, locked run: %d steps", adaptive.Steps, locked.Steps)
	if adaptive.Steps < 8 || locked.Steps < 8 {
		t.Fatalf("implausible step counts: adaptive %d, locked %d", adaptive.Steps, locked.Steps)
	}
	if adaptive.MaxOrder != 5 {
		t.Fatalf("adaptive run reached order %d, want 5", adaptive.MaxOrder)
	}
	if locked.MaxOrder != 1 {
		t.Fatalf("locked run reached order %d, want 1 throughout", locked.MaxOrder)
	}
	if adaptive.Steps*3 > locked.Steps {
		t.Fatalf("the adaptive run took %d steps against the locked run's %d: order adaptation did not engage",
			adaptive.Steps, locked.Steps)
	}
}

// TestIntegrateBDFVarAccuracy checks the adaptive driver on a smooth
// problem against the analytic decay, over a full oscillator period
// with a two-dimensional state, and backwards in time.
func TestIntegrateBDFVarAccuracy(t *testing.T) {
	end, err := IntegrateBDFVar(decay, 0, 1, mustFloats(t, []float64{1}),
		BDFVarOptions{RelTol: 1e-8, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateBDFVar: %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 := IntegrateBDFVar(oscillator, 0, 2*math.Pi, mustFloats(t, []float64{1, 0}),
		BDFVarOptions{RelTol: 1e-8, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateBDFVar 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 := IntegrateBDFVar(decay, 1, 0, mustFloats(t, []float64{math.Exp(-1)}),
		BDFVarOptions{RelTol: 1e-8, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateBDFVar backward: %v", err)
	}
	if math.Abs(back.FloatAt(0)-1) > 1e-5 {
		t.Fatalf("backward y(0) = %.14g, want 1 ± 1e-5", back.FloatAt(0))
	}
}

// TestBDFVarCoefficientsMatchBDF2 pins the coefficient recurrence: at
// order two the divided-difference form must reproduce the shipped
// bdf2Coefficients, on equal steps and on skewed ones, in α, β and the
// predictor seed alike.
func TestBDFVarCoefficientsMatchBDF2(t *testing.T) {
	patterns := []struct{ tNext, t, tNm1, tNm2 float64 }{
		{3, 2, 1, 0},
		{1.3, 0.75, 0.4, -0.1},
		{5, 1, 0.5, -2},
	}
	vals := []float64{2.5, -3, 7} // y at tNm2, tNm1, t
	for _, p := range patterns {
		hist := &bdfVarHistory{}
		for i, tt := range []float64{p.tNm2, p.tNm1, p.t} {
			hist.push(tt, []float64{vals[i]})
		}
		beta := make([]float64, 1)
		seed := make([]float64, 1)
		alpha := bdfVarCoefficients(2, p.tNext, hist, beta, seed,
			make([]float64, bdfVarKeep), make([]float64, bdfVarKeep), make([]float64, bdfVarKeep))
		beta2 := make([]float64, 1)
		seed2 := make([]float64, 1)
		alpha2 := bdf2Coefficients(p.t, p.tNext, p.tNm1, p.tNm2, []float64{vals[2]},
			[]float64{vals[1]}, []float64{vals[0]}, beta2, seed2)
		tol := func(v float64) float64 { return 1e-12 * math.Max(1, math.Abs(v)) }
		if math.Abs(alpha-alpha2) > tol(alpha2) {
			t.Fatalf("pattern %v: alpha = %.16g, bdf2 gives %.16g", p, alpha, alpha2)
		}
		if math.Abs(beta[0]-beta2[0]) > tol(beta2[0]) {
			t.Fatalf("pattern %v: beta = %.16g, bdf2 gives %.16g", p, beta[0], beta2[0])
		}
		if math.Abs(seed[0]-seed2[0]) > tol(seed2[0]) {
			t.Fatalf("pattern %v: seed = %.16g, bdf2 gives %.16g", p, seed[0], seed2[0])
		}
	}
}

// TestBDFVarMilneConstantMatchesBDF2 pins the variable-step Milne
// constant against the shipped bdf2Milne at order two.
func TestBDFVarMilneConstantMatchesBDF2(t *testing.T) {
	patterns := []struct{ tNext, t, tNm1, tNm2 float64 }{
		{3, 2, 1, 0},
		{1.3, 0.75, 0.4, -0.1},
		{5, 1, 0.5, -2},
	}
	for _, p := range patterns {
		hist := &bdfVarHistory{}
		for _, tt := range []float64{p.tNm2, p.tNm1, p.t} {
			hist.push(tt, []float64{0})
		}
		alpha := bdfVarCoefficients(2, p.tNext, hist, make([]float64, 1), make([]float64, 1),
			make([]float64, bdfVarKeep), make([]float64, bdfVarKeep), make([]float64, bdfVarKeep))
		_, tOldest := hist.back(2)
		got := 1 / (1 + alpha*(p.tNext-tOldest))
		want := bdf2Milne(p.t, p.tNext, p.tNm1, p.tNm2)
		if math.Abs(got-want) > 1e-14*math.Max(1, math.Abs(want)) {
			t.Fatalf("pattern %v: milne constant %.16g, bdf2Milne gives %.16g", p, got, want)
		}
	}
}

// TestIntegrateBDFVarFixedOrderLinear pins the exactness of the fixed
// orders on y' = k·t^(k−1), whose solution y = t^k only order k
// reproduces exactly: the k-step formula carries the k-th derivative
// the problem is built from, and any lower order drops it, so each
// locked run must land on 1 at the end AND report that it ran at the
// locked order, which together rule out a hook that silently
// integrates at order 1.
func TestIntegrateBDFVarFixedOrderLinear(t *testing.T) {
	for order := 1; order <= 5; order++ {
		k := float64(order)
		f := func(t float64, y *core.Array) (*core.Array, error) {
			out, err := core.Zeros(core.Float, 1)
			if err != nil {
				return nil, err
			}
			out.SetFloatAt(0, k*math.Pow(t, k-1))
			return out, nil
		}
		stats := &BDFVarStats{}
		end, err := integrateBDFVar("TestIntegrateBDFVarFixedOrderLinear", f, 0, 1,
			mustFloats(t, []float64{0}), BDFVarOptions{MaxSteps: 10000, Stats: stats}, order, true)
		if err != nil {
			t.Fatalf("locked order %d: %v", order, err)
		}
		if math.Abs(end[0]-1) > 5e-5 {
			t.Fatalf("locked order %d: y(1) = %.16g, want 1", order, end[0])
		}
		if stats.MaxOrder != order {
			t.Fatalf("locked order %d ran at max order %d", order, stats.MaxOrder)
		}
	}
}

// TestBDFVarExactPolynomialPerOrder drives the coefficient recurrence
// directly: a single order-k step from exact history on the degree-k
// polynomial p(t) = t^k must return p at the new time to rounding, on
// skewed steps, because the variable-step formula is exact for degree
// k when the past is exact.
func TestBDFVarExactPolynomialPerOrder(t *testing.T) {
	const tNext = 1.3
	// Back-value times on skewed step gaps, newest first.
	patterns := [][6]float64{
		{1, 0.7, 0.35, 0.1, -0.2, -1},
		{1, 0.9, 0.75, 0.5, 0.2, -0.1},
	}
	for order := 1; order <= 5; order++ {
		for _, g := range patterns {
			times := g[:order+1]
			hist := &bdfVarHistory{}
			for _, tt := range times {
				hist.push(tt, []float64{math.Pow(tt, float64(order))})
			}
			beta := make([]float64, 1)
			seed := make([]float64, 1)
			w := &odeWork{}
			alpha := bdfVarCoefficients(order, tNext, hist, beta, seed,
				make([]float64, bdfVarKeep), make([]float64, bdfVarKeep), make([]float64, bdfVarKeep))
			z := make([]float64, 1)
			err := odeNewton("TestBDFVarExactPolynomialPerOrder",
				func(t float64, y *core.Array) (*core.Array, error) {
					return core.FromFloats([]float64{float64(order) * math.Pow(t, float64(order-1))}, 1)
				}, w, tNext, alpha, 1, beta, seed, z, 1e-13, 1e-13)
			if err != nil {
				t.Fatalf("order %d gaps %v: odeNewton: %v", order, times, err)
			}
			want := math.Pow(tNext, float64(order))
			if math.Abs(z[0]-want) > 1e-11*math.Max(1, math.Abs(want)) {
				t.Fatalf("order %d gaps %v: z = %.16g, want %.16g to rounding", order, times, z[0], want)
			}
		}
	}
}

// TestIntegrateBDFVarErrors 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, and a
// nonsensical order cap is refused.
func TestIntegrateBDFVarErrors(t *testing.T) {
	y0 := mustFloats(t, []float64{1})
	same, err := integrateBDFVar("TestIntegrateBDFVarErrors", decay, 1, 1, y0, BDFVarOptions{}, 5, false)
	if err != nil {
		t.Fatalf("zero span: %v", err)
	}
	if math.Abs(same[0]-1) > 0 {
		t.Fatalf("zero span moved the state to %v", same[0])
	}
	wrongShape := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{1, 1}, 2)
	}
	if _, err := IntegrateBDFVar(wrongShape, 0, 1, y0, BDFVarOptions{}); err == nil {
		t.Fatal("expected an error when f returns the wrong shape")
	}
	matrixState := mustFloats(t, []float64{1, 1}, 1, 2)
	if _, err := IntegrateBDFVar(decay, 0, 1, matrixState, BDFVarOptions{}); err == nil {
		t.Fatal("expected an error for a rank-2 state")
	}
	if _, err := IntegrateBDFVar(decay, 0, 1, mustFloats(t, nil), BDFVarOptions{}); err == nil {
		t.Fatal("expected an error for an empty state")
	}
	if _, err := IntegrateBDFVar(decay, 0, 1, y0, BDFVarOptions{MaxSteps: 2}); err == nil {
		t.Fatal("expected an error for an exhausted step budget")
	}
	if _, err := integrateBDFVar("TestIntegrateBDFVarErrors", decay, 0, 1, y0, BDFVarOptions{}, 6, false); err == nil {
		t.Fatal("expected an error for an order cap above five")
	}
	if _, err := integrateBDFVar("TestIntegrateBDFVarErrors", decay, 0, 1, y0, BDFVarOptions{}, 0, false); err == nil {
		t.Fatal("expected an error for an order cap below one")
	}
	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 := IntegrateBDFVar(boom, 0, 1, y0, BDFVarOptions{}); err == nil {
		t.Fatal("expected the operator error to propagate")
	}
	// An f that survives the two probe evaluations and fails on the
	// starter's own evaluation is refused at once.
	calls := 0
	counted := func(t float64, y *core.Array) (*core.Array, error) {
		calls++
		if calls > 2 {
			return nil, base.Errf("detector tripped")
		}
		return core.FromFloats([]float64{0}, 1)
	}
	if _, err := IntegrateBDFVar(counted, 0, 1, mustFloats(t, []float64{1}), BDFVarOptions{}); err == nil {
		t.Fatal("expected the starter's f failure to surface")
	}
}