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

// Regression pins for integrate/pde2d.go: the three-buffer leapfrog
// rotation, the pristine read of the Taylor start, the published-sample
// floor and schedule, and the complex-velocity guard. The oracles here
// are derived from the stencil itself, not from the library's own
// paths.

package integrate

import (
	"math"
	"strings"
	"testing"

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

// pde2dSteps re-derives the documented schedule contract independently
// of the implementation's pdeSchedule: ceil(tFinal/dt) steps, rounded up
// to a multiple of the sampling interval samples-1, all of size
// tFinal/steps, so every published time j·tFinal/(samples−1) is a step
// boundary. The returned every = steps/(samples−1) steps sit between two
// published samples.
func pde2dSteps(tFinal, dt float64, samples int) (steps, every int, h float64) {
	steps = max(int(math.Ceil(tFinal/dt)), samples-1)
	if rem := steps % (samples - 1); rem != 0 {
		steps += samples - 1 - rem
	}
	return steps, steps / (samples - 1), tFinal / float64(steps)
}

// pinSineMode builds sin(kx·π·x)·sin(ky·π·y) on a rows×cols grid whose
// ring is zero, sampled with x = c·dx and y = r·dy.
func pinSineMode(t *testing.T, rows, cols, kx, ky int) *core.Array {
	t.Helper()
	vals := make([]float64, rows*cols)
	for r := range rows {
		for c := range cols {
			vals[r*cols+c] = math.Sin(float64(kx)*math.Pi*float64(c)/float64(cols-1)) *
				math.Sin(float64(ky)*math.Pi*float64(r)/float64(rows-1))
		}
	}
	a, err := core.FromFloats(vals, rows, cols)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return a
}

// modeMu is the eigenvalue the five-point Laplacian gives that
// mode: Δmode = −μ·mode with μ = 4/dx²·sin²(kxπ/(2(cols−1))) +
// 4/dy²·sin²(kyπ/(2(rows−1))).
func modeMu(rows, cols, kx, ky int, dx, dy float64) float64 {
	sx := math.Sin(float64(kx) * math.Pi / (2 * float64(cols-1)))
	sy := math.Sin(float64(ky) * math.Pi / (2 * float64(rows-1)))
	return 4/(dx*dx)*sx*sx + 4/(dy*dy)*sy*sy
}

// point3x3 is the exact trajectory of the single interior point of
// a 3x3 grid with the zero ring: its four neighbours are all boundary
// values, so the five-point Laplacian is exactly −μ·u with
// μ = 2/dx² + 2/dy² and the leapfrog collapses to the scalar recurrence
// u_{s+1} = 2u_s − u_{s−1} + (c·h)²·(−μ·u_s), started by the Taylor step
// u¹ = u⁰ + h·v⁰ + (c·h)²/2·(−μ·u⁰). Its closed form is cos(s·θ) with
// cos θ = 1 − (c·h)²·μ/2.
func point3x3(steps int, h, c, dx, dy, u0, v0 float64) (hist []float64, theta float64) {
	mu := 2/(dx*dx) + 2/(dy*dy)
	lapW := (c * h) * (c * h)
	theta = math.Acos(1 - 0.5*lapW*mu)
	hist = make([]float64, steps+1)
	hist[0] = u0
	hist[1] = u0 + h*v0 + 0.5*lapW*(-mu*u0)
	for s := 2; s <= steps; s++ {
		hist[s] = 2*hist[s-1] - hist[s-2] + lapW*(-mu*hist[s-1])
	}
	return hist, theta
}

// singlePoint3x3 builds the 3x3 initial state whose one interior
// point carries u = 1 and whose ring is zero.
func singlePoint3x3(t *testing.T) *core.Array {
	t.Helper()
	vals := make([]float64, 9)
	vals[4] = 1
	a, err := core.FromFloats(vals, 3, 3)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return a
}

// TestWave2DLeapfrogScalarTrajectory pins the three-buffer rotation. On
// the 3x3 grid the interior point has no interior neighbour, so the
// whole trajectory IntegrateWave2D returns must equal the scalar
// leapfrog step by step. The old rotation (prev, cur = cur, next) left
// next aliased to cur after two swaps, which replaces the second-order
// recurrence with the first-order map u next = u + lapW*lap(u): step 201 came
// back as +0.7253864115 where the leapfrog gives −0.1460225271.
func TestWave2DLeapfrogScalarTrajectory(t *testing.T) {
	const dx, dy, c = 0.5, 0.5, 1.0
	const tFinal, dt, samples = 2.0, 0.01, 202
	steps, every, h := pde2dSteps(tFinal, dt, samples)
	if steps != 201 || every != 1 {
		t.Fatalf("schedule %d steps of %v, every %d, want 201 steps every 1", steps, h, every)
	}
	u0 := singlePoint3x3(t)
	v0 := core.New(core.Float, 3, 3)
	hist, err := IntegrateWave2D(u0, v0, c, dx, dy, tFinal, dt, samples)
	if err != nil {
		t.Fatalf("IntegrateWave2D: %v", err)
	}
	ref, theta := point3x3(steps, h, c, dx, dy, 1, 0)
	worst, worstAt := 0.0, 0
	for j := range samples {
		got := hist.FloatAt(j*9 + 4)
		if e := math.Abs(got - ref[j*every]); e > worst {
			worst, worstAt = e, j
		}
	}
	if worst > 1e-12 {
		t.Fatalf("sample %d = %.12f, want the scalar leapfrog %.12f (worst deviation %.3e over the whole history)",
			worstAt, hist.FloatAt(worstAt*9+4), ref[worstAt*every], worst)
	}
	// The recurrence itself is the closed form cos(s·θ), so the oracle is
	// pinned to the discrete characteristic and not to the implementation.
	worst, worstAt = 0.0, 0
	for s := range steps + 1 {
		want := math.Cos(float64(s) * theta)
		if e := math.Abs(ref[s] - want); e > worst {
			worst, worstAt = e, s
		}
	}
	if worst > 1e-11 {
		t.Fatalf("the reference recurrence deviates from cos(s·θ) by %.3e at step %d", worst, worstAt)
	}
}

// TestWave2DTaylorStartReadsUntouchedState pins the start-up read. The
// first step is the Taylor start, and it must read the untouched u⁰:
// computed into the same slice it reads (the old form wrote cur in place),
// the Laplacian at an interior point picks up already-updated left and
// upper neighbours. On this exact eigenmode with v⁰ = 0 the first
// published sample is cos θ·u⁰ to rounding, and the in-place start moved
// it by 1.4e-7 (relative to the unit amplitude at the centre).
func TestWave2DTaylorStartReadsUntouchedState(t *testing.T) {
	const n = 9
	const c = 1.0
	dx := 1.0 / float64(n-1)
	const tFinal, dt, samples = 0.008, 0.004, 3
	steps, every, h := pde2dSteps(tFinal, dt, samples)
	if steps != 2 || every != 1 {
		t.Fatalf("schedule %d steps of %v, every %d, want 2 steps every 1", steps, h, every)
	}
	u0 := pinSineMode(t, n, n, 1, 1)
	v0 := core.New(core.Float, n, n)
	hist, err := IntegrateWave2D(u0, v0, c, dx, dx, tFinal, dt, samples)
	if err != nil {
		t.Fatalf("IntegrateWave2D: %v", err)
	}
	mu := modeMu(n, n, 1, 1, dx, dx)
	cosTheta := 1 - 0.5*(c*h)*(c*h)*mu
	worst, worstAt := 0.0, 0
	for r := range n {
		for cc := range n {
			i := r*n + cc
			want := cosTheta * u0.FloatAt(i)
			if e := math.Abs(hist.FloatAt(n*n+i) - want); e > worst {
				worst, worstAt = e, i
			}
		}
	}
	if worst > 1e-13 {
		t.Fatalf("the Taylor start deviates from cos θ·u⁰ by %.3e at point %d (row %d, column %d): %.12f, want %.12f",
			worst, worstAt, worstAt/n, worstAt%n, hist.FloatAt(n*n+worstAt), cosTheta*u0.FloatAt(worstAt))
	}
}

// TestWave2DSamplesFloorTwo pins the floors of the published-sample
// argument. One sample leaves steps/(samples−1) with a zero divisor, so
// both 2-D entry points must refuse it the way the 1-D solvers do
// ("at least two samples are needed"), not panic with an integer divide
// by zero as the earlier revision did.
func TestWave2DSamplesFloorTwo(t *testing.T) {
	u0 := singlePoint3x3(t)
	v0 := core.New(core.Float, 3, 3)
	for _, samples := range []int{1, 0} {
		_, err := IntegrateHeat2D(u0, 1, 0.25, 0.25, 0.1, 0.01, samples, 0, 0, 0, 0)
		if err == nil {
			t.Fatalf("IntegrateHeat2D accepted samples = %d", samples)
		}
		if !strings.Contains(err.Error(), "at least two samples are needed") {
			t.Fatalf("IntegrateHeat2D samples = %d: error %q, want the at-least-two wording", samples, err)
		}
		_, err = IntegrateWave2D(u0, v0, 1, 0.25, 0.25, 0.1, 0.01, samples)
		if err == nil {
			t.Fatalf("IntegrateWave2D accepted samples = %d", samples)
		}
		if !strings.Contains(err.Error(), "at least two samples are needed") {
			t.Fatalf("IntegrateWave2D samples = %d: error %q, want the at-least-two wording", samples, err)
		}
	}
	// Two samples are still a valid call: the endpoints alone.
	if _, err := IntegrateWave2D(u0, v0, 1, 0.25, 0.25, 0.1, 0.01, 2); err != nil {
		t.Fatalf("IntegrateWave2D with samples = 2: %v", err)
	}
	if _, err := IntegrateHeat2D(u0, 1, 0.25, 0.25, 0.1, 0.01, 2, 0, 0, 0, 0); err != nil {
		t.Fatalf("IntegrateHeat2D with samples = 2: %v", err)
	}
}

// TestWave2DComplexVelocityRefusal pins the dtype guard on the initial
// velocity. The earlier revision checked only the rank and shape of v0,
// so a complex array reached v0.FloatAt and panicked inside the Taylor
// start ("index out of range [6] with length 0"); the 1-D solver refuses
// the same input with a message, which is the behaviour mirrored here.
func TestWave2DComplexVelocityRefusal(t *testing.T) {
	u0, err := core.FromFloats(make([]float64, 25), 5, 5)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	v0 := core.New(core.Complex, 5, 5)
	_, err = IntegrateWave2D(u0, v0, 1, 0.25, 0.25, 0.1, 0.01, 3)
	if err == nil {
		t.Fatal("IntegrateWave2D accepted a complex velocity array")
	}
	if !strings.Contains(err.Error(), "complex velocities are not supported") {
		t.Fatalf("IntegrateWave2D: error %q, want the unsupported-velocity wording", err)
	}
	// The same wording the 1-D wave solver uses for the same input.
	u1 := core.New(core.Float, 5)
	v1 := core.New(core.Complex, 5)
	_, err = IntegrateWave1D(u1, v1, 1, 0.25, 0.1, 0.01, 3)
	if err == nil || !strings.Contains(err.Error(), "complex velocities are not supported") {
		t.Fatalf("IntegrateWave1D: error %v, want the same refusal", err)
	}
}

// TestWave2DSlotsLandOnTheirTimes pins the published-sample schedule with
// an interval wider than one step. With samples = 5 over 8 steps the
// interval is every = 2, so slot j must hold the state after 2j steps, at
// the time 2j·h = j·tFinal/(samples−1). The earlier revision published
// the Taylor start (step 1) in slot 1 instead of step 2, shifting every
// slot from 1 to samples−2 off its published time.
func TestWave2DSlotsLandOnTheirTimes(t *testing.T) {
	const dx, dy, c = 0.5, 0.5, 1.0
	const tFinal, dt, samples = 0.08, 0.01, 5
	steps, every, h := pde2dSteps(tFinal, dt, samples)
	if steps != 8 || every != 2 {
		t.Fatalf("schedule %d steps of %v, every %d, want 8 steps every 2", steps, h, every)
	}
	u0 := singlePoint3x3(t)
	v0 := core.New(core.Float, 3, 3)
	hist, err := IntegrateWave2D(u0, v0, c, dx, dy, tFinal, dt, samples)
	if err != nil {
		t.Fatalf("IntegrateWave2D: %v", err)
	}
	ref, _ := point3x3(steps, h, c, dx, dy, 1, 0)
	for j := range samples {
		// The published time must be the step boundary j·every, which is
		// what makes the slots a uniform time grid with both endpoints.
		time, boundary := float64(j)*tFinal/float64(samples-1), float64(j*every)*h
		if e := math.Abs(time - boundary); e > 1e-15 {
			t.Fatalf("slot %d: published time %v is %d steps (%v) into the run, off by %.3e",
				j, time, j*every, boundary, e)
		}
		got, want := hist.FloatAt(j*9+4), ref[j*every]
		if e := math.Abs(got - want); e > 1e-12 {
			t.Fatalf("slot %d (t = %v) = %.12f, want the state after %d steps %.12f (off by %.3e)",
				j, time, got, j*every, want, e)
		}
	}
}

// TestWave2DRingHeldAtZero pins the documented ring contract against the
// buffer recycling of the rotation fix: the solver holds the boundary
// ring at zero, so a caller's own ring values may not reach the stencil
// at any published sample. A run whose state is all ones (ring 1) must
// therefore agree, from slot 1 on, with the run whose ring is zero, and
// the returned ring must be zero there. The earlier revision read the
// caller's ring into the first two Laplacians, and the three-buffer
// rotation makes that slip possible every third step unless the recycled
// buffer is cleared.
func TestWave2DRingHeldAtZero(t *testing.T) {
	const n = 5
	const c = 1.0
	dx := 1.0 / float64(n-1)
	const tFinal, dt, samples = 0.04, 0.01, 5
	steps, every, _ := pde2dSteps(tFinal, dt, samples)
	if steps != 4 || every != 1 {
		t.Fatalf("schedule %d steps, every %d, want 4 steps every 1", steps, every)
	}
	ones := make([]float64, n*n)
	for i := range ones {
		ones[i] = 1
	}
	cleared := make([]float64, n*n)
	copy(cleared, ones)
	for r := range n {
		cleared[r*n], cleared[r*n+n-1] = 0, 0
	}
	for cc := range n {
		cleared[cc], cleared[(n-1)*n+cc] = 0, 0
	}
	ringed, err := core.FromFloats(ones, n, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	bare, err := core.FromFloats(cleared, n, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	v0 := core.New(core.Float, n, n)
	got, err := IntegrateWave2D(ringed, v0, c, dx, dx, tFinal, dt, samples)
	if err != nil {
		t.Fatalf("IntegrateWave2D: %v", err)
	}
	want, err := IntegrateWave2D(bare, v0, c, dx, dx, tFinal, dt, samples)
	if err != nil {
		t.Fatalf("IntegrateWave2D: %v", err)
	}
	for j := 1; j < samples; j++ {
		for i := range n * n {
			r, cc := i/n, i%n
			g := got.FloatAt(j*n*n + i)
			if e := math.Abs(g - want.FloatAt(j*n*n+i)); e > 1e-15 {
				t.Fatalf("slot %d point (row %d, column %d) = %.16f, want %.16f: the caller's ring reached the stencil (off by %.3e)",
					j, r, cc, g, want.FloatAt(j*n*n+i), e)
			}
			if (r == 0 || r == n-1 || cc == 0 || cc == n-1) && g != 0 {
				t.Fatalf("slot %d point (row %d, column %d) = %v, want the ring held at zero", j, r, cc, g)
			}
		}
	}
}

// TestWave2DStandingModeMatchesDiscreteEigenvalue is the acceptance test
// for the leapfrog. The (1,1) sine mode of the 9x9 grid is an exact
// eigenfunction of the five-point Laplacian, so with zero initial
// velocity every published sample must be cos(ω·t_j)·u⁰, where the
// discrete characteristic of the scheme is
// cos(ω·h) = 1 − (c·h)²·μ/2 and μ is the eigenvalue derived above. The
// earlier revision returned +0.9250145127 at t = 1 where the discrete
// solution is −0.2935539474 (monotone decay instead of oscillation), and
// its in-place Taylor start missed the first sample by 1.4e-7.
func TestWave2DStandingModeMatchesDiscreteEigenvalue(t *testing.T) {
	const n = 9
	const c = 1.0
	dx := 1.0 / float64(n-1)
	const tFinal, dt, samples = 1.0, 0.004, 253
	steps, every, h := pde2dSteps(tFinal, dt, samples)
	if steps != 252 || every != 1 {
		t.Fatalf("schedule %d steps of %v, every %d, want 252 steps every 1", steps, h, every)
	}
	u0 := pinSineMode(t, n, n, 1, 1)
	v0 := core.New(core.Float, n, n)
	hist, err := IntegrateWave2D(u0, v0, c, dx, dx, tFinal, dt, samples)
	if err != nil {
		t.Fatalf("IntegrateWave2D: %v", err)
	}
	mu := modeMu(n, n, 1, 1, dx, dx)
	theta := math.Acos(1 - 0.5*(c*h)*(c*h)*mu) // ω = θ/h, and t_j = j·every·h
	worst, worstAt, worstSlot := 0.0, 0, 0
	for j := range samples {
		amp := math.Cos(float64(j*every) * theta)
		for i := range n * n {
			got := hist.FloatAt(j*n*n + i)
			if e := math.Abs(got - amp*u0.FloatAt(i)); e > worst {
				worst, worstAt, worstSlot = e, i, j
			}
		}
	}
	if worst > 1e-11 {
		t.Fatalf("sample %d (t = %v) point %d = %.12f, want cos(ω·t)·u⁰ = %.12f (worst %.3e over every published sample)",
			worstSlot, float64(worstSlot)*tFinal/float64(samples-1), worstAt,
			hist.FloatAt(worstSlot*n*n+worstAt),
			math.Cos(float64(worstSlot*every)*theta)*u0.FloatAt(worstAt), worst)
	}
}

// TestIntegrateRefusesInfiniteIntegrand pins the non-finite gate: an
// integrand returning +Inf used to poison the error sum with NaN,
// whose comparisons are false, and the integral came back as Inf with
// a nil error.
func TestIntegrateRefusesInfiniteIntegrand(t *testing.T) {
	inf := func(x float64) (float64, error) { return math.Inf(1), nil }
	if _, _, err := IntegrateFunction(inf, 0, 1, QuadratureOptions{}); err == nil {
		t.Fatal("IntegrateFunction with an infinite integrand returned no error")
	}
}