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

package signal

import (
	"math"
	"slices"
	"testing"

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

// Benchmarks for the kernels whose scheduling and per-element work was
// reworked: the spectral Poisson solves, the Bluestein route for a
// non-power-of-two length, the DCT/DST table path, the CWT wavelet
// build, the nonuniform transform and the rank filters. Every input is a
// fixed literal pattern, so the numbers are comparable between runs and
// machines.

// benchPerfFloats builds a deterministic float array of the requested
// rank from a fixed literal pattern.
func benchPerfFloats(b *testing.B, n int, shape ...int) *core.Array {
	b.Helper()
	v := make([]float64, n)
	for i := range v {
		v[i] = math.Sin(float64(i)*0.017) + 0.25*float64(i%29) - 3
	}
	a, err := core.FromFloats(v, shape...)
	if err != nil {
		b.Fatal(err)
	}
	return a
}

// poissonSource builds the right-hand side of a Poisson solve on a
// rows×cols grid: a smooth interior bump that vanishes on the boundary,
// the shape the Dirichlet solve is built for.
func poissonSource(b *testing.B, rows, cols int) *core.Array {
	b.Helper()
	v := make([]float64, rows*cols)
	for r := range rows {
		for c := range cols {
			v[r*cols+c] = math.Sin(math.Pi*float64(r)/float64(rows-1)) *
				math.Sin(2*math.Pi*float64(c)/float64(cols-1))
		}
	}
	a, err := core.FromFloats(v, rows, cols)
	if err != nil {
		b.Fatal(err)
	}
	return a
}

// BenchmarkSolvePoissonDirichlet measures the 512×512 sine-basis solve:
// the per-mode division by the stencil eigenvalues and the two
// separable DST-I passes. The grid is large enough for the eigenvalue
// loop to be a visible share of the total.
func BenchmarkSolvePoissonDirichlet(b *testing.B) {
	f := poissonSource(b, 512, 512)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := SolvePoissonDirichlet(f, 1, 1); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkSolvePoissonNeumann is the cosine-basis twin on a 256×256
// grid.
func BenchmarkSolvePoissonNeumann(b *testing.B) {
	f := poissonSource(b, 256, 256)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := SolvePoissonNeumann(f, 1, 1); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkFFTBluestein1000 measures a 1-D transform at a length that
// is not a power of two, the only route to Bluestein's chirp-z
// transform. The chirp and the kernel spectrum depend only on the
// length and the direction, which is what the plan cache holds.
func BenchmarkFFTBluestein1000(b *testing.B) {
	x := benchPerfFloats(b, 1000, 1000)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := FFT(x); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkFFTBluestein10000 is the same route at a length whose padded
// size is large enough for the kernel transform to dominate.
func BenchmarkFFTBluestein10000(b *testing.B) {
	x := benchPerfFloats(b, 10000, 10000)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := FFT(x); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkDCT2_1024 measures the type-2 cosine transform of a
// 1024-sample vector: the per-sample phase table, one padded inverse
// FFT and the per-bin rotation table.
func BenchmarkDCT2_1024(b *testing.B) {
	x := benchPerfFloats(b, 1024, 1024)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := DCT(x, 2); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkDST1_1023 measures the type-1 sine transform, whose phase
// table is a constant fill and whose bin rotations are the widest.
func BenchmarkDST1_1023(b *testing.B) {
	x := benchPerfFloats(b, 1023, 1023)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := DST(x, 1); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkCWTMorletScales16 measures sixteen Morlet scales over a
// 512-sample signal: the per-scale wavelet build (a cosine, a sine and
// an exponential per sample, with the time axis wrapped once per
// sample) and the FFT pair around it.
func BenchmarkCWTMorletScales16(b *testing.B) {
	x := benchPerfFloats(b, 512, 512)
	scales := make([]float64, 16)
	for i := range scales {
		scales[i] = 1 + float64(i)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := CWT(x, Morlet, scales, 0.25); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkNUFFTType1_4096 measures the type-1 nonuniform transform:
// 4096 samples spread onto an eight-times oversampled grid of 2048
// points, one inverse FFT and the per-bin deconvolution.
func BenchmarkNUFFTType1_4096(b *testing.B) {
	count := 4096
	xs := make([]float64, count)
	cs := make([]complex128, count)
	for j := range count {
		xs[j] = float64(j)/float64(count) - 0.5 + 1e-4
		cs[j] = complex(math.Sin(float64(j)*0.05), math.Cos(float64(j)*0.11))
	}
	x, err := core.FromFloats(xs, count)
	if err != nil {
		b.Fatal(err)
	}
	c, err := core.FromComplexes(cs, count)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := NUFFTType1(x, c, 512); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkMedianFilterWide31 measures the running median over a
// 32768-sample signal with a 31-sample window: every output is the
// middle order statistic of its own window.
func BenchmarkMedianFilterWide31(b *testing.B) {
	x := benchPerfFloats(b, 1<<15, 1<<15)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := MedianFilter(x, 31); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkRankFilterWide31_7 is the general-rank twin: the same window
// read at rank 7 rather than at the middle.
func BenchmarkRankFilterWide31_7(b *testing.B) {
	x := benchPerfFloats(b, 1<<15, 1<<15)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := RankFilter(x, 31, 7); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkMedianFilter2D measures the 3×3 running median over a
// 256×256 image, whose windows truncate at every edge.
func BenchmarkMedianFilter2D(b *testing.B) {
	img := benchPerfFloats(b, 256*256, 256, 256)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := MedianFilter2D(img, 3); err != nil {
			b.Fatal(err)
		}
	}
}

// rankFilterPerWindow is the per-window sort RankFilter used before the
// sliding mirror, kept here as the reference the current kernel is
// checked against bit for bit. A window is collected and sorted whole
// for every output point, and the rank reads the sorted slice.
func rankFilterPerWindow(x *core.Array, window, k int) *core.Array {
	n := x.Len()
	src := widenFloats(x)
	out := core.New(core.Float, n)
	dst := out.RawFloats()
	win := make([]float64, window)
	half := window / 2
	for i := range n {
		lo := max(i-half, 0)
		hi := min(i+half+1, n)
		count := hi - lo
		nan := false
		for j := lo; j < hi; j++ {
			v := src[j]
			if math.IsNaN(v) {
				nan = true
				break
			}
			win[j-lo] = v
		}
		if nan {
			dst[i] = math.NaN()
			continue
		}
		slices.Sort(win[:count])
		dst[i] = win[k*count/window]
	}
	return out
}

// TestRankFilterMirrorMatchesPerWindow pins the sliding mirror to the
// per-window sort bit for bit. The windows where the sort's own tie
// order decides the answer, a zero next to a −0 and any NaN, are the
// reason the mirror is dropped there; the patterns below hold ties,
// signed zeros, duplicates, NaNs, constant runs and ramps, over every
// signal length and window the gate admits.
func TestRankFilterMirrorMatchesPerWindow(t *testing.T) {
	seed := uint64(0x2545F4914F6CDD1D)
	next := func() uint64 {
		seed ^= seed << 13
		seed ^= seed >> 7
		seed ^= seed << 17
		return seed
	}
	fill := func(pattern string, n int) []float64 {
		vals := make([]float64, n)
		for i := range vals {
			u := next()
			switch pattern {
			case "random":
				vals[i] = float64(int64(u%2001)-1000) / 8
			case "duplicates":
				vals[i] = float64(u % 4)
			case "signed zeros":
				switch u % 5 {
				case 0:
					vals[i] = math.Copysign(0, -1)
				case 1:
					vals[i] = 0
				default:
					vals[i] = float64(u%3) - 1
				}
			case "NaN":
				vals[i] = float64(u%7) - 3
				if u%11 == 0 {
					vals[i] = math.NaN()
				}
			case "constant":
				vals[i] = 2.5
			default:
				vals[i] = float64(i)
			}
		}
		return vals
	}
	values := func(t *testing.T, vals []float64) *core.Array {
		t.Helper()
		a, err := core.FromFloats(vals, len(vals))
		if err != nil {
			t.Fatal(err)
		}
		return a
	}
	for _, pattern := range []string{"random", "duplicates", "signed zeros", "NaN", "constant", "ramp"} {
		for n := 1; n <= 48; n++ {
			x := values(t, fill(pattern, n))
			for window := 3; window <= n; window += 2 {
				for k := range window {
					got, err := RankFilter(x, window, k)
					if err != nil {
						t.Fatalf("%s n=%d window=%d k=%d: %v", pattern, n, window, k, err)
					}
					want := rankFilterPerWindow(x, window, k)
					g, w := got.RawFloats(), want.RawFloats()
					for i := range g {
						if math.Float64bits(g[i]) != math.Float64bits(w[i]) {
							t.Fatalf("%s n=%d window=%d k=%d sample %d: got %v want %v",
								pattern, n, window, k, i, g[i], w[i])
						}
					}
				}
			}
		}
	}
	// Long runs, so the mirror has slid far before the check.
	for _, pattern := range []string{"random", "duplicates", "signed zeros", "NaN"} {
		x := values(t, fill(pattern, 5000))
		for _, window := range []int{3, 5, 31, 997} {
			for _, k := range []int{0, 1, window / 2, window - 2, window - 1} {
				got, err := RankFilter(x, window, k)
				if err != nil {
					t.Fatal(err)
				}
				want := rankFilterPerWindow(x, window, k)
				g, w := got.RawFloats(), want.RawFloats()
				for i := range g {
					if math.Float64bits(g[i]) != math.Float64bits(w[i]) {
						t.Fatalf("%s window=%d k=%d sample %d: got %v want %v",
							pattern, window, k, i, g[i], w[i])
					}
				}
			}
		}
	}
}

// BenchmarkFFTRealInput4096 measures the 1-D forward transform of a
// real (float64) array: the payload is materialised rather than shared,
// so the entry point adds no defensive copy of its own.
func BenchmarkFFTRealInput4096(b *testing.B) {
	x := benchPerfFloats(b, 4096, 4096)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := FFT(x); err != nil {
			b.Fatal(err)
		}
	}
}