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

// resampleTone builds n samples of a·cos(2π·cycles·i/n) and its
// analytic continuation for comparisons.
func resampleTone(t *testing.T, a float64, cycles, n int) *core.Array {
	t.Helper()
	vals := make([]float64, n)
	for i := range n {
		vals[i] = a * math.Cos(2*math.Pi*float64(cycles)*float64(i)/float64(n))
	}
	out, err := core.FromFloats(vals, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return out
}

// TestDecimateTone checks that an in-band tone survives decimation
// with its amplitude and phase, sampled on the new grid. Output
// starts once the filter has full context, at input index
// ceil(delay/factor)·factor + delay.
func TestDecimateTone(t *testing.T) {
	const cycles, n, factor = 4, 256, 4
	out, err := Decimate(resampleTone(t, 1, cycles, n), factor, 0)
	if err != nil {
		t.Fatalf("Decimate: %v", err)
	}
	// taps = 32·factor+1 by default, delay = 16·factor, so the first
	// kept compensated index is 64 and sample i sits at input index
	// 64 + i·factor.
	const first = 64
	vals := out.RawFloats()[:out.Len()]
	for i, got := range vals {
		want := math.Cos(2 * math.Pi * float64(cycles) * float64(first+i*factor) / float64(n))
		if math.Abs(got-want) > 0.01 {
			t.Fatalf("sample %d = %.6g, want %.6g", i, got, want)
		}
	}
}

// TestDecimateAliasedTone checks the anti-alias job: a tone past the
// new Nyquist must come out suppressed, not folded down. The first
// and last taps samples carry filter transients and are skipped.
func TestDecimateAliasedTone(t *testing.T) {
	// 44 cycles over 256 samples sits past the new Nyquist of 32 when
	// decimating by 4; the stopband is around 80 dB, so a small leak
	// is honest, a fold-back to a visible tone is not.
	out, err := Decimate(resampleTone(t, 1, 44, 256), 4, 0)
	if err != nil {
		t.Fatalf("Decimate: %v", err)
	}
	vals := out.RawFloats()[:out.Len()]
	peak := 0.0
	for i, v := range vals {
		if i*4 < 128 || i*4 > 256-128 {
			continue
		}
		if a := math.Abs(v); a > peak {
			peak = a
		}
	}
	if peak > 0.02 {
		t.Fatalf("aliased tone survived at peak %.4g, the anti-alias filter leaked", peak)
	}
}

// TestResampleUpThenDown checks the rational path on a tone: up by 3
// lands on a finer grid with the tone intact, and the round trip back
// down recovers the input samples.
func TestResampleUpThenDown(t *testing.T) {
	const cycles, n = 5, 60
	in := resampleTone(t, 1, cycles, n)
	up, err := Resample(in, 3, 1, 0)
	if err != nil {
		t.Fatalf("Resample up: %v", err)
	}
	if up.Len() != 3*n {
		t.Fatalf("up-sampled length %d, want %d", up.Len(), 3*n)
	}
	vals := up.RawFloats()[:up.Len()]
	// Skip the edges, where the filter context is partial.
	for i := 24; i < up.Len()-24; i++ {
		want := math.Cos(2 * math.Pi * float64(cycles) * float64(i) / float64(3*n))
		if math.Abs(vals[i]-want) > 0.02 {
			t.Fatalf("up sample %d = %.6g, want %.6g", i, vals[i], want)
		}
	}
	down, err := Resample(up, 1, 3, 0)
	if err != nil {
		t.Fatalf("Resample down: %v", err)
	}
	if down.Len() != n {
		t.Fatalf("round-trip length %d, want %d", down.Len(), n)
	}
	inVals := in.RawFloats()[:n]
	dVals := down.RawFloats()[:down.Len()]
	for i := 12; i < n-12; i++ {
		if math.Abs(inVals[i]-dVals[i]) > 0.02 {
			t.Fatalf("round trip sample %d = %.6g, want %.6g", i, dVals[i], inVals[i])
		}
	}
}

// TestResampleFourierTone checks the exact band-limited resample: a
// tone stays the same amplitude on a doubled grid, point for point.
func TestResampleFourierTone(t *testing.T) {
	const cycles, n, size = 4, 64, 192
	out, err := ResampleFourier(resampleTone(t, 0.75, cycles, n), size)
	if err != nil {
		t.Fatalf("ResampleFourier: %v", err)
	}
	if out.Len() != size {
		t.Fatalf("length %d, want %d", out.Len(), size)
	}
	vals := out.RawFloats()[:out.Len()]
	for i, got := range vals {
		want := 0.75 * math.Cos(2*math.Pi*float64(cycles)*float64(i)/float64(size))
		if math.Abs(got-want) > 1e-12 {
			t.Fatalf("sample %d = %.12g, want %.12g", i, got, want)
		}
	}
}

// TestResampleFourierDown checks the truncating direction: a mixed
// two-tone series resampled to a third of its length keeps the low
// tone and drops the one past the new Nyquist.
func TestResampleFourierDown(t *testing.T) {
	const n = 96
	vals := make([]float64, n)
	for i := range n {
		low := math.Cos(2 * math.Pi * 3 * float64(i) / float64(n))
		high := 0.5 * math.Cos(2*math.Pi*30*float64(i)/float64(n))
		vals[i] = low + high
	}
	in, _ := core.FromFloats(vals, n)
	out, err := ResampleFourier(in, n/3)
	if err != nil {
		t.Fatalf("ResampleFourier: %v", err)
	}
	outVals := out.RawFloats()[:out.Len()]
	for i, got := range outVals {
		want := math.Cos(2 * math.Pi * 3 * float64(i) / float64(n/3))
		if math.Abs(got-want) > 1e-12 {
			t.Fatalf("sample %d = %.12g, want %.12g", i, got, want)
		}
	}
}

// TestResampleRefusals checks the shape and argument guards.
func TestResampleRefusals(t *testing.T) {
	bad := core.New(core.Float, 2, 2)
	if _, err := Decimate(bad, 2, 0); err == nil {
		t.Fatal("matrix accepted by Decimate")
	}
	if _, err := Resample(bad, 2, 1, 0); err == nil {
		t.Fatal("matrix accepted by Resample")
	}
	if _, err := ResampleFourier(bad, 8); err == nil {
		t.Fatal("matrix accepted by ResampleFourier")
	}
	one := core.New(core.Float, 16)
	if _, err := Decimate(one, 1, 0); err == nil {
		t.Fatal("identity factor accepted by Decimate")
	}
	if _, err := Resample(one, 1, 1, 0); err == nil {
		t.Fatal("identity rate accepted by Resample")
	}
	if _, err := Resample(one, 3, 1, 200); err == nil {
		t.Fatal("taps larger than the series accepted by Resample")
	}
	if _, err := ResampleFourier(one, 0); err == nil {
		t.Fatal("zero size accepted by ResampleFourier")
	}
}