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

package signal

import (
	"math"

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

// The public window catalogue. Every builder returns a fresh
// preallocated []float64 of n samples. Two length conventions exist and
// the periodic flag picks between them: the symmetric window divides
// its argument by n−1, so its first and last samples coincide (the
// right shape for a finite impulse-response design) and the periodic
// window divides by n, which makes the sequence one exact period of
// its underlying continuous shape (the right shape for spectral
// estimates, where the segment is treated as one period and the
// doubled lobes of the symmetric tail would leak). periodic is the
// last argument, false everywhere it does not matter; a one-sample
// window is the single value 1 in both conventions. Every builder
// refuses n below 1.

// WindowBox returns the untapered box: n ones, the window that
// filters nothing. The periodic flag changes nothing here and exists
// only for signature uniformity across the catalogue.
func WindowBox(n int, periodic bool) ([]float64, error) {
	w, _, one, err := windowSetup("WindowBox", n, periodic)
	if err != nil || one {
		return w, err
	}
	for i := range w {
		w[i] = 1
	}
	return w, nil
}

// WindowHann returns the Hann window, the raised cosine
// 0.5 − 0.5·cos(2πx), the gentlest of the generalised cosines: zero
// at the edges in both conventions, −6 dB per octave sidelobe roll-off.
func WindowHann(n int, periodic bool) ([]float64, error) {
	return generalCosine("WindowHann", n, hannCoeffs, periodic)
}

// WindowHamming returns the Hamming window, the raised cosine on a
// pedestal 0.54 − 0.46·cos(2πx): the nonzero pedestal cancels the
// Hann window's first sidelobe, at the cost of a floor the outer
// sidelobes never drop below.
func WindowHamming(n int, periodic bool) ([]float64, error) {
	return generalCosine("WindowHamming", n, hammingCoeffs, periodic)
}

// WindowBlackman returns the (exact) Blackman window
// 0.42 − 0.5·cos(2πx) + 0.08·cos(4πx): two cosines instead of
// Hann's one buy sidelobes below −58 dB at the price of a doubled
// main lobe.
func WindowBlackman(n int, periodic bool) ([]float64, error) {
	return generalCosine("WindowBlackman", n, blackmanCoeffs, periodic)
}

// WindowBlackmanHarris returns the four-term Blackman-Harris window,
// the minimum-sidelobe member of the generalised-cosine family with
// four terms: sidelobes below −92 dB, a main lobe three Hann lobes
// wide. The usual choice when dynamic range matters more than
// resolution.
func WindowBlackmanHarris(n int, periodic bool) ([]float64, error) {
	return generalCosine("WindowBlackmanHarris", n, blackmanHarrisCoeffs, periodic)
}

// WindowFlatTop returns the flat-top window, the five-term generalised
// cosine whose main lobe is flat to within a hundredth of a decibel:
// the amplitude of a spectral line reads true to the window's ripple
// no matter where the line falls between bins, which is what the wide
// lobe buys. The edge samples are slightly negative, so the window is
// for amplitude metrology, not for filtering.
func WindowFlatTop(n int, periodic bool) ([]float64, error) {
	return generalCosine("WindowFlatTop", n, flatTopCoeffs, periodic)
}

// WindowBartlett returns the Bartlett window, the triangle 1 − |2x − 1|:
// the piecewise-linear taper, zero at both edges, whose sidelobes sit
// between the box's and Hann's. It is the Fejér kernel of the box and
// is non-negative everywhere, which the generalised cosines are not.
func WindowBartlett(n int, periodic bool) ([]float64, error) {
	w, den, one, err := windowSetup("WindowBartlett", n, periodic)
	if err != nil || one {
		return w, err
	}
	for i := range w {
		w[i] = 1 - math.Abs(2*float64(i)/den-1)
	}
	return w, nil
}

// WindowKaiser returns the Kaiser window of parameter beta: the
// modified Bessel taper I0(beta·sqrt(1 − r²))/I0(beta) over the
// normalised radius r = 2x − 1, the adjustable compromise between main
// lobe width and sidelobe height. beta 0 is the box; near 5 the
// sidelobes sit around −30 dB, near 9 around −60 dB, and the usual
// rule of thumb spends about 2.2·beta decibels of stopband. beta must
// be finite and non-negative.
func WindowKaiser(n int, beta float64, periodic bool) ([]float64, error) {
	const name = "WindowKaiser"
	if beta < 0 || math.IsNaN(beta) || math.IsInf(beta, 0) {
		return nil, base.Errf("%s: beta must be finite and non-negative, got %g", name, beta)
	}
	w, den, one, err := windowSetup(name, n, periodic)
	if err != nil || one {
		return w, err
	}
	i0b := kaiserI0(beta)
	for i := range w {
		r := 2*float64(i)/den - 1
		// The radius can leave the unit disk by a rounding step at
		// the edges; the squared radius is clamped so the square
		// root stays real.
		w[i] = kaiserI0(beta*math.Sqrt(math.Max(0, 1-r*r))) / i0b
	}
	return w, nil
}

// WindowCosine returns the cosine (sine) window sin(πx): one positive
// half-period whose derivative vanishes at neither edge, the taper of
// the MDCT and of the minimum-tap Blackman derivations.
func WindowCosine(n int, periodic bool) ([]float64, error) {
	w, den, one, err := windowSetup("WindowCosine", n, periodic)
	if err != nil || one {
		return w, err
	}
	for i := range w {
		w[i] = math.Sin(math.Pi * float64(i) / den)
	}
	return w, nil
}

// The generalised-cosine coefficient sets, with the signs carried in
// the table: sample i of the symmetric n-window is
// Σ_k c_k·cos(2πk·i/(n−1)). The flat-top set is written as the exact
// fractions of 19 the amplitude-calibration standard defines it by.
var (
	hannCoeffs           = []float64{0.5, -0.5}
	hammingCoeffs        = []float64{0.54, -0.46}
	blackmanCoeffs       = []float64{0.42, -0.5, 0.08}
	blackmanHarrisCoeffs = []float64{0.35875, -0.48829, 0.14128, -0.01168}
	flatTopCoeffs        = []float64{4.096 / 19, -7.916 / 19, 5.268 / 19, -1.588 / 19, 0.132 / 19}
)

// generalCosine evaluates the generalised cosine family: the sum of
// signed cosine terms c_k over x = i/den, the denominator chosen by
// the symmetric or periodic convention. The term arguments keep the
// exact shape 2πk·i/den the package's spectral estimates have always
// fed their Hann and Hamming windows, so the periodic two-term
// members reproduce the legacy windowTaper outputs bit for bit.
func generalCosine(name string, n int, coeffs []float64, periodic bool) ([]float64, error) {
	w, den, one, err := windowSetup(name, n, periodic)
	if err != nil || one {
		return w, err
	}
	for i := range w {
		acc := coeffs[0]
		for k, c := range coeffs[1:] {
			// The explicit conversion is the FMA fence: the v4 build
			// contracts a bare product-plus-add and the window bits
			// drift a ulp from the portable build's, which the
			// bit-identity pin holds.
			acc += float64(c * math.Cos(2*math.Pi*float64(k+1)*float64(i)/den))
		}
		w[i] = acc
	}
	return w, nil
}

// windowSetup validates the requested length, prepares the window
// buffer and returns the denominator the window argument divides by:
// n−1 for the symmetric convention, n for the periodic one. The
// one-sample window is the constant 1 in both, so the flag-one return
// hands back that finished window and no builder reaches a zero
// denominator.
func windowSetup(name string, n int, periodic bool) (w []float64, den float64, one bool, err error) {
	if n < 1 {
		return nil, 0, false, base.Errf("%s: n must be at least 1, got %d", name, n)
	}
	if n == 1 {
		return []float64{1}, 1, true, nil
	}
	if periodic {
		return make([]float64, n), float64(n), false, nil
	}
	return make([]float64, n), float64(n - 1), false, nil
}

// windowTaper builds the named window of the given length for the
// spectral estimates, routing through the catalogue's periodic forms:
// "hann" and "hamming" are WindowHann and WindowHamming at
// periodic=true, "box" is WindowBox, and the outputs are the same
// bits the dedicated loops this replaces produced for every length
// the callers accept (they refuse n below 2, where the two
// conventions differ). The legacy names stay because WelchPSD, STFT
// and Spectrogram publish them.
func windowTaper(window string, n int) ([]float64, error) {
	switch window {
	case "box":
		return WindowBox(n, true)
	case "hann":
		return WindowHann(n, true)
	case "hamming":
		return WindowHamming(n, true)
	default:
		return nil, base.Errf("unknown window %q, want hann, hamming or box", window)
	}
}