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

package signal_test

// Runnable godoc examples for the flagship workflows. Every one of them
// runs on a fixed input and prints a fixed result, so `go test` checks
// the documentation against the code.

import (
	"fmt"
	"log"
	"math"

	tensor "sourcedock.dev/petrbalvin/tensor"
	"sourcedock.dev/petrbalvin/tensor/signal"
)

// A two-tone signal at 4 Hz and 12 Hz, sampled at 64 Hz: Welch's
// averaged, windowed estimate puts both tones on their own bins.
func ExampleWelchPSD() {
	const n, fs = 64, 64.0
	raw := make([]float64, n)
	for i := range raw {
		t := float64(i) / fs
		raw[i] = math.Cos(2*math.Pi*4*t) + 0.5*math.Cos(2*math.Pi*12*t)
	}
	x, err := tensor.FromFloats(raw, n)
	if err != nil {
		log.Fatal(err)
	}
	freqs, psd, err := signal.WelchPSD(x, fs, 32, 16, "hann")
	if err != nil {
		log.Fatal(err)
	}
	peak := 0
	for k := range psd.Len() {
		if psd.FloatAt(k) > psd.FloatAt(peak) {
			peak = k
		}
	}
	fmt.Printf("peak %.0f Hz, %.3f power, %d bins\n",
		freqs.FloatAt(peak), psd.FloatAt(peak), psd.Len())
	// Output: peak 4 Hz, 0.167 power, 17 bins
}

// A 5 Hz tone in a 20 Hz passband: the one-pass filter lags the tone,
// the forward-and-backward sweep leaves it where it was.
func ExampleFiltfilt() {
	const n, fs, f = 200, 100.0, 5.0
	raw := make([]float64, n)
	for i := range raw {
		raw[i] = math.Sin(2 * math.Pi * f * float64(i) / fs)
	}
	x, err := tensor.FromFloats(raw, n)
	if err != nil {
		log.Fatal(err)
	}
	b, a, err := signal.ButterworthLowPass(2, fs, 20)
	if err != nil {
		log.Fatal(err)
	}
	once, err := signal.FilterApply(b, a, x)
	if err != nil {
		log.Fatal(err)
	}
	twice, err := signal.Filtfilt(b, a, x)
	if err != nil {
		log.Fatal(err)
	}
	peak := func(y *tensor.Array) int {
		at := 0
		for i := range y.Len() {
			if y.FloatAt(i) > y.FloatAt(at) {
				at = i
			}
		}
		return at
	}
	fmt.Printf("input peaks at %d, one pass at %d, filtfilt at %d\n",
		peak(x), peak(once), peak(twice))
	// Output: input peaks at 5, one pass at 6, filtfilt at 5
}

// The DB2 transform packs the coefficients as [A, D2, D1]. The periodic
// boundary keeps the energy, so the inverse returns the signal.
func ExampleDaubechiesDWT() {
	const levels = 2
	x, err := tensor.FromFloats([]float64{1, 2, 3, 4, 5, 6, 7, 8}, 8)
	if err != nil {
		log.Fatal(err)
	}
	coef, err := signal.DaubechiesDWT(x, signal.DB2, levels, signal.DWTPeriodic)
	if err != nil {
		log.Fatal(err)
	}
	back, err := signal.DaubechiesIDWT(coef, signal.DB2, levels, signal.DWTPeriodic)
	if err != nil {
		log.Fatal(err)
	}
	energy := func(v *tensor.Array) float64 {
		total := 0.0
		for i := range v.Len() {
			total += v.FloatAt(i) * v.FloatAt(i)
		}
		return total
	}
	worst := 0.0
	for i := range x.Len() {
		worst = math.Max(worst, math.Abs(back.FloatAt(i)-x.FloatAt(i)))
	}
	fmt.Printf("%d levels, %d coefficients, round trip exact: %v\n",
		levels, coef.Len(), worst < 1e-12)
	fmt.Printf("energy kept to %.3f\n", energy(coef)/energy(x))
	// Output: 2 levels, 8 coefficients, round trip exact: true
	// energy kept to 1.000
}

// The Hilbert envelope of a tone whose amplitude swings at 1 Hz: the
// modulus of the analytic signal traces the swing without smoothing lag.
func ExampleEnvelope() {
	const n, fs = 256, 128.0
	raw := make([]float64, n)
	for i := range raw {
		t := float64(i) / fs
		// A 20 Hz carrier under a 1 Hz amplitude swing.
		raw[i] = (1 + 0.5*math.Cos(2*math.Pi*t)) * math.Cos(2*math.Pi*20*t)
	}
	x, err := tensor.FromFloats(raw, n)
	if err != nil {
		log.Fatal(err)
	}
	env, err := signal.Envelope(x)
	if err != nil {
		log.Fatal(err)
	}
	// The swing peaks at t = 0 s and t = 1 s, and bottoms at t = 0.5 s.
	at := func(seconds float64) int { return int(seconds * fs) }
	fmt.Printf("t=0 s %.3f, t=0.5 s %.3f, t=1 s %.3f\n",
		env.FloatAt(at(0)), env.FloatAt(at(0.5)), env.FloatAt(at(1)))
	// Output: t=0 s 1.500, t=0.5 s 0.500, t=1 s 1.500
}

// A constant position observed three times: the linear Kalman filter
// fuses each measurement into the running estimate and reports the
// filtered means, one row per measurement.
func ExampleKalmanFilter() {
	z, err := tensor.FromFloats([]float64{1, 2, 3}, 3)
	if err != nil {
		log.Fatal(err)
	}
	f, err := tensor.FromFloats([]float64{1}, 1, 1)
	if err != nil {
		log.Fatal(err)
	}
	h, err := tensor.FromFloats([]float64{1}, 1, 1)
	if err != nil {
		log.Fatal(err)
	}
	res, err := signal.KalmanFilter(z, f, h, signal.KalmanOptions{})
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("states %.3f %.3f %.3f\n", res.States.FloatAt(0), res.States.FloatAt(1), res.States.FloatAt(2))
	fmt.Printf("log likelihood %.2f\n", res.LogLikelihood)
	// Output: states 0.500 1.000 1.500
	// log likelihood -5.95
}

// An AR(1) series built from a fixed linear congruential driver: the
// Yule-Walker fit recovers the coefficient that generated it.
func ExampleEstimateAR() {
	const n = 512
	raw := make([]float64, n)
	seed := 1.0
	for i := range raw {
		// The minimal-standard generator: deterministic, so the
		// example never calls a random source.
		seed = math.Mod(16807*seed, 2147483647)
		raw[i] = seed/2147483647 - 0.5
		if i > 0 {
			raw[i] += 0.5 * raw[i-1]
		}
	}
	x, err := tensor.FromFloats(raw, n)
	if err != nil {
		log.Fatal(err)
	}
	res, err := signal.EstimateAR(x, 1)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("AR(%d): phi_1 %.3f\n", res.P, res.AR[0])
	// Output: AR(1): phi_1 0.495
}