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

// TestLombScargleFindsPeak recovers the known 0.1-cycle-per-sample
// period from unevenly sampled data: the grid point nearest the true
// frequency must carry the largest power by a clear margin.
func TestLombScargleFindsPeak(t *testing.T) {
	const trueFreq = 0.1
	n := 200
	times := make([]float64, n)
	values := make([]float64, n)
	for i := range n {
		// Deterministic uneven sampling: every third tick skipped.
		times[i] = float64(i) * 1.3
		values[i] = math.Sin(2*math.Pi*trueFreq*times[i]) + 0.3*math.Cos(2*math.Pi*0.31*times[i])
	}
	freqs, power, err := LombScargle(mustFloats(t, times, n), mustFloats(t, values, n),
		0.02, 0.45, 400)
	if err != nil {
		t.Fatalf("LombScargle: %v", err)
	}
	best, bestF := 0, 0.0
	for i := range power.Len() {
		if power.FloatAt(i) > float64(best) {
			best = i
			bestF = freqs.FloatAt(i)
		}
	}
	if math.Abs(bestF-trueFreq) > 0.45/400*3 {
		t.Fatalf("peak at %.5f, want within three bins of %.5f", bestF, trueFreq)
	}
	// The second tone must also stand out on the grid.
	power2 := 0.0
	for i := range power.Len() {
		if math.Abs(freqs.FloatAt(i)-0.31) < 0.01 && power.FloatAt(i) > power2 {
			power2 = power.FloatAt(i)
		}
	}
	if power2 <= 0 {
		t.Fatal("the secondary tone left no peak")
	}
}

// TestLombScargleScale pins the classical normalisation: a pure
// unit-amplitude sinusoid sampled evenly at exactly its period peaks
// at n/(4·var) ≈ n/2 for unit-amplitude data of unit variance… checked
// against the directly evaluated defining sum instead of a hand rule.
func TestLombScargleScale(t *testing.T) {
	n := 60
	times := make([]float64, n)
	values := make([]float64, n)
	for i := range n {
		times[i] = float64(i)
		values[i] = math.Sin(2 * math.Pi * float64(i) / 12)
	}
	_, power, err := LombScargle(mustFloats(t, times, n), mustFloats(t, values, n),
		1.0/12.0, 1.0/12.0, 1)
	if err != nil {
		t.Fatalf("LombScargle: %v", err)
	}
	// Direct reference: the two orthogonal sums at the exact tone.
	mean := 0.0
	for i := range n {
		mean += values[i]
	}
	mean /= float64(n)
	variance := 0.0
	for i := range n {
		d := values[i] - mean
		variance += d * d
	}
	variance /= float64(n - 1)
	omega := 2 * math.Pi / 12
	sc, ss := 0.0, 0.0
	for i := range n {
		sc += math.Cos(omega * float64(i))
		ss += math.Sin(omega * float64(i))
	}
	tau := 0.5 * math.Atan2(ss, sc) / omega
	sumCos, sumSin, sumCosSq, sumSinSq := 0.0, 0.0, 0.0, 0.0
	for i := range n {
		arg := omega * (float64(i) - tau)
		c, s := math.Cos(arg), math.Sin(arg)
		sumCos += (values[i] - mean) * c
		sumSin += (values[i] - mean) * s
		sumCosSq += c * c
		sumSinSq += s * s
	}
	want := (sumCos*sumCos/sumCosSq + sumSin*sumSin/sumSinSq) / (2 * variance)
	if got := power.FloatAt(0); math.IsNaN(got) || math.IsInf(got, 0) {
		t.Fatalf("power = %v, want a finite value comparable to %.12g", got, want)
	}
	if math.Abs(power.FloatAt(0)-want) > 1e-9 {
		t.Fatalf("power = %.12g, want the direct sum %.12g", power.FloatAt(0), want)
	}
}

// TestLombScargleErrors pins the validation contract.
func TestLombScargleErrors(t *testing.T) {
	times := mustFloats(t, []float64{0, 1, 2, 3}, 4)
	values := mustFloats(t, []float64{1, 2, 1, 2}, 4)
	pair := mustFloats(t, []float64{0, 1}, 2)
	if _, _, err := LombScargle(pair, mustFloats(t, []float64{1, 2}, 2), 0.1, 1, 5); err == nil {
		t.Fatal("expected an error for a two-point time base")
	}
	short := mustFloats(t, []float64{0, 1, 2}, 3)
	flat := mustFloats(t, []float64{1, 1, 1}, 3)
	if _, _, err := LombScargle(short, flat, 0.1, 1, 5); err == nil {
		t.Fatal("expected an error for zero variance")
	}
	if _, _, err := LombScargle(times, mustFloats(t, []float64{1, 2}, 2), 0.1, 1, 5); err == nil {
		t.Fatal("expected an error for a length mismatch")
	}
	if _, _, err := LombScargle(times, values, 0, 1, 5); err == nil {
		t.Fatal("expected an error for a non-positive minFreq")
	}
	if _, _, err := LombScargle(times, values, 1, 0.5, 5); err == nil {
		t.Fatal("expected an error for an inverted range")
	}
	if _, _, err := LombScargle(times, values, 0.1, 1, 0); err == nil {
		t.Fatal("expected an error for zero frequency points")
	}
}

// TestWelchPSDWhiteNoise checks the level: filtered deterministic
// samples with unit variance must estimate a flat band near one.
func TestWelchPSDWhiteNoise(t *testing.T) {
	n := 4096
	x := make([]float64, n)
	g := core.NewGenerator(7)
	draws, err := core.Normal(g, n, 0, 1)
	if err != nil {
		t.Fatalf("Normal: %v", err)
	}
	for i := range n {
		x[i] = draws.FloatAt(i)
	}
	_, psd, err := WelchPSD(mustFloats(t, x, n), 1000, 256, 128, "hann")
	if err != nil {
		t.Fatalf("WelchPSD: %v", err)
	}
	mean := 0.0
	for i := range psd.Len() {
		mean += psd.FloatAt(i) / float64(psd.Len())
	}
	// A flat unit-variance band integrates to σ² over fs/2, so the
	// level sits at 2σ²/fs in PSD units of x²/Hz.
	if r := mean / (2 / 1000.0); math.Abs(r-1) > 0.15 {
		t.Fatalf("mean PSD = %.6f, want 2σ²/fs = %.6f (ratio %.3f)", mean, 2/1000.0, r)
	}
}

// TestWelchPSDSinusoid pins the peak location and the Parseval
// balance: the PSD integrated over frequency returns the signal
// variance.
func TestWelchPSDSinusoid(t *testing.T) {
	const fs = 128.0
	const tone = 16.0
	n := 2048
	x := make([]float64, n)
	for i := range n {
		x[i] = math.Sin(2 * math.Pi * tone * float64(i) / fs)
	}
	freqs, psd, err := WelchPSD(mustFloats(t, x, n), fs, 256, 128, "hann")
	if err != nil {
		t.Fatalf("WelchPSD: %v", err)
	}
	peak, best := 0, 0.0
	for i := range psd.Len() {
		if psd.FloatAt(i) > best {
			best = psd.FloatAt(i)
			peak = i
		}
	}
	if math.Abs(freqs.FloatAt(peak)-tone) > fs/256 {
		t.Fatalf("PSD peak at %.4f Hz, want %.1f", freqs.FloatAt(peak), tone)
	}
	// Parseval: sum(psd)·df ≈ variance for a windowed estimate on a
	// signal with negligible edge leakage.
	total, df := 0.0, fs/256
	for i := range psd.Len() {
		total += psd.FloatAt(i)
	}
	total *= df
	sq := 0.0
	for i := range n {
		sq += x[i] * x[i]
	}
	variance := sq / float64(n)
	if math.Abs(total-variance)/variance > 0.1 {
		t.Fatalf("Parseval: ∫PSD = %.4f, variance = %.4f (rel %.3f)", total, variance,
			math.Abs(total-variance)/variance)
	}
}

// TestWelchPSDErrors pins the validation contract.
func TestWelchPSDErrors(t *testing.T) {
	x := mustFloats(t, make([]float64, 64), 64)
	if _, _, err := WelchPSD(x, 100, 128, 0, "hann"); err == nil {
		t.Fatal("expected an error for a segment longer than the signal")
	}
	if _, _, err := WelchPSD(x, 100, 32, 32, "hann"); err == nil {
		t.Fatal("expected an error for a full overlap")
	}
	if _, _, err := WelchPSD(x, 100, 32, -1, "hann"); err == nil {
		t.Fatal("expected an error for a negative overlap")
	}
	if _, _, err := WelchPSD(x, 0, 32, 16, "hann"); err == nil {
		t.Fatal("expected an error for a zero sampling rate")
	}
	if _, _, err := WelchPSD(x, 100, 32, 16, "kaiser"); err == nil {
		t.Fatal("expected an error for an unknown window")
	}
	rank2, _ := core.FromFloats([]float64{1, 2, 3, 4}, 2, 2)
	if _, _, err := WelchPSD(rank2, 100, 2, 1, "hann"); err == nil {
		t.Fatal("expected an error for a rank-2 signal")
	}
}

// TestWelchPSDSegmentCount pins the segment arithmetic: a signal of
// n samples with non-overlapping segments fills exactly
// (n−overlap)/(segment−overlap) segments, the last one included. The
// third of the three segments here carries the only energy, so a
// dropped final segment would leave the Nyquist bin empty, and the
// average over three segments puts |X|²=64 at exactly 64/(3·fs·wPower).
func TestWelchPSDSegmentCount(t *testing.T) {
	n := 24
	x := make([]float64, n) // segments 1 and 2 are silent
	for i := 16; i < n; i++ {
		x[i] = 1 // the alternating ±1 pattern shifted to +1: 16+i even
		if (i-16)%2 == 1 {
			x[i] = -1
		}
	}
	freqs, psd, err := WelchPSD(mustFloats(t, x, n), 1, 8, 0, "box")
	if err != nil {
		t.Fatalf("WelchPSD: %v", err)
	}
	if psd.Len() != 5 {
		t.Fatalf("bins = %d, want 5", psd.Len())
	}
	for k := range 5 {
		want := 0.0
		if k == 4 {
			want = 64.0 / (3 * 1 * 8) // one segment of three carries |X[4]|² = 64
		}
		if math.Abs(psd.FloatAt(k)-want) > 1e-9 {
			t.Fatalf("psd[%d] = %.9g, want %.9g", k, psd.FloatAt(k), want)
		}
		if k == 4 && math.Abs(freqs.FloatAt(k)-0.5) > 1e-9 {
			t.Fatalf("bin 4 sits at %.6f Hz, want the 0.5 Nyquist", freqs.FloatAt(k))
		}
	}
	// A signal exactly one segment long is a legal single-segment
	// periodogram, not an error.
	solo, psd1, err := WelchPSD(mustFloats(t, x[16:], 8), 1, 8, 0, "box")
	if err != nil {
		t.Fatalf("WelchPSD single segment: %v", err)
	}
	if solo.Len() != 5 || psd1.Len() != 5 {
		t.Fatalf("single-segment shapes = %d/%d, want 5/5", solo.Len(), psd1.Len())
	}
	if math.Abs(psd1.FloatAt(4)-8) > 1e-9 {
		t.Fatalf("single-segment psd[4] = %.9g, want 64/(1·1·8) = 8", psd1.FloatAt(4))
	}
}

// TestLombScargleSingleFrequency pins the nFreq == 1 contract: the
// grid holds the single frequency minFreq and a finite power (the
// old grid formula divided 0/0 and produced NaN).
func TestLombScargleSingleFrequency(t *testing.T) {
	n := 40
	times := make([]float64, n)
	values := make([]float64, n)
	for i := range n {
		times[i] = 1.3 * float64(i)
		values[i] = math.Sin(2 * math.Pi * 0.1 * times[i])
	}
	freqs, power, err := LombScargle(mustFloats(t, times, n), mustFloats(t, values, n), 0.1, 0.7, 1)
	if err != nil {
		t.Fatalf("LombScargle: %v", err)
	}
	if freqs.Len() != 1 || power.Len() != 1 {
		t.Fatalf("shapes %d/%d, want 1/1", freqs.Len(), power.Len())
	}
	if freqs.FloatAt(0) != 0.1 {
		t.Fatalf("single frequency = %g, want minFreq 0.1", freqs.FloatAt(0))
	}
	if p := power.FloatAt(0); math.IsNaN(p) || math.IsInf(p, 0) || p <= 0 {
		t.Fatalf("single-frequency power = %g, want a finite positive value", p)
	}
}

// TestLombScargleConstantTimesErrors pins the degenerate time base: a
// constant time base carries no phase information and used to drive
// the sine fit to 0/0 NaN.
func TestLombScargleConstantTimesErrors(t *testing.T) {
	times := mustFloats(t, []float64{2, 2, 2, 2, 2}, 5)
	values := mustFloats(t, []float64{1, -1, 1, -1, 1}, 5)
	if _, _, err := LombScargle(times, values, 0.1, 1, 5); err == nil {
		t.Fatal("expected an error for an all-equal time base")
	}
}