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

// naiveDFT computes the transform quadratically, the reference the fast
// paths are verified against.
func naiveDFT(vals []complex128, sign float64) []complex128 {
	n := len(vals)
	out := make([]complex128, n)
	for k := range n {
		var s complex128
		for j, x := range vals {
			angle := sign * 2 * math.Pi * float64(j*k%n) / float64(n)
			s += x * complex(math.Cos(angle), math.Sin(angle))
		}
		out[k] = s
	}
	return out
}

func assertCplxApprox(t *testing.T, name string, got []complex128, want []complex128) {
	t.Helper()
	for i := range want {
		if math.Abs(real(got[i])-real(want[i])) > 1e-9 || math.Abs(imag(got[i])-imag(want[i])) > 1e-9 {
			t.Fatalf("%s[%d]: got %v, want %v", name, i, got[i], want[i])
		}
	}
}

func fftValues(t *testing.T, a *core.Array) []complex128 {
	t.Helper()
	out, err := FFT(a)
	if err != nil {
		t.Fatalf("FFT: %v", err)
	}
	vals := make([]complex128, out.Len())
	for i := range vals {
		vals[i], _ = core.ComplexAt(out, i)
	}
	return vals
}

func TestFFTKnownValues(t *testing.T) {
	// FFT of [1, 2, 3, 4] is the textbook [10, -2+2i, -2, -2-2i].
	a := mustFromInts(t, []int64{1, 2, 3, 4}, 4)
	got := fftValues(t, a)
	assertCplxApprox(t, "FFT", got, []complex128{10, complex(-2, 2), -2, complex(-2, -2)})

	// A constant input concentrates in the DC bin.
	c := mustFromFloats(t, []float64{2.5, 2.5, 2.5, 2.5}, 4)
	got = fftValues(t, c)
	assertCplxApprox(t, "FFT DC", got[:1], []complex128{10})
	for i := 1; i < 4; i++ {
		if math.Abs(real(got[i])) > 1e-9 || math.Abs(imag(got[i])) > 1e-9 {
			t.Fatalf("FFT constant input leaked into bin %d: %v", i, got[i])
		}
	}
}

func TestFFTAgainstNaive(t *testing.T) {
	// Powers of two take the radix-2 path; the odd and mixed lengths take
	// Bluestein. Both must agree with the quadratic reference.
	lengths := []int{1, 2, 4, 8, 16, 3, 5, 6, 7, 12, 13, 20}
	for _, n := range lengths {
		vals := make([]complex128, n)
		g := core.NewGenerator(int64(n))
		for i := range vals {
			f, _ := core.Floats(g, 2)
			re, _ := core.FloatAt(f, 0)
			im, _ := core.FloatAt(f, 1)
			vals[i] = complex(re*10-5, im*10-5)
		}
		src, err := core.FromComplexes(vals, n)
		if err != nil {
			t.Fatalf("FromComplexes(%d): %v", n, err)
		}
		assertCplxApprox(t, "FFT radix/bluestein", fftValues(t, src), naiveDFT(vals, -1))
	}
}

func TestIFFTRoundTrip(t *testing.T) {
	all := []float64{0.5, -1.25, 3, -0.5, 2, 1, -2, 4, 0.25, -3, 1.5, 2.75}
	for _, n := range []int{8, 5, 12} {
		src := mustFromFloats(t, all[:n], n)
		fwd, err := FFT(src)
		if err != nil {
			t.Fatalf("FFT(%d): %v", n, err)
		}
		back, err := IFFT(fwd)
		if err != nil {
			t.Fatalf("IFFT(%d): %v", n, err)
		}
		for i := range n {
			want, _ := core.FloatAt(src, i)
			got, _ := core.ComplexAt(back, i)
			if math.Abs(real(got)-want) > 1e-9 || math.Abs(imag(got)) > 1e-9 {
				t.Fatalf("IFFT(%d)[%d]: got %v, want %v", n, i, got, want)
			}
		}
	}
}

func TestFFTErrors(t *testing.T) {
	empty := mustFromInts(t, nil, 0)
	if _, err := FFT(empty); err == nil || !strings.Contains(err.Error(), "empty array") {
		t.Fatalf("FFT empty: %v", err)
	}
	if _, err := IFFT(empty); err == nil || !strings.Contains(err.Error(), "empty array") {
		t.Fatalf("IFFT empty: %v", err)
	}
	m := mustFromInts(t, []int64{1, 2, 3, 4}, 2, 2)
	if _, err := FFT(m); err == nil || !strings.Contains(err.Error(), "needs a 1-D array") {
		t.Fatalf("FFT 2-D: %v", err)
	}
}