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

package signal

import (
	"math"
	"strings"
	"testing"

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

// filtfiltTestFilter is the shared design of the pins: a fourth-order
// Butterworth low-pass at a fifth of the sampling rate, stable and
// mildly selective.
func filtfiltTestFilter(t *testing.T) (b, a []float64) {
	t.Helper()
	b, a, err := ButterworthLowPass(4, 100, 20)
	if err != nil {
		t.Fatalf("ButterworthLowPass: %v", err)
	}
	return b, a
}

// TestFiltfiltZeroPhase pins the defining property from two sides.
// First, the response to an interior impulse is exactly symmetric
// about it, which is the time-domain face of zero phase. Second, a
// passband sinusoid comes out with no phase shift against its input,
// the frequency-domain face, where the one-pass filter visibly lags.
func TestFiltfiltZeroPhase(t *testing.T) {
	b, a := filtfiltTestFilter(t)
	const (
		n = 501
		m = 250
	)
	vals := make([]float64, n)
	vals[m] = 1
	got, err := Filtfilt(b, a, mustFloats(t, vals))
	if err != nil {
		t.Fatalf("Filtfilt: %v", err)
	}
	// The pad regions (twice the filter's pad at each end, once per
	// pass) host the seam transients; everything inside stays clean.
	const clean = 100
	for d := 1; d <= clean; d++ {
		left := got.FloatAt(m - d)
		right := got.FloatAt(m + d)
		if math.Abs(left-right) > 1e-12*math.Max(1, math.Abs(left)) {
			t.Fatalf("impulse response asymmetric at offset %d: %g vs %g", d, left, right)
		}
	}
	// Phase of a passband sinusoid: fit A·sin + B·cos over the
	// interior and read the phase back. Zero phase means B ≈ 0.
	// wTone is the angular frequency of the 5 hertz test tone at
	// 100 hertz.
	const wTone = 2 * math.Pi * 5.0 / 100.0
	fitPhase := func(y *core.Array, from, to int) (amp, phase float64) {
		var ss, sc, cc, ys, yc float64
		for i := from; i < to; i++ {
			s, c := math.Sincos(wTone * float64(i))
			v := y.FloatAt(i)
			ss += s * s
			sc += s * c
			cc += c * c
			ys += v * s
			yc += v * c
		}
		a1 := (cc*ys - sc*yc) / (ss*cc - sc*sc)
		b1 := (ss*yc - sc*ys) / (ss*cc - sc*sc)
		return math.Hypot(a1, b1), math.Atan2(b1, a1)
	}
	sine := make([]float64, n)
	for i := range sine {
		sine[i] = math.Sin(wTone * float64(i))
	}
	out, err := Filtfilt(b, a, mustFloats(t, sine))
	if err != nil {
		t.Fatalf("Filtfilt: %v", err)
	}
	amp, phase := fitPhase(out, clean, n-clean)
	if math.Abs(phase) > 1e-9 {
		t.Fatalf("passband phase %g rad, want 0", phase)
	}
	wantAmp := cmag(freqResponse(b, a, wTone))
	wantAmp *= wantAmp
	if math.Abs(amp-wantAmp) > 1e-9 {
		t.Fatalf("passband amplitude %.10f, want |H|^2 = %.10f", amp, wantAmp)
	}
	// The one-pass contrast: same fit shows the causal group delay.
	one, err := FilterApply(b, a, mustFloats(t, sine))
	if err != nil {
		t.Fatalf("FilterApply: %v", err)
	}
	_, onePhase := fitPhase(one, clean, n-clean)
	if math.Abs(onePhase) < 0.1 {
		t.Fatalf("one-pass phase %g rad should be plainly nonzero", onePhase)
	}
}

// TestFiltfiltMatchesManualComposition pins the architecture: the
// documented pad, two FilterApply passes with a time-reversal between
// them, and the crop reproduce Filtfilt sample for sample, bit for
// bit.
func TestFiltfiltMatchesManualComposition(t *testing.T) {
	b, a := filtfiltTestFilter(t)
	g := core.NewGenerator(11)
	const n = 300
	vals := make([]float64, n)
	for i := range vals {
		vals[i] = g.NormalUnit()
	}
	x := mustFloats(t, vals)
	got, err := Filtfilt(b, a, x)
	if err != nil {
		t.Fatalf("Filtfilt: %v", err)
	}
	// The pad of the documentation, built by hand.
	nfilt := max(len(b), len(a))
	pad := 3 * (nfilt - 1)
	p := n + 2*pad
	padded := make([]float64, p)
	copy(padded[pad:pad+n], vals)
	for k := range pad {
		padded[k] = vals[pad-k]
		padded[pad+n+k] = vals[n-2-k]
	}
	y1, err := FilterApply(b, a, mustFloats(t, padded))
	if err != nil {
		t.Fatalf("FilterApply forward: %v", err)
	}
	rev := make([]float64, p)
	for i := range p {
		rev[i] = y1.FloatAt(p - 1 - i)
	}
	y2, err := FilterApply(b, a, mustFloats(t, rev))
	if err != nil {
		t.Fatalf("FilterApply backward: %v", err)
	}
	for ts := range n {
		want := y2.FloatAt(p - 1 - pad - ts)
		if got.FloatAt(ts) != want {
			t.Fatalf("sample %d: %v, want %v", ts, got.FloatAt(ts), want)
		}
	}
}

// TestFiltfiltSuppressesMoreThanOnePass pins the doubled magnitude
// response: in the stopband the filtfilt attenuation is the square of
// the one-pass attenuation at the same frequency.
func TestFiltfiltSuppressesMoreThanOnePass(t *testing.T) {
	b, a := filtfiltTestFilter(t)
	const fs = 100.0
	const fStop = 40.0
	wStop := 2 * math.Pi * fStop / fs
	const n = 800
	vals := make([]float64, n)
	for i := range vals {
		vals[i] = math.Sin(wStop * float64(i))
	}
	x := mustFloats(t, vals)
	two, err := Filtfilt(b, a, x)
	if err != nil {
		t.Fatalf("Filtfilt: %v", err)
	}
	one, err := FilterApply(b, a, x)
	if err != nil {
		t.Fatalf("FilterApply: %v", err)
	}
	rms := func(y *core.Array) float64 {
		var s float64
		for i := 200; i < n-200; i++ {
			s += y.FloatAt(i) * y.FloatAt(i)
		}
		// Whole periods sit in the window, so the RMS of a unit
		// sinusoid of amplitude A is A/√2 and the factor below
		// returns the amplitude.
		return math.Sqrt2 * math.Sqrt(s/(n-400))
	}
	gotTwo := rms(two)
	gotOne := rms(one)
	wantTwo := cmag(freqResponse(b, a, wStop))
	wantTwo *= wantTwo
	if math.Abs(gotTwo-wantTwo) > 0.05*wantTwo+1e-9 {
		t.Fatalf("filtfilt stopband amplitude %.3g, want |H|^2 = %.3g", gotTwo, wantTwo)
	}
	wantOne := cmag(freqResponse(b, a, wStop))
	if math.Abs(gotOne-wantOne) > 0.05*wantOne+1e-9 {
		t.Fatalf("one-pass stopband amplitude %.3g, want |H| = %.3g", gotOne, wantOne)
	}
	if gotTwo >= gotOne {
		t.Fatalf("filtfilt %.3g not quieter than one pass %.3g", gotTwo, gotOne)
	}
}

// TestFiltfiltConstantTracks pins the DC behaviour: away from the
// cropped transient regions a constant comes back at the constant,
// squared DC gain being exactly one for the Butterworth design.
func TestFiltfiltConstantTracks(t *testing.T) {
	b, a := filtfiltTestFilter(t)
	const n = 400
	vals := make([]float64, n)
	for i := range vals {
		vals[i] = 2.5
	}
	got, err := Filtfilt(b, a, mustFloats(t, vals))
	if err != nil {
		t.Fatalf("Filtfilt: %v", err)
	}
	for i := 100; i < n-100; i++ {
		if math.Abs(got.FloatAt(i)-2.5) > 1e-9 {
			t.Fatalf("sample %d: %g, want 2.5", i, got.FloatAt(i))
		}
	}
}

// TestFiltfiltDtype pins the dtype contract: float32 keeps its dtype,
// int widens to float.
func TestFiltfiltDtype(t *testing.T) {
	b, a := filtfiltTestFilter(t)
	f32, err := core.FromFloat32s([]float32{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13}, 13)
	if err != nil {
		t.Fatal(err)
	}
	out, err := Filtfilt(b, a, f32)
	if err != nil {
		t.Fatalf("Filtfilt float32: %v", err)
	}
	if out.Dtype() != core.Float32 {
		t.Fatalf("float32 input produced dtype %s", out.Dtype())
	}
	ints, err := core.FromInts([]int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13}, 13)
	if err != nil {
		t.Fatal(err)
	}
	out, err = Filtfilt(b, a, ints)
	if err != nil {
		t.Fatalf("Filtfilt int: %v", err)
	}
	if out.Dtype() != core.Float {
		t.Fatalf("int input produced dtype %s", out.Dtype())
	}
}

// TestFiltfiltErrors pins the input gates: the shape and coefficient
// checks FilterApply applies, and the pad-length refusal for short
// signals.
func TestFiltfiltErrors(t *testing.T) {
	b, a := filtfiltTestFilter(t)
	nfilt := max(len(b), len(a))
	short := mustFloats(t, make([]float64, 3*(nfilt-1)))
	if _, err := Filtfilt(b, a, short); err == nil {
		t.Error("signal at exactly the pad length accepted")
	}
	rank2 := mustFloats(t, []float64{1, 2, 3, 4}, 2, 2)
	if _, err := Filtfilt(b, a, rank2); err == nil {
		t.Error("rank-2 accepted")
	}
	cx := mustComplexes(t, []complex128{1, 2, 3, 4, 5}, 5)
	if _, err := Filtfilt(b, a, cx); err == nil {
		t.Error("complex accepted")
	}
	exact := mustFloats(t, make([]float64, 3*(nfilt-1)+1))
	if _, err := Filtfilt([]float64{}, a, exact); err == nil || !strings.Contains(err.Error(), "coefficient lists") {
		t.Errorf("empty b: want the coefficient gate, got %v", err)
	}
	if _, err := Filtfilt([]float64{1}, []float64{0, 1}, exact); err == nil {
		t.Error("a[0] = 0 accepted")
	}
	if _, err := Filtfilt(b, a, mustFloats(t, nil)); err == nil {
		t.Error("empty signal accepted")
	}
}