// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package integrate

import (
	"math"
	"testing"

	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// Regression tests for the integrators: the
// absolute cubature tolerance, the collapse floor of the ODE steppers
// and the sample times of the PDE evolutions.

// TestCubatureScalesWithMagnitude pins the stopping rule: a smooth
// integrand of large magnitude must converge, not exhaust the budget
// chasing an absolute bound below the rounding floor of the sum.
func TestCubatureScalesWithMagnitude(t *testing.T) {
	const want = 1e6 // ∫∫ 1e6 over the unit square
	got, err := IntegrateND(func([]float64) float64 { return want }, []float64{0, 0}, []float64{1, 1}, CubatureOptions{})
	if err != nil {
		t.Fatalf("IntegrateND of a constant: %v", err)
	}
	if math.Abs(got-want) > 1e-6*want {
		t.Fatalf("IntegrateND = %v, want %v", got, want)
	}
	// A small integral keeps its absolute accuracy.
	small, err := IntegrateND(func([]float64) float64 { return 1e-8 }, []float64{0, 0}, []float64{1, 1}, CubatureOptions{})
	if err != nil {
		t.Fatalf("IntegrateND of a small constant: %v", err)
	}
	if math.Abs(small-1e-8) > 1e-12 {
		t.Fatalf("IntegrateND = %v, want 1e-8", small)
	}
}

// TestODETinySpan pins the collapse rule: a span far below the absolute
// time scale is integrable, and the stepper must not refuse it.
func TestODETinySpan(t *testing.T) {
	zero := func(float64, *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{0}, 1)
	}
	y0, err := core.FromFloats([]float64{1}, 1)
	if err != nil {
		t.Fatal(err)
	}
	for _, span := range []float64{1e-13, 1e-15, 1e-20} {
		end, err := IntegrateODE(zero, 0, span, y0, ODEOptions{MaxSteps: 100000})
		if err != nil {
			t.Fatalf("span %g: %v", span, err)
		}
		if got := end.FloatAt(0); got != 1 {
			t.Fatalf("span %g: y = %v, want 1", span, got)
		}
	}
	// A real decay over a tiny span: y' = −y, y(1e-12) = exp(−1e-12).
	const span = 1e-12
	decay := func(_ float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{-y.FloatAt(0)}, 1)
	}
	end, err := IntegrateODE(decay, 0, span, y0, ODEOptions{MaxSteps: 100000})
	if err != nil {
		t.Fatalf("decay over %g: %v", span, err)
	}
	if got, want := end.FloatAt(0), math.Exp(-span); math.Abs(got-want) > 1e-12 {
		t.Fatalf("y = %.17g, want %.17g", got, want)
	}
}

// sineMode returns the interior grid of sin(π·x) on [0, 1] with n
// points, the eigenmode of the Dirichlet Laplacian.
func sineMode(t *testing.T, n int) *core.Array {
	t.Helper()
	u := make([]float64, n)
	for i := range n {
		x := float64(i+1) / float64(n+1)
		u[i] = math.Sin(math.Pi * x)
	}
	a, err := core.FromFloats(u, n)
	if err != nil {
		t.Fatal(err)
	}
	return a
}

// TestPDESchedule pins the step schedule directly, because the
// physics tests below pass with either schedule for a fine enough
// step: the published times must be j·tFinal/(samples−1) exactly, so
// the step count is a multiple of samples−1 and the last step lands on
// tFinal. `dt` is an upper bound, never a divisor to be honoured
// blindly.
func TestPDESchedule(t *testing.T) {
	cases := []struct {
		tFinal, dt float64
		samples    int
	}{
		{1, 0.3, 3}, {1, 0.3, 5}, {1, 0.7, 4}, {2.5, 0.4, 3}, {0.1, 0.3, 2}, {1, 1, 3}, {1, 0.01, 8},
	}
	for _, tc := range cases {
		steps, h := pdeSchedule(tc.tFinal, tc.dt, tc.samples)
		if steps <= 0 || h <= 0 {
			t.Fatalf("pdeSchedule(%v, %v, %d) = %d steps of %v", tc.tFinal, tc.dt, tc.samples, steps, h)
		}
		if float64(steps)*h != tc.tFinal {
			t.Errorf("pdeSchedule(%v, %v, %d): %d steps of %v reach %v",
				tc.tFinal, tc.dt, tc.samples, steps, h, float64(steps)*h)
		}
		if steps%(tc.samples-1) != 0 {
			t.Errorf("pdeSchedule(%v, %v, %d): %d steps do not divide by %d",
				tc.tFinal, tc.dt, tc.samples, steps, tc.samples-1)
		}
		if h > tc.dt {
			t.Errorf("pdeSchedule(%v, %v, %d): the step %v exceeds the bound %v",
				tc.tFinal, tc.dt, tc.samples, h, tc.dt)
		}
		every := steps / (tc.samples - 1)
		for j := range tc.samples {
			got := float64(j*every) * h
			want := tc.tFinal * float64(j) / float64(tc.samples-1)
			if math.Abs(got-want) > 1e-12*tc.tFinal {
				t.Errorf("sample %d at %v, want %v", j, got, want)
			}
		}
	}
}

// TestHeatSamplesLandOnTheirTimes checks the values at the published
// times; the schedule itself is pinned above.
// published states are the states at t = 0, tFinal/2 and tFinal, which
// the single sine mode turns into an exact decay ratio. The step is
// chosen for accuracy (r = κ·h/dx² ≈ 0.45) and does not divide tFinal,
// so the schedule has to round the count up.
func TestHeatSamplesLandOnTheirTimes(t *testing.T) {
	const (
		kappa = 1.0
		n     = 128
	)
	u0 := sineMode(t, n)
	dx := 1.0 / float64(n+1)
	dt := 0.9 * 0.5 * dx * dx
	got, err := IntegrateHeat1D(u0, kappa, dx, 1, dt, 3, 0, 0)
	if err != nil {
		t.Fatalf("IntegrateHeat1D: %v", err)
	}
	if s := got.Shape(); s[0] != 3 || s[1] != n {
		t.Fatalf("shape %v, want [3 %d]", s, n)
	}
	decay := func(tm float64) float64 { return math.Exp(-kappa * math.Pi * math.Pi * tm) }
	for j := range n {
		first := got.FloatAt(j)
		if math.Abs(first-u0.FloatAt(j)) > 1e-12 {
			t.Fatalf("sample 0 is not the initial state at %d", j)
		}
		mid := got.FloatAt(n + j)
		wantMid := first * decay(0.5)
		if rel := math.Abs(mid-wantMid) / wantMid; rel > 1e-3 {
			t.Fatalf("sample 1 at %d: %v, want %v (relative %.2g): the time is not 0.5", j, mid, wantMid, rel)
		}
		last := got.FloatAt(2*n + j)
		wantLast := first * decay(1.0)
		if rel := math.Abs(last-wantLast) / wantLast; rel > 1e-3 {
			t.Fatalf("sample 2 at %d: %v, want %v (relative %.2g): the time is not 1", j, last, wantLast, rel)
		}
	}
}

// TestWaveSamplesLandOnTheirTimes does the same for the wave equation
// (the schedule is pinned above):
// the standing mode is cos(π·c·t)·sin(π·x), so the sample at tFinal = 1
// with c = 1 is the initial state negated and the one at 0.5 is zero.
func TestWaveSamplesLandOnTheirTimes(t *testing.T) {
	const n = 128
	u0 := sineMode(t, n)
	v0, err := core.FromFloats(make([]float64, n), n)
	if err != nil {
		t.Fatal(err)
	}
	dx := 1.0 / float64(n+1)
	got, err := IntegrateWave1D(u0, v0, 1, dx, 1, 0.9*dx, 3)
	if err != nil {
		t.Fatalf("IntegrateWave1D: %v", err)
	}
	for j := range n {
		last := got.FloatAt(2*n + j)
		want := -u0.FloatAt(j) // cos(π·1) = −1
		if math.Abs(last-want) > 5e-3*math.Abs(want) {
			t.Fatalf("sample 2 at %d: %v, want %v: the last sample is not at t = 1", j, last, want)
		}
		mid := got.FloatAt(n + j)
		if math.Abs(mid) > 1e-2*math.Abs(u0.FloatAt(j)) {
			t.Fatalf("sample 1 at %d: %v, want near zero at t = 0.5", j, mid)
		}
	}
}