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

package signal

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

import (
	"math"
	"testing"
)

// poissonGrid samples a function of (x, y) on the periodic square
// [0, lx] × [0, ly] as a (rows × cols) array, matching the sampling
// SolvePoissonPeriodic documents.
func poissonGrid(t *testing.T, fn func(x, y float64) float64, rows, cols int, lx, ly float64) *core.Array {
	t.Helper()
	vals := make([]float64, rows*cols)
	for r := range rows {
		for c := range cols {
			vals[r*cols+c] = fn(float64(c)*lx/float64(cols), float64(r)*ly/float64(rows))
		}
	}
	return mustFloats(t, vals, rows, cols)
}

// poissonError returns the largest absolute difference between the
// solution and an exact function of (x, y) on the same grid.
func poissonError(u *core.Array, want func(x, y float64) float64, rows, cols int, lx, ly float64) float64 {
	worst := 0.0
	for r := range rows {
		for c := range cols {
			d := math.Abs(u.FloatAt(r*cols+c) - want(float64(c)*lx/float64(cols), float64(r)*ly/float64(rows)))
			worst = math.Max(worst, d)
		}
	}
	return worst
}

// TestSolvePoissonPeriodicModes checks the plane-wave answers the
// diagonalisation gets exactly: for u = sin x·sin y the Laplacian
// gives −2u, and for u = sin 3x·cos 2y it gives −13u. Both land on
// round-off.
func TestSolvePoissonPeriodicModes(t *testing.T) {
	const lx, ly = 2 * math.Pi, 2 * math.Pi
	cases := []struct {
		f, u func(x, y float64) float64
	}{
		{
			f: func(x, y float64) float64 { return 2 * math.Sin(x) * math.Sin(y) },
			u: func(x, y float64) float64 { return math.Sin(x) * math.Sin(y) },
		},
		{
			f: func(x, y float64) float64 { return 13 * math.Sin(3*x) * math.Cos(2*y) },
			u: func(x, y float64) float64 { return math.Sin(3*x) * math.Cos(2*y) },
		},
	}
	for k, tc := range cases {
		grid := poissonGrid(t, tc.f, 16, 16, lx, ly)
		u, err := SolvePoissonPeriodic(grid, lx, ly)
		if err != nil {
			t.Fatalf("case %d: SolvePoissonPeriodic: %v", k, err)
		}
		if u.Dtype() == core.Complex {
			t.Fatalf("case %d: the solution must be real", k)
		}
		if got := poissonError(u, tc.u, 16, 16, lx, ly); got > 1e-10 {
			t.Fatalf("case %d: max error %.3g, want round-off", k, got)
		}
	}
}

// TestSolvePoissonPeriodicSpectral shows the convergence is spectral:
// for u = sin x·e^{sin y}, whose Fourier coefficients decay faster
// than any power, refining the grid drives the error from visible to
// round-off in one refinement.
func TestSolvePoissonPeriodicSpectral(t *testing.T) {
	const lx, ly = 2 * math.Pi, 2 * math.Pi
	f := func(x, y float64) float64 {
		return math.Sin(x) * math.Exp(math.Sin(y)) * (math.Sin(y)*math.Sin(y) + math.Sin(y))
	}
	u := func(x, y float64) float64 { return math.Sin(x) * math.Exp(math.Sin(y)) }
	errAt := func(n int) float64 {
		sol, err := SolvePoissonPeriodic(poissonGrid(t, f, n, n, lx, ly), lx, ly)
		if err != nil {
			t.Fatalf("SolvePoissonPeriodic(%d): %v", n, err)
		}
		return poissonError(sol, u, n, n, lx, ly)
	}
	coarse, fine, finer := errAt(16), errAt(32), errAt(64)
	if coarse < 1e-9 {
		t.Skipf("coarse grid already at round-off (%v)", coarse)
	}
	if fine >= coarse {
		t.Fatalf("refining the grid did not help: %.3g then %.3g", coarse, fine)
	}
	if fine > 1e-11 || finer > 1e-11 {
		t.Fatalf("spectral accuracy not reached: %.3g, %.3g", fine, finer)
	}
}

// TestSolvePoissonPeriodicErrors pins the validation contract: the
// zero-mean compatibility condition, the shape and size of the grid,
// and positive side lengths.
func TestSolvePoissonPeriodicErrors(t *testing.T) {
	const lx, ly = 2 * math.Pi, 2 * math.Pi
	constant := poissonGrid(t, func(x, y float64) float64 { return 1 }, 8, 8, lx, ly)
	if _, err := SolvePoissonPeriodic(constant, lx, ly); err == nil {
		t.Fatal("expected an error for a right-hand side with nonzero mean")
	}
	rank1 := mustFloats(t, []float64{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}, 8)
	if _, err := SolvePoissonPeriodic(rank1, lx, ly); err == nil {
		t.Fatal("expected an error for a rank-1 right-hand side")
	}
	small := poissonGrid(t, func(x, y float64) float64 { return x }, 1, 8, lx, ly)
	if _, err := SolvePoissonPeriodic(small, lx, ly); err == nil {
		t.Fatal("expected an error for a grid dimension below 2")
	}
	grid := poissonGrid(t, func(x, y float64) float64 { return math.Sin(x) }, 8, 8, lx, ly)
	if _, err := SolvePoissonPeriodic(grid, 0, ly); err == nil {
		t.Fatal("expected an error for a zero side length")
	}
}

// TestSolvePoissonPeriodicDtypeContract pins the dtype contract: the
// solver computes in float64 end to end, so a float32, int or complex
// right-hand side is an error rather than a silently all-zero
// solution (RawFloats is nil for those dtypes).
func TestSolvePoissonPeriodicDtypeContract(t *testing.T) {
	const lx, ly = 2 * math.Pi, 2 * math.Pi
	f32, _ := core.FromFloat32s([]float32{1, 2, 3, 4, 5, 6, 7, 8, 8, 7, 6, 5, 4, 3, 2, 1}, 4, 4)
	if _, err := SolvePoissonPeriodic(f32, lx, ly); err == nil {
		t.Fatal("expected an error for a float32 right-hand side")
	}
	ints, _ := core.FromInts([]int64{1, 2, 3, 4, 5, 6, 7, 8, 8, 7, 6, 5, 4, 3, 2, 1}, 4, 4)
	if _, err := SolvePoissonPeriodic(ints, lx, ly); err == nil {
		t.Fatal("expected an error for an int right-hand side")
	}
	cplx, _ := core.FromComplexes(make([]complex128, 16), 4, 4)
	if _, err := SolvePoissonPeriodic(cplx, lx, ly); err == nil {
		t.Fatal("expected an error for a complex right-hand side")
	}
}