package dsp

import "math"

// GoertzelEnergy computes the energy at a specific frequency in a block
// of samples using the Goertzel algorithm. Returns the magnitude squared.
func GoertzelEnergy(samples []float64, sampleRate, targetFreqHz float64) float64 {
	n := len(samples)
	if n == 0 {
		return 0
	}
	k := int(0.5 + float64(n)*targetFreqHz/sampleRate)
	w := 2 * math.Pi * float64(k) / float64(n)
	coeff := 2 * math.Cos(w)
	var s0, s1, s2 float64
	for _, x := range samples {
		s0 = x + coeff*s1 - s2
		s2 = s1
		s1 = s0
	}
	return s1*s1 + s2*s2 - coeff*s1*s2
}

// BandEnergy computes aggregate energy in a frequency band by averaging
// Goertzel results at the center and ±span/4 offsets.
func BandEnergy(samples []float64, sampleRate, centerHz, spanHz float64) float64 {
	e0 := GoertzelEnergy(samples, sampleRate, centerHz)
	e1 := GoertzelEnergy(samples, sampleRate, centerHz-spanHz/4)
	e2 := GoertzelEnergy(samples, sampleRate, centerHz+spanHz/4)
	return (e0 + e1 + e2) / 3
}