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

// Command spectral estimates the frequency content of a signal two
// ways: Welch's averaged periodogram on evenly sampled data, and the
// Lomb-Scargle periodogram on the same signal observed at irregular
// times, where an FFT cannot run at all. Both must find the two
// buried sinusoids.
//
// Usage: go run ./examples/spectral
package main

import (
	"fmt"
	"log"
	"math"

	"sourcedock.dev/petrbalvin/tensor"
)

func main() {
	const (
		fs      = 100.0
		seconds = 4.0
		f1      = 5.0
		f2      = 13.0
	)
	n := int(fs * seconds)
	gen := tensor.NewGenerator(11)

	// The signal: two sinusoids plus noise.
	t := make([]float64, n)
	x := make([]float64, n)
	for i := range n {
		t[i] = float64(i) / fs
		x[i] = math.Sin(2*math.Pi*f1*t[i]) + 0.6*math.Sin(2*math.Pi*f2*t[i]) + 0.4*gen.NormalUnit()
	}
	xArr, err := tensor.FromFloats(x, n)
	if err != nil {
		log.Fatal(err)
	}

	// Welch: average periodograms over Hann-windowed segments, the
	// variance-suppressed estimate an FFT alone cannot give.
	freqs, psd, err := tensor.WelchPSD(xArr, fs, 256, 128, "hann")
	if err != nil {
		log.Fatal(err)
	}
	wf1, wf2 := twoPeaks(peakFrequencies(freqs, psd, 2))
	fmt.Printf("welch          peaks at %.2f Hz and %.2f Hz (want %.1f and %.1f)\n", wf1, wf2, f1, f2)

	// Lomb-Scargle: keep every second sample at jittered times, the
	// uneven regime the DFT does not define. The mean rate stays at
	// 50 Hz, comfortably above both sources' Nyquist needs, while the
	// jitter is what makes the ordinary FFT inapplicable.
	times := make([]float64, 0, n/2)
	values := make([]float64, 0, n/2)
	for i := 0; i < n; i += 2 {
		jitter := 0.6 * gen.Unit() / fs
		times = append(times, t[i]+jitter)
		values = append(values, x[i])
	}
	tArr, err := tensor.FromFloats(times, len(times))
	if err != nil {
		log.Fatal(err)
	}
	vArr, err := tensor.FromFloats(values, len(values))
	if err != nil {
		log.Fatal(err)
	}
	lsFreqs, power, err := tensor.LombScargle(tArr, vArr, 1.0, 30.0, 3000)
	if err != nil {
		log.Fatal(err)
	}
	lf1, lf2 := twoPeaks(peakFrequencies(lsFreqs, power, 2))
	fmt.Printf("lomb-scargle   peaks at %.2f Hz and %.2f Hz (want %.1f and %.1f)\n", lf1, lf2, f1, f2)

	for _, got := range []float64{wf1, wf2, lf1, lf2} {
		if math.Abs(got-f1) > 0.3 && math.Abs(got-f2) > 0.3 {
			log.Fatalf("a peak landed at %.2f Hz, away from both sources", got)
		}
	}
}

// peakFrequencies returns the abscissae of the count largest local
// maxima of a periodogram, descending by height and kept at least
// 1.5 Hz apart so a sidelobe of a tall peak cannot shadow a real one.
func peakFrequencies(freqs, power *tensor.Array, count int) []float64 {
	n := freqs.Len()
	// Three-point boxcar smooth: the periodogram's noise is white, a
	// genuine peak is not.
	smooth := make([]float64, n)
	for i := range n {
		lo := max(i-1, 0)
		hi := min(i+1, n-1)
		s := 0.0
		for j := lo; j <= hi; j++ {
			s += power.FloatAt(j)
		}
		smooth[i] = s / float64(hi-lo+1)
	}
	type peak struct {
		f, h float64
	}
	var peaks []peak
	for i := 1; i < n-1; i++ {
		if smooth[i] > smooth[i-1] && smooth[i] >= smooth[i+1] {
			peaks = append(peaks, peak{freqs.FloatAt(i), smooth[i]})
		}
	}
	for i := 1; i < len(peaks); i++ {
		for j := i; j > 0 && peaks[j-1].h < peaks[j].h; j-- {
			peaks[j-1], peaks[j] = peaks[j], peaks[j-1]
		}
	}
	out := make([]float64, 0, count)
	for _, p := range peaks {
		if len(out) == count {
			break
		}
		far := true
		for _, f := range out {
			if math.Abs(p.f-f) < 1.5 {
				far = false
				break
			}
		}
		if far {
			out = append(out, p.f)
		}
	}
	return out
}

// twoPeaks unpacks the two-element result of peakFrequencies.
func twoPeaks(fs []float64) (float64, float64) { return fs[0], fs[1] }