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

package signal

import (
	"math"
	"testing"

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

// mustF builds a float array or fails the test.
func mustF(t *testing.T, vals []float64, shape ...int) *core.Array {
	t.Helper()
	a, err := core.FromFloats(vals, shape...)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return a
}

// TestFFTRejectsComplexMatrices pins the rank gate on the
// complex path, which used to flatten a matrix silently where the
// real path refused.
func TestFFTRejectsComplexMatrices(t *testing.T) {
	c, err := core.FromComplexes([]complex128{1, 2, 3, 4}, 2, 2)
	if err != nil {
		t.Fatalf("FromComplexes: %v", err)
	}
	if _, err := FFT(c); err == nil {
		t.Fatal("FFT: expected a rank error for a complex matrix")
	}
	if _, err := IFFT(c); err == nil {
		t.Fatal("IFFT: expected a rank error for a complex matrix")
	}
}

// TestNUFFTRejectsNaNCoordinate pins the named refusal: NaN
// defeats both range comparisons and used to poison the whole grid.
func TestNUFFTRejectsNaNCoordinate(t *testing.T) {
	x := mustF(t, []float64{math.NaN(), 0.1}, 2)
	c := mustF(t, []float64{1, 1}, 2)
	if _, err := NUFFTType1(x, c, 4); err == nil {
		t.Fatal("NUFFTType1: expected an error for a NaN coordinate")
	}
}

// TestInfiniteParameterGates pins that a +Inf parameter cannot
// pass a positive-only gate and turn the answer into NaN or silent
// zeros.
func TestInfiniteParameterGates(t *testing.T) {
	times := mustF(t, []float64{0.1, 0.4, 0.7, 1.0, 1.3}, 5)
	vals := mustF(t, []float64{1, -0.5, 0.8, -0.2, 0.6}, 5)
	if _, _, err := LombScargle(times, vals, math.Inf(1), math.Inf(1), 4); err == nil {
		t.Fatal("LombScargle: expected an error for an infinite frequency range")
	}
	x := mustF(t, make([]float64, 16), 16)
	for i := range 16 {
		x.SetFloatAt(i, math.Sin(float64(i)/3))
	}
	if _, err := CWT(x, Morlet, []float64{1}, math.Inf(1)); err == nil {
		t.Fatal("CWT: expected an error for an infinite sample spacing")
	}
	if _, err := CWT(x, Morlet, []float64{math.Inf(1)}, 0.1); err == nil {
		t.Fatal("CWT: expected an error for an infinite scale")
	}
	if _, _, err := WelchPSD(x, math.Inf(1), 8, 0, "hann"); err == nil {
		t.Fatal("WelchPSD: expected an error for an infinite fs")
	}
	if _, err := Spectrogram(x, math.Inf(1), STFTOptions{Segment: 8}); err == nil {
		t.Fatal("Spectrogram: expected an error for an infinite fs")
	}
	if _, _, err := ButterworthLowPass(2, math.Inf(1), 100); err == nil {
		t.Fatal("ButterworthLowPass: expected an error for an infinite fs")
	}
	if _, _, err := ChebyshevLowPass(2, math.Inf(1), 100, 1); err == nil {
		t.Fatal("ChebyshevLowPass: expected an error for an infinite fs")
	}
}

// TestPoissonDirichletMinimumGrid pins the documented 3×3
// minimum, whose length-one interior used to hit the DST-I floor.
func TestPoissonDirichletMinimumGrid(t *testing.T) {
	f := mustF(t, []float64{0, 0, 0, 0, 1, 0, 0, 0, 0}, 3, 3)
	u, err := SolvePoissonDirichlet(f, 1, 1)
	if err != nil {
		t.Fatalf("SolvePoissonDirichlet on the 3×3 minimum: %v", err)
	}
	// The single interior unknown solves −(4u)/h² = 1 at h = 1/2:
	// u = -1/16... the sign follows the convention −Δu = f, so the
	// interior value is f/16 with the operator's sign.
	if v := u.FloatAt(4); math.IsNaN(v) || math.IsInf(v, 0) {
		t.Fatalf("interior value = %v", v)
	}
	// (3, k) and (k, 3) grids take the one-point DST on one axis only.
	f2 := mustF(t, []float64{0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0}, 3, 4)
	if _, err := SolvePoissonDirichlet(f2, 1, 1); err != nil {
		t.Fatalf("SolvePoissonDirichlet on 3×4: %v", err)
	}
}

// TestMaxPoolNaNEveryRank pins the NaN-propagation promise on
// the four max-pool entry points the earlier pin left uncovered.
func TestMaxPoolNaNEveryRank(t *testing.T) {
	x1 := mustF(t, []float64{1, math.NaN(), 3, 4, 5, 6}, 1, 1, 6)
	got1, err := MaxPool1D(x1, 2, 1, 0)
	if err != nil {
		t.Fatalf("MaxPool1D: %v", err)
	}
	if !math.IsNaN(got1.FloatAt(0)) {
		t.Fatalf("MaxPool1D[0] = %v, want NaN", got1.FloatAt(0))
	}
	a1, err := AdaptiveMaxPool1D(x1, 2)
	if err != nil {
		t.Fatalf("AdaptiveMaxPool1D: %v", err)
	}
	if !math.IsNaN(a1.FloatAt(0)) {
		t.Fatalf("AdaptiveMaxPool1D[0] = %v, want NaN", a1.FloatAt(0))
	}
	g1, err := GlobalMaxPool1D(x1)
	if err != nil {
		t.Fatalf("GlobalMaxPool1D: %v", err)
	}
	if !math.IsNaN(g1.FloatAt(0)) {
		t.Fatalf("GlobalMaxPool1D[0] = %v, want NaN", g1.FloatAt(0))
	}
	x3 := make([]float64, 2*3*4*4*4)
	for i := range x3 {
		x3[i] = float64(i % 7)
	}
	x3[40] = math.NaN()
	c3, err := core.FromFloats(x3, 2, 3, 4, 4, 4)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	g3, err := GlobalMaxPool3D(c3)
	if err != nil {
		t.Fatalf("GlobalMaxPool3D: %v", err)
	}
	if !math.IsNaN(g3.FloatAt(0)) {
		t.Fatalf("GlobalMaxPool3D[0] = %v, want NaN", g3.FloatAt(0))
	}
	m3, err := MaxPool3D(c3, [3]int{2, 2, 2}, [3]int{1, 1, 1}, [3]int{0, 0, 0})
	if err != nil {
		t.Fatalf("MaxPool3D: %v", err)
	}
	// Flat 40 is (d, h, w) = (2, 2, 0); the window that covers it
	// starts at (1, 1, 0), output index 12 of the 3×3×3 grid.
	if !math.IsNaN(m3.FloatAt(12)) {
		t.Fatalf("MaxPool3D[12] = %v, want NaN", m3.FloatAt(12))
	}
	a3, err := AdaptiveMaxPool3D(c3, 1, 1, 1)
	if err != nil {
		t.Fatalf("AdaptiveMaxPool3D: %v", err)
	}
	if !math.IsNaN(a3.FloatAt(0)) {
		t.Fatalf("AdaptiveMaxPool3D[0] = %v, want NaN", a3.FloatAt(0))
	}
}

// TestGlobalPoolsRejectEmptySpatial pins the named refusal the
// adaptive ranks already make.
func TestGlobalPoolsRejectEmptySpatial(t *testing.T) {
	empty1, err := core.Zeros(core.Float, 2, 3, 0)
	if err != nil {
		t.Fatalf("Zeros: %v", err)
	}
	if _, err := GlobalMaxPool1D(empty1); err == nil {
		t.Fatal("GlobalMaxPool1D: expected an error for an empty spatial dimension")
	}
	if _, err := GlobalAvgPool1D(empty1); err == nil {
		t.Fatal("GlobalAvgPool1D: expected an error for an empty spatial dimension")
	}
	empty3, err := core.Zeros(core.Float, 2, 3, 0, 4, 4)
	if err != nil {
		t.Fatalf("Zeros: %v", err)
	}
	if _, err := GlobalMaxPool3D(empty3); err == nil {
		t.Fatal("GlobalMaxPool3D: expected an error for an empty spatial dimension")
	}
	if _, err := GlobalAvgPool3D(empty3); err == nil {
		t.Fatal("GlobalAvgPool3D: expected an error for an empty spatial dimension")
	}
}

// TestResampleMatchesDirectSum pins the polyphase window
// restriction: the bounded loop must produce bit-identical output to
// the full scan it replaced.
func TestResampleMatchesDirectSum(t *testing.T) {
	n := 61
	src := make([]float64, n)
	for i := range n {
		src[i] = math.Sin(float64(i)*0.37) + 0.2*math.Cos(float64(i)*1.7)
	}
	x := mustF(t, src, n)
	for _, c := range []struct{ up, down int }{{2, 3}, {3, 2}, {1, 2}, {2, 1}, {5, 7}} {
		got, err := Resample(x, c.up, c.down, 13)
		if err != nil {
			t.Fatalf("Resample(%d/%d): %v", c.up, c.down, err)
		}
		// The reference: the original full scan over every input
		// index, computing the same taps by hand.
		taps := 13
		fc := 0.5 * math.Min(1/float64(c.up), 1/float64(c.down))
		h := kaiserSinc(taps, fc)
		delay := (taps - 1) / 2
		outLen := (n*c.up + c.down - 1) / c.down
		if got.Len() != outLen {
			t.Fatalf("Resample(%d/%d): length %d, want %d", c.up, c.down, got.Len(), outLen)
		}
		for m := range outLen {
			centre := delay + m*c.down
			total := 0.0
			for k := range n {
				j := centre - k*c.up
				if j < 0 || j >= taps {
					continue
				}
				total += h[j] * src[k]
			}
			if want := float64(c.up) * total; got.FloatAt(m) != want {
				t.Fatalf("Resample(%d/%d)[%d] = %g, want %g", c.up, c.down, m, got.FloatAt(m), want)
			}
		}
	}
}