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

// advectGrid builds the interior cells of [0, 1] with n cells and the
// matching cell centres.
func advectGrid(n int) (dx float64, centres []float64) {
	dx = 1 / float64(n+1)
	centres = make([]float64, n)
	for i := range n {
		centres[i] = float64(i+1) * dx
	}
	return dx, centres
}

// advectL1 returns the L1 error of the final sample against exact.
func advectL1(t *testing.T, states *core.Array, n int, exact func(x float64) float64) float64 {
	t.Helper()
	last := (states.Shape()[0] - 1) * n
	_, centres := advectGrid(n)
	sum := 0.0
	for i := range n {
		sum += math.Abs(states.FloatAt(last+i) - exact(centres[i]))
	}
	return sum / float64(n)
}

// TestAdvectionUpwindMonotone pins the discrete maximum principle of
// both schemes: a monotone profile transported at CFL 0.9 stays
// monotone and inside its initial range, sample after sample.
func TestAdvectionUpwindMonotone(t *testing.T) {
	for _, limited := range []bool{false, true} {
		name := "upwind"
		if limited {
			name = "Koren"
		}
		n := 96
		dx, centres := advectGrid(n)
		u0 := make([]float64, n)
		for i := range n {
			u0[i] = 1 - centres[i]
		}
		u0Arr, err := core.FromFloats(u0, n)
		if err != nil {
			t.Fatal(err)
		}
		cfl := 0.9
		a := 1.0
		dt := cfl * dx
		run := func() (*core.Array, error) {
			if limited {
				return IntegrateAdvection1D(u0Arr, a, dx, 0.3, dt, 4, 1, 0)
			}
			return IntegrateUpwindAdvection1D(u0Arr, a, dx, 0.3, dt, 4, 1, 0)
		}
		states, err := run()
		if err != nil {
			t.Fatalf("%s: %v", name, err)
		}
		rows := states.Shape()[0]
		for r := range rows {
			prev := math.Inf(1)
			for i := range n {
				v := states.FloatAt(r*n + i)
				if v < -1e-12 || v > 1+1e-12 {
					t.Fatalf("%s: sample %d cell %d left the range [0, 1]: %g", name, r, i, v)
				}
				if v > prev+1e-12 {
					t.Fatalf("%s: sample %d stops being non-increasing at cell %d: %g after %g",
						name, r, i, v, prev)
				}
				prev = v
			}
		}
		// The transported ramp also keeps its shape: the last sample
		// tracks the shifted exact ramp to within the scheme's smear.
		last := (rows - 1) * n
		shift := 0.3
		for i := range n {
			x := centres[i] - shift
			want := 1.0
			if x > 0 {
				want = 1 - x
			}
			if d := math.Abs(states.FloatAt(last+i) - want); d > 0.02 {
				t.Fatalf("%s: ramp cell %d: %g, want %.6g", name, i, states.FloatAt(last+i), want)
			}
		}
	}
}

// TestAdvectionSquareWaveLimiterBeatsUpwind pins the limiter's reason
// for being: a square pulse carried at CFL 0.9 comes out visibly
// sharper under the Koren flux than under plain upwind, measured as
// the L1 error ratio against the exact shifted pulse.
func TestAdvectionSquareWaveLimiterBeatsUpwind(t *testing.T) {
	n := 256
	dx, centres := advectGrid(n)
	u0 := make([]float64, n)
	for i := range n {
		u0[i] = 0.0
		if centres[i] >= 0.3 && centres[i] <= 0.7 {
			u0[i] = 1.0
		}
	}
	u0Arr, err := core.FromFloats(u0, n)
	if err != nil {
		t.Fatal(err)
	}
	a := 1.0
	cfl := 0.9
	dt := cfl * dx
	const tFinal = 0.2
	upwind, err := IntegrateUpwindAdvection1D(u0Arr, a, dx, tFinal, dt, 2, 0, 0)
	if err != nil {
		t.Fatalf("upwind: %v", err)
	}
	koren, err := IntegrateAdvection1D(u0Arr, a, dx, tFinal, dt, 2, 0, 0)
	if err != nil {
		t.Fatalf("Koren: %v", err)
	}
	exact := func(x float64) float64 {
		if x >= 0.5 && x <= 0.9 {
			return 1.0
		}
		return 0.0
	}
	errUpwind := advectL1(t, upwind, n, exact)
	errKoren := advectL1(t, koren, n, exact)
	t.Logf("square pulse: upwind L1 %.4g, Koren L1 %.4g, ratio %.2f", errUpwind, errKoren, errUpwind/errKoren)
	if errUpwind/errKoren < 2.0 {
		t.Fatalf("limiter advantage %.2f, want at least 2x over upwind", errUpwind/errKoren)
	}
	// Both stay monotone in the sense that matters for a pulse: no
	// undershoot below the initial range.
	last := n
	for i := range n {
		for _, states := range []*core.Array{upwind, koren} {
			if v := states.FloatAt(last + i); v < -1e-12 || v > 1+1e-12 {
				t.Fatalf("scheme left the pulse range: %g at cell %d", v, i)
			}
		}
	}
}

// TestAdvectionSmoothTransportOrder pins the documented orders: at
// fixed CFL 0.9 the upwind L1 error halves as the grid halves (first
// order) and the Koren error improves second order or better (measured
// ratio past 3 at every pair), staying below the upwind error
// throughout. Three grid levels separate the two rates: a pair alone
// cannot tell a second-order Koren from a degraded one.
func TestAdvectionSmoothTransportOrder(t *testing.T) {
	a := 1.0
	cfl := 0.9
	const tFinal = 0.3
	gaussian := func(x float64) float64 { return math.Exp(-math.Pow((x-0.35)/0.1, 2)) }
	previousUpwind, previousKoren := 0.0, 0.0
	for _, n := range []int{100, 200, 400} {
		dx, centres := advectGrid(n)
		u0 := make([]float64, n)
		for i := range n {
			u0[i] = gaussian(centres[i])
		}
		u0Arr, err := core.FromFloats(u0, n)
		if err != nil {
			t.Fatal(err)
		}
		dt := cfl * dx / math.Abs(a)
		upwind, err := IntegrateUpwindAdvection1D(u0Arr, a, dx, tFinal, dt, 2, 0, 0)
		if err != nil {
			t.Fatalf("upwind n=%d: %v", n, err)
		}
		koren, err := IntegrateAdvection1D(u0Arr, a, dx, tFinal, dt, 2, 0, 0)
		if err != nil {
			t.Fatalf("Koren n=%d: %v", n, err)
		}
		shift := a * tFinal
		exact := func(x float64) float64 { return gaussian(x - shift) }
		eUpwind := advectL1(t, upwind, n, exact)
		eKoren := advectL1(t, koren, n, exact)
		t.Logf("n=%3d: upwind L1 %.3g, Koren L1 %.3g", n, eUpwind, eKoren)
		if eKoren > eUpwind {
			t.Fatalf("n=%d: Koren error %.3g above upwind %.3g", n, eKoren, eUpwind)
		}
		if previousUpwind > 0 {
			if r := previousUpwind / eUpwind; r < 1.5 || r > 3.0 {
				t.Fatalf("n=%d: upwind refinement ratio %.2f, want about 2", n, r)
			}
			if r := previousKoren / eKoren; r < 3.0 || r > 6.0 {
				t.Fatalf("n=%d: Koren refinement ratio %.2f, want the second-order rate the limiter carries", n, r)
			}
		}
		previousUpwind, previousKoren = eUpwind, eKoren
	}
}

// TestAdvectionDiffusionMatchesHeatWhenAZero pins the reduction: with
// a = 0 the advection-diffusion solver performs exactly the
// Crank-Nicolson steps of IntegrateHeat1D, so the two histories agree
// to the last bit.
func TestAdvectionDiffusionMatchesHeatWhenAZero(t *testing.T) {
	n := 32
	dx := 1 / float64(n+1)
	u0 := make([]float64, n)
	for i := range n {
		u0[i] = math.Sin(math.Pi * float64(i+1) * dx)
	}
	u0Arr, err := core.FromFloats(u0, n)
	if err != nil {
		t.Fatal(err)
	}
	heat, err := IntegrateHeat1D(u0Arr, 0.05, dx, 0.5, 0.004, 3, 0, 0)
	if err != nil {
		t.Fatalf("IntegrateHeat1D: %v", err)
	}
	adv, err := IntegrateAdvectionDiffusion1D(u0Arr, 0, 0.05, dx, 0.5, 0.004, 3, 0, 0)
	if err != nil {
		t.Fatalf("IntegrateAdvectionDiffusion1D: %v", err)
	}
	for i := range heat.Len() {
		if heat.FloatAt(i) != adv.FloatAt(i) {
			t.Fatalf("sample %d differs: heat %.17g, advection-diffusion %.17g",
				i, heat.FloatAt(i), adv.FloatAt(i))
		}
	}
}

// TestAdvectionDiffusionConvergence pins the documented orders of the
// combination: the split step is first order in time and second order
// in space, so the coupled refinement at fixed CFL shows the two
// mixed, an L1 error shrinking by roughly 2.5 to 3 per halving. The
// exact solution is the drifting heat kernel
// sqrt(w0/w)·exp(−(x−x0−at)²/w), w = w0 + 4Dt, whose boundary values
// are zero to well below the measured errors.
func TestAdvectionDiffusionConvergence(t *testing.T) {
	const (
		a      = 0.3
		probD  = 0.005
		x0     = 0.3
		tFinal = 0.6
	)
	exact := func(t, x float64) float64 {
		w := 0.01 + 4*probD*t
		return math.Sqrt(0.01/w) * math.Exp(-math.Pow(x-x0-a*t, 2)/w)
	}
	run := func(n int, cfl float64) float64 {
		dx := 1 / float64(n+1)
		u0 := make([]float64, n)
		for i := range n {
			u0[i] = exact(0, float64(i+1)*dx)
		}
		u0Arr, err := core.FromFloats(u0, n)
		if err != nil {
			t.Fatal(err)
		}
		states, err := IntegrateAdvectionDiffusion1D(u0Arr, a, probD, dx, tFinal, cfl*dx/a, 2, 0, 0)
		if err != nil {
			t.Fatalf("n=%d: %v", n, err)
		}
		return advectL1(t, states, n, func(x float64) float64 { return exact(tFinal, x) })
	}
	for _, cfl := range []float64{0.9, 0.3} {
		previous := 0.0
		for _, n := range []int{64, 128, 256} {
			e := run(n, cfl)
			t.Logf("CFL %.2f n=%3d: L1 %.4g", cfl, n, e)
			if previous > 0 {
				if r := previous / e; r < 2.2 || r > 3.7 {
					t.Fatalf("CFL %.2f: refinement ratio %.2f at n=%d, want the mixed time-space rate about 2.7",
						cfl, r, n)
				}
			}
			previous = e
		}
	}
}

func TestAdvectionCFLRefusal(t *testing.T) {
	n := 10
	dx := 1 / float64(n+1)
	u0, _ := core.FromFloats(make([]float64, n), n)
	// CFL = 1.5 for a = 1.
	dt := 1.5 * dx
	if _, err := IntegrateAdvection1D(u0, 1, dx, 0.1, dt, 2, 0, 0); err == nil || !stringsContains(err, "CFL violated") {
		t.Fatalf("Koren CFL violation: %v", err)
	}
	if _, err := IntegrateUpwindAdvection1D(u0, -1, dx, 0.1, dt, 2, 0, 0); err == nil || !stringsContains(err, "CFL violated") {
		t.Fatalf("upwind CFL violation: %v", err)
	}
	if _, err := IntegrateAdvectionDiffusion1D(u0, 1, 0.01, dx, 0.1, dt, 2, 0, 0); err == nil || !stringsContains(err, "CFL violated") {
		t.Fatalf("advection-diffusion CFL violation: %v", err)
	}
	// The boundary value at the CFL edge is accepted.
	if _, err := IntegrateAdvection1D(u0, 1, dx, 0.1, dx, 2, 0, 0); err != nil {
		t.Fatalf("CFL = 1 refused: %v", err)
	}
}

func TestAdvectionErrors(t *testing.T) {
	if _, err := IntegrateAdvection1D(mustFloats(t, []float64{1, 2, 3}, 3, 1), 1, 0.1, 1, 0.01, 2, 0, 0); err == nil || !stringsContains(err, "rank-1") {
		t.Fatalf("a rank-2 initial state: %v", err)
	}
	if _, err := IntegrateAdvection1D(mustFloats(t, []float64{1, 2, 3}), math.NaN(), 0.1, 1, 0.01, 2, 0, 0); err == nil || !stringsContains(err, "finite") {
		t.Fatalf("a NaN speed: %v", err)
	}
	if _, err := IntegrateAdvection1D(mustFloats(t, []float64{1, 2, 3}), 1, 0.1, 1, 0.01, 2, math.Inf(1), 0); err == nil || !stringsContains(err, "finite") {
		t.Fatalf("an infinite ghost value: %v", err)
	}
	if _, err := IntegrateUpwindAdvection1D(mustFloats(t, []float64{1, 2, 3}), 1, -0.1, 1, 0.01, 2, 0, 0); err == nil || !stringsContains(err, "positive") {
		t.Fatalf("a negative spacing: %v", err)
	}
	if _, err := IntegrateAdvection1D(mustFloats(t, []float64{1, 2, 3}), 1, 0.1, 1, 0.01, 1, 0, 0); err == nil || !stringsContains(err, "two samples") {
		t.Fatalf("one sample: %v", err)
	}
	if _, err := IntegrateAdvection1D(mustFloats(t, []float64{math.NaN()}), 1, 0.1, 1, 0.01, 2, 0, 0); err == nil || !stringsContains(err, "non-finite") {
		t.Fatalf("a NaN initial cell: %v", err)
	}
	if _, err := IntegrateAdvectionDiffusion1D(mustFloats(t, []float64{1, 2, 3}), 1, 0, 0.1, 1, 0.01, 2, 0, 0); err == nil || !stringsContains(err, "diffusivity") {
		t.Fatalf("zero diffusivity: %v", err)
	}
	// A single interior cell is refused by the limiter's stencil need
	// only for the limited boundary faces, so n = 1 must still run on
	// the upwind path with first order there.
	if _, err := IntegrateUpwindAdvection1D(mustFloats(t, []float64{0.5}), 1, 0.1, 0.05, 0.05, 2, 1, 0); err != nil {
		t.Fatalf("a single-cell upwind run: %v", err)
	}
}

// TestAdvectionLeftwardTransport pins the mirror branch: with a < 0
// the inflow is the right boundary, and both schemes must transport a
// monotone leftward ramp without new extrema, with the ghost value
// feeding in from the right.
func TestAdvectionLeftwardTransport(t *testing.T) {
	n := 96
	dx, centres := advectGrid(n)
	u0 := make([]float64, n)
	for i := range n {
		u0[i] = centres[i]
	}
	u0Arr, err := core.FromFloats(u0, n)
	if err != nil {
		t.Fatal(err)
	}
	a := -1.0
	dt := 0.9 * dx
	for _, limited := range []bool{false, true} {
		name := "upwind"
		run := func() (*core.Array, error) {
			if limited {
				name = "Koren"
				return IntegrateAdvection1D(u0Arr, a, dx, 0.2, dt, 3, 0, 1)
			}
			return IntegrateUpwindAdvection1D(u0Arr, a, dx, 0.2, dt, 3, 0, 1)
		}
		states, err := run()
		if err != nil {
			t.Fatalf("%s: %v", name, err)
		}
		for r := range states.Shape()[0] {
			for i := range n {
				v := states.FloatAt(r*n + i)
				if v < -1e-9 || v > 1+1e-9 {
					t.Fatalf("%s: sample %d cell %d left the range: %g", name, r, i, v)
				}
				if i > 0 && states.FloatAt(r*n+i) < states.FloatAt(r*n+i-1)-1e-9 {
					t.Fatalf("%s: sample %d grows a new extremum at cell %d: %g after %g", name, r, i, v, states.FloatAt(r*n+i-1))
				}
			}
		}
		// The exact ramp is x + 0.2, cut at the inflow value 1.
		last := (states.Shape()[0] - 1) * n
		for i := range n {
			want := math.Min(centres[i]+0.2, 1)
			if d := math.Abs(states.FloatAt(last+i) - want); d > 0.02 {
				t.Fatalf("%s: cell %d: %.6g, want %.6g", name, i, states.FloatAt(last+i), want)
			}
		}
	}
}