use std::collections::HashMap;
use crate::fft::correction::{dbm_from_vrms, vrms_from_vpeak};
use crate::fft::window::{FftWindow, window_coefficients};
use num_complex::Complex64;
use realfft::RealFftPlanner;
use rustfft::FftPlanner;
use rustfft::num_complex::Complex;
const FFT_EPS: f64 = 1e-18;
const AUTO_FFT_CAP: usize = 8192;
const MAX_FFT_SIZE: usize = 1 << 20;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct WindowKey {
size: usize,
window: FftWindow,
}
#[derive(Clone, Copy)]
pub struct AmplitudeResult {
pub freq: f64,
pub dbm: f64,
}
#[derive(Clone)]
pub struct FftResult {
pub freqs: Vec<f64>,
pub bins: Vec<Complex<f64>>,
pub window: FftWindow,
pub fft_len: usize,
amplitudes: Vec<AmplitudeResult>,
}
impl FftResult {
fn clear(&mut self) {
self.freqs.clear();
self.bins.clear();
self.fft_len = 0;
self.amplitudes.clear();
}
pub fn as_dbm(&mut self, impedance_ohms: Option<f64>) -> &[AmplitudeResult] {
self.as_dbm_internal(impedance_ohms, false)
}
pub fn as_dbm_single_sided(&mut self, impedance_ohms: Option<f64>) -> &[AmplitudeResult] {
self.as_dbm_internal(impedance_ohms, true)
}
fn as_dbm_internal(
&mut self,
impedance_ohms: Option<f64>,
single_sided: bool,
) -> &[AmplitudeResult] {
let size = self.fft_len;
if size < 2 || self.freqs.is_empty() || self.bins.is_empty() {
self.amplitudes.clear();
return &self.amplitudes;
}
let half = size / 2;
self.amplitudes.clear();
self.amplitudes.reserve(self.freqs.len());
for (idx, (freq, value)) in self.freqs.iter().zip(self.bins.iter()).enumerate() {
let mag = (value.re * value.re + value.im * value.im).sqrt();
let dbm = if let Some(impedance_ohms) = impedance_ohms {
if impedance_ohms <= 0.0 {
f64::NEG_INFINITY
} else {
let scale = if single_sided
&& !(idx == 0 || (idx == half - 1 && size.is_multiple_of(2)))
{
2.0
} else {
1.0
};
let v_peak = scale * mag / (size as f64 * self.window.coherent_gain());
dbm_from_vrms(vrms_from_vpeak(v_peak), impedance_ohms)
}
} else {
10.0 * (mag * mag).max(FFT_EPS).log10()
};
self.amplitudes.push(AmplitudeResult { freq: *freq, dbm });
}
&self.amplitudes
}
}
impl Default for FftResult {
fn default() -> Self {
Self {
freqs: Vec::new(),
bins: Vec::new(),
window: FftWindow::Rectangular,
fft_len: 0,
amplitudes: Vec::new(),
}
}
}
pub struct ReusableFftStateReal {
window_cache: HashMap<WindowKey, Vec<f64>>,
planner: RealFftPlanner<f64>,
input: Vec<f64>,
output: Vec<Complex<f64>>,
scratch: Vec<Complex<f64>>,
result: FftResult,
}
impl Default for ReusableFftStateReal {
fn default() -> Self {
Self::new()
}
}
impl ReusableFftStateReal {
pub fn new() -> Self {
Self {
window_cache: HashMap::new(),
planner: RealFftPlanner::new(),
input: Vec::new(),
output: Vec::new(),
scratch: Vec::new(),
result: FftResult::default(),
}
}
pub(super) fn fft_points_mag(
&mut self,
samples: impl ExactSizeIterator<Item = f64>,
sample_rate: f64,
fft_len: Option<usize>,
window: FftWindow,
) -> &mut FftResult {
if samples.len() < 2 || sample_rate <= 0.0 {
self.result.clear();
return &mut self.result;
}
let Some(size) = fft_size(samples.len(), fft_len) else {
self.result.clear();
return &mut self.result;
};
let key = WindowKey { size, window };
let window_coeffs = self
.window_cache
.entry(key)
.or_insert_with(|| window_coefficients(size, window));
self.input.clear();
self.input.extend(
samples
.take(size)
.zip(window_coeffs.iter())
.map(|(value, coeff)| value * coeff),
);
let fft = self.planner.plan_fft_forward(size);
if self.output.len() != fft.complex_len() {
self.output.resize(fft.complex_len(), Complex::default());
}
if self.scratch.len() != fft.get_scratch_len() {
self.scratch
.resize(fft.get_scratch_len(), Complex::default());
}
let _ = fft.process_with_scratch(&mut self.input, &mut self.output, &mut self.scratch);
let half = size / 2;
self.result.freqs.clear();
self.result.bins.clear();
self.result.freqs.reserve(half);
self.result.bins.reserve(half);
for idx in 0..half {
self.result
.freqs
.push((idx as f64) * sample_rate / (size as f64));
self.result.bins.push(self.output[idx]);
}
self.result.window = window;
self.result.fft_len = size;
&mut self.result
}
}
pub struct ReusableFftStateComplex {
window_cache: HashMap<WindowKey, Vec<f64>>,
planner: FftPlanner<f64>,
buffer: Vec<Complex<f64>>,
scratch: Vec<Complex<f64>>,
result: FftResult,
}
impl Default for ReusableFftStateComplex {
fn default() -> Self {
Self::new()
}
}
impl ReusableFftStateComplex {
pub fn new() -> Self {
Self {
window_cache: HashMap::new(),
planner: FftPlanner::new(),
buffer: Vec::new(),
scratch: Vec::new(),
result: FftResult::default(),
}
}
pub(super) fn fft_points_complex(
&mut self,
samples: impl ExactSizeIterator<Item = Complex64>,
sample_rate: f64,
center_frequency: f64,
fft_len: Option<usize>,
window: FftWindow,
) -> &mut FftResult {
if samples.len() < 2 || sample_rate <= 0.0 {
self.result.clear();
return &mut self.result;
}
let Some(size) = fft_size(samples.len(), fft_len) else {
self.result.clear();
return &mut self.result;
};
let key = WindowKey { size, window };
let window_coeffs = self
.window_cache
.entry(key)
.or_insert_with(|| window_coefficients(size, window));
self.buffer.clear();
self.buffer.extend(
samples
.take(size)
.zip(window_coeffs.iter())
.map(|(value, coeff)| Complex::new(value.re * *coeff, value.im * *coeff)),
);
let fft = self.planner.plan_fft_forward(size);
if self.scratch.len() != fft.get_inplace_scratch_len() {
self.scratch
.resize(fft.get_inplace_scratch_len(), Complex::default());
}
fft.process_with_scratch(&mut self.buffer, &mut self.scratch);
let half = size / 2;
self.result.freqs.clear();
self.result.bins.clear();
self.result.freqs.reserve(size);
self.result.bins.reserve(size);
for idx in 0..size {
let shifted_idx = (idx + half) % size;
let freq = ((idx as isize - half as isize) as f64) * sample_rate / (size as f64)
+ center_frequency;
self.result.freqs.push(freq);
self.result.bins.push(self.buffer[shifted_idx]);
}
self.result.window = window;
self.result.fft_len = size;
&mut self.result
}
}
fn fft_size(len: usize, override_len: Option<usize>) -> Option<usize> {
if len < 2 {
return None;
}
let target = match override_len {
Some(value) => value.min(len).min(MAX_FFT_SIZE),
None => len.min(AUTO_FFT_CAP),
};
if target < 2 {
return None;
}
let mut size = target.next_power_of_two();
if size > target {
size /= 2;
}
if size < 2 {
return None;
}
Some(size.min(MAX_FFT_SIZE))
}