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

// heatMode builds sin(πx)·sin(πy) on an (n+2)×(n+2) grid over [0,1]²
// including the zero boundary ring, the lowest interior mode.
func heatMode(t *testing.T, n int) *core.Array {
	t.Helper()
	vals := make([]float64, (n+2)*(n+2))
	for r := range n + 2 {
		for c := range n + 2 {
			vals[r*(n+2)+c] = math.Sin(math.Pi*float64(c)/float64(n+1)) *
				math.Sin(math.Pi*float64(r)/float64(n+1))
		}
	}
	a, err := core.FromFloats(vals, n+2, n+2)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return a
}

// TestIntegrateHeat2DModeDecay checks the ADI solver on the lowest
// mode: with zero boundaries the amplitude decays like
// exp(−κ·2π²·t), and the scheme's O(dt²+h²) error must stay inside a
// one percent band on a 32-interior grid.
func TestIntegrateHeat2DModeDecay(t *testing.T) {
	const n = 32
	u0 := heatMode(t, n)
	const kappa, dt, tFinal = 1.0, 0.005, 0.1
	history, err := IntegrateHeat2D(u0, kappa, 1.0/float64(n+1), 1.0/float64(n+1),
		tFinal, dt, 2, 0, 0, 0, 0)
	if err != nil {
		t.Fatalf("IntegrateHeat2D: %v", err)
	}
	if history.Shape()[0] != 2 || history.Shape()[1] != n+2 {
		t.Fatalf("shape %v, want [%d %d %d]", history.Shape(), 2, n+2, n+2)
	}
	final := history.Shape()[0] - 1
	// The interior peak of the final state versus the exact decay.
	peak := 0.0
	for r := 1; r <= n; r++ {
		for c := 1; c <= n; c++ {
			if v := history.FloatAt(final*(n+2)*(n+2) + r*(n+2) + c); v > peak {
				peak = v
			}
		}
	}
	want := math.Exp(-kappa * 2 * math.Pi * math.Pi * tFinal)
	if math.Abs(peak-want) > 0.01 {
		t.Fatalf("final peak %.5g, want %.5g", peak, want)
	}
	// The boundary ring is held at zero.
	for c := range n + 2 {
		if history.FloatAt(final*(n+2)*(n+2)+c) != 0 ||
			history.FloatAt(final*(n+2)*(n+2)+(n+1)*(n+2)+c) != 0 {
			t.Fatalf("boundary ring moved at column %d", c)
		}
	}
}

// TestIntegrateWave2DStandingWave checks the explicit solver on the
// lowest standing mode with zero initial velocity against the closed
// form of the leapfrog itself. The mode is an exact eigenfunction of the
// five-point Laplacian, with eigenvalue mu, and the discrete
// characteristic of the scheme is cos theta = 1 - (c*h)²·mu/2, so after
// s steps the state is cos(s·theta)·u0. The run below takes 400 steps of
// period/400, so its amplitude is cos(400·theta) = -0.4155, not 1: the
// mode has not returned to its start at this time.
func TestIntegrateWave2DStandingWave(t *testing.T) {
	const n = 32
	u0 := heatMode(t, n) // sin(πx)sin(πy) with zero boundary ring
	v0 := core.New(core.Float, n+2, n+2)
	const c = 1.0
	dx := 1.0 / float64(n+1)
	period := math.Sqrt2 / (c * math.Pi)
	const steps = 400 // samples = 2 and dt = period/400 force this many steps
	dt := period / float64(steps)
	history, err := IntegrateWave2D(u0, v0, c, dx, dx, period, dt, 2)
	if err != nil {
		t.Fatalf("IntegrateWave2D: %v", err)
	}
	h := period / float64(steps)
	// The discrete eigenvalue of the (1,1) mode, mu = 8/dx²·sin²(π·dx/2)
	// on this square grid, and the phase 400 steps accumulate.
	sine := math.Sin(math.Pi * dx / 2)
	mu := 8 / (dx * dx) * sine * sine
	amp := math.Cos(float64(steps) * math.Acos(1-0.5*(c*h)*(c*h)*mu))
	final := history.Shape()[0] - 1
	worst := 0.0
	for r := range n + 2 {
		for cc := range n + 2 {
			i := final*(n+2)*(n+2) + r*(n+2) + cc
			if e := math.Abs(history.FloatAt(i) - amp*u0.FloatAt(r*(n+2)+cc)); e > worst {
				worst = e
			}
		}
	}
	if worst > 1e-11 {
		t.Fatalf("after %.0f steps the worst deviation from cos(%.6f)·u0 is %.4g, want the leapfrog characteristic",
			float64(steps), amp, worst)
	}
}

// TestIntegrateWave2DCFLRefusal checks the stability budget: a step
// past the CFL limit is an error, not a silent blow-up.
func TestIntegrateWave2DCFLRefusal(t *testing.T) {
	const n = 32
	u0 := heatMode(t, n)
	v0 := core.New(core.Float, n+2, n+2)
	dx := 1.0 / float64(n+1)
	// c·dt·sqrt(1/dx²+1/dy²) = 1·0.05·45.25 ≈ 2.26 > 1.
	if _, err := IntegrateWave2D(u0, v0, 1, dx, dx, 0.05, 0.05, 2); err == nil {
		t.Fatal("a CFL-violating step accepted")
	}
	if _, err := IntegrateWave2D(u0, core.New(core.Float, 3, 3), 1, dx, dx, 0.05, 0.001, 2); err == nil {
		t.Fatal("mismatched velocity grid accepted")
	}
}

// pde2dAnisoMode builds the (p, q) discrete Dirichlet eigenmode of the
// five-point Laplacian on a rows×cols grid: sin(π·p·c/(cols−1)) ·
// sin(π·q·r/(rows−1)), which vanishes on all four boundary lines.
func pde2dAnisoMode(t *testing.T, rows, cols, p, q int) *core.Array {
	t.Helper()
	vals := make([]float64, rows*cols)
	for r := range rows {
		for c := range cols {
			vals[r*cols+c] = math.Sin(math.Pi*float64(p)*float64(c)/float64(cols-1)) *
				math.Sin(math.Pi*float64(q)*float64(r)/float64(rows-1))
		}
	}
	a, err := core.FromFloats(vals, rows, cols)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return a
}

// anisoMu returns the dimensionless eigenvalues of the undivided second
// difference along each axis for the (p, q) mode above.
func anisoMu(rows, cols, p, q int) (mx, my float64) {
	return 4 * math.Pow(math.Sin(math.Pi*float64(p)/2/float64(cols-1)), 2),
		4 * math.Pow(math.Sin(math.Pi*float64(q)/2/float64(rows-1)), 2)
}

// TestIntegrateWave2DAnisotropicGrid pins the explicit solver on a grid
// whose spacings differ between the axes, the case the square-grid tests
// cannot see. The (1, 2) mode is an exact eigenfunction of the
// five-point Laplacian with eigenvalue λx + λy, where λx carries dx and
// λy carries dy, so the leapfrog state after s steps is cos(s·θ)·u0 with
// cos θ = 1 − (c·h)²·(λx + λy)/2. A stencil that divides the y
// neighbours by dx² instead of dy², or the reverse, moves those
// eigenvalues and the amplitude with them.
func TestIntegrateWave2DAnisotropicGrid(t *testing.T) {
	const (
		rows, cols = 26, 34
		dx, dy     = 0.02, 0.05
		c          = 1.0
		steps      = 60
	)
	if dx == dy {
		t.Fatal("the case needs spacings that differ between the axes")
	}
	mx, my := anisoMu(rows, cols, 1, 2)
	lx, ly := mx/(dx*dx), my/(dy*dy)
	dt := 0.5 / (c * math.Sqrt(1/(dx*dx)+1/(dy*dy))) // CFL = 1/2
	tFinal := dt * float64(steps)
	u0 := pde2dAnisoMode(t, rows, cols, 1, 2)
	v0 := core.New(core.Float, rows, cols)
	history, err := IntegrateWave2D(u0, v0, c, dx, dy, tFinal, dt, 2)
	if err != nil {
		t.Fatalf("IntegrateWave2D: %v", err)
	}
	h := tFinal / float64(steps)
	amp := math.Cos(float64(steps) * math.Acos(1-0.5*(c*h)*(c*h)*(lx+ly)))
	if math.Abs(amp) < 0.1 {
		t.Fatalf("the run decays to %.3g; the case needs an amplitude the comparison can see", amp)
	}
	final := history.Shape()[0] - 1
	worst := 0.0
	for i := range rows * cols {
		worst = math.Max(worst, math.Abs(history.FloatAt(final*rows*cols+i)-amp*u0.FloatAt(i)))
	}
	if worst > 1e-11 {
		t.Fatalf("after %d steps the worst deviation from cos(θ·%d)·u0 is %.4g (amplitude %.4g), want the leapfrog characteristic",
			steps, steps, worst, amp)
	}
}

// TestIntegrateHeat2DAnisotropicGrid pins the ADI solver the same way:
// on an eigenmode the two half steps compose into one amplification
// factor per step, (1 − rx·μx)(1 − ry·μy)/((1 + rx·μx)(1 + ry·μy)),
// with rx = κ·h/(2·dx²) and ry = κ·h/(2·dy²) against the dimensionless
// second-difference eigenvalues. Exchanging the two spacings moves the
// factor by orders of magnitude, so a mislabelled axis cannot pass.
func TestIntegrateHeat2DAnisotropicGrid(t *testing.T) {
	const (
		rows, cols = 26, 34
		dx, dy     = 0.02, 0.05
		kappa      = 1.0
		steps      = 10
	)
	if dx == dy {
		t.Fatal("the case needs spacings that differ between the axes")
	}
	mx, my := anisoMu(rows, cols, 1, 2)
	h := 0.0044
	tFinal := h * float64(steps)
	u0 := pde2dAnisoMode(t, rows, cols, 1, 2)
	history, err := IntegrateHeat2D(u0, kappa, dx, dy, tFinal, h, 2, 0, 0, 0, 0)
	if err != nil {
		t.Fatalf("IntegrateHeat2D: %v", err)
	}
	rx := kappa * h / (2 * dx * dx)
	ry := kappa * h / (2 * dy * dy)
	amp := math.Pow((1-rx*mx)*(1-ry*my)/((1+rx*mx)*(1+ry*my)), float64(steps))
	if math.Abs(amp) < 0.05 {
		t.Fatalf("the run decays to %.3g; the case needs an amplitude the comparison can see", amp)
	}
	final := history.Shape()[0] - 1
	worst := 0.0
	for i := range rows * cols {
		worst = math.Max(worst, math.Abs(history.FloatAt(final*rows*cols+i)-amp*u0.FloatAt(i)))
	}
	if worst > 1e-12 {
		t.Fatalf("after %d steps the worst deviation from the ADI amplification %.6g·u0 is %.4g, want the anisotropic factor",
			steps, amp, worst)
	}
}