Initial version

2 files changed, 514 insertions, 0 deletions

diff --git a/src/decoder.rs b/src/decoder.rs
new file mode 100644
index 0000000..657bff7
--- /dev/null
+++ b/src/decoder.rs
@@ -0,0 +1,442 @@
+// A Rust implementation of the ggmorse signal processing pipeline.
+// This includes a resampler, band-pass filter, STFT for pitch detection,
+// a Goertzel filter for tone extraction, and the core decoding logic.
+
+use anyhow::{Result, bail};
+use collections::VecDeque;
+use rubato::{InterpolationParameters, InterpolationType, Resampler, SincFixedIn, WindowFunction};
+use rustfft::{Fft, FftPlanner, num_complex::Complex};
+use std::sync::Arc;
+
+// --- DSP Constants (ported from ggmorse) ---
+const FREQ_MIN_HZ: f32 = 200.0;
+const FREQ_MAX_HZ: f32 = 1200.0;
+const HISTORY_S: f32 = 3.0; // How many seconds of audio to keep for analysis
+
+// --- Biquad Filter (Corrected Implementation) ---
+#[derive(Debug, Clone, Copy)]
+pub enum FilterType {
+    HighPass,
+    LowPass,
+}
+
+pub struct BiquadFilter {
+    a0: f32,
+    a1: f32,
+    a2: f32,
+    b1: f32,
+    b2: f32,
+    x1: f32,
+    x2: f32, // Delayed inputs
+    y1: f32,
+    y2: f32, // Delayed outputs
+}
+
+impl BiquadFilter {
+    pub fn new(filter_type: FilterType, cutoff_hz: f32, sample_rate: u32) -> Self {
+        let mut filter = Self {
+            a0: 1.0,
+            a1: 0.0,
+            a2: 0.0,
+            b1: 0.0,
+            b2: 0.0,
+            x1: 0.0,
+            x2: 0.0,
+            y1: 0.0,
+            y2: 0.0,
+        };
+        let c = (std::f32::consts::PI * cutoff_hz / sample_rate as f32).tan();
+        let sqrt2 = 2.0f32.sqrt();
+
+        match filter_type {
+            FilterType::LowPass => {
+                let d = 1.0 / (1.0 + sqrt2 * c + c * c);
+                filter.a0 = c * c * d;
+                filter.a1 = 2.0 * filter.a0;
+                filter.a2 = filter.a0;
+                filter.b1 = 2.0 * (c * c - 1.0) * d;
+                filter.b2 = (1.0 - sqrt2 * c + c * c) * d;
+            }
+            FilterType::HighPass => {
+                let d = 1.0 / (1.0 + sqrt2 * c + c * c);
+                filter.a0 = d;
+                filter.a1 = -2.0 * d;
+                filter.a2 = d;
+                filter.b1 = 2.0 * (c * c - 1.0) * d;
+                filter.b2 = (1.0 - sqrt2 * c + c * c) * d;
+            }
+        }
+        filter
+    }
+
+    pub fn process(&mut self, input: &mut [f32]) {
+        for sample in input.iter_mut() {
+            let x0 = *sample;
+            let y0 = self.a0 * x0 + self.a1 * self.x1 + self.a2 * self.x2
+                - self.b1 * self.y1
+                - self.b2 * self.y2;
+
+            self.x2 = self.x1;
+            self.x1 = x0;
+            self.y2 = self.y1;
+            self.y1 = y0;
+
+            *sample = y0;
+        }
+    }
+}
+
+// --- Goertzel Filter ---
+struct Goertzel {
+    coeff: f32,
+    history: VecDeque<f32>,
+    window: Vec<f32>,
+}
+
+impl Goertzel {
+    fn new(target_freq: f32, sample_rate: u32, window_size: usize) -> Self {
+        let k = (0.5 + (window_size as f32 * target_freq) / sample_rate as f32) as usize;
+        let omega = (2.0 * std::f32::consts::PI * k as f32) / window_size as f32;
+        let coeff = 2.0 * omega.cos();
+
+        let window = (0..window_size)
+            .map(|i| {
+                0.54 - 0.46 * (2.0 * std::f32::consts::PI * i as f32 / window_size as f32).cos()
+            })
+            .collect();
+
+        Self {
+            coeff,
+            history: VecDeque::with_capacity(window_size),
+            window,
+        }
+    }
+
+    fn run(&self, samples: &[f32]) -> f32 {
+        let mut q0;
+        let mut q1 = 0.0;
+        let mut q2 = 0.0;
+
+        for (i, &sample) in samples.iter().enumerate() {
+            q0 = self.coeff * q1 - q2 + sample * self.window[i];
+            q2 = q1;
+            q1 = q0;
+        }
+
+        q1 * q1 + q2 * q2 - self.coeff * q1 * q2
+    }
+
+    fn process_stream(&mut self, samples: &[f32]) -> Vec<f32> {
+        self.history.extend(samples.iter());
+        let mut power = Vec::new();
+        while self.history.len() >= self.window.len() {
+            let chunk: Vec<f32> = self
+                .history
+                .iter()
+                .take(self.window.len())
+                .copied()
+                .collect();
+            power.push(self.run(&chunk));
+            self.history.pop_front();
+        }
+        power
+    }
+}
+
+// --- Main Decoder ---
+pub struct MorseDecoder {
+    resampler: SincFixedIn<f32>,
+    filter_hp: BiquadFilter,
+    filter_lp: BiquadFilter,
+    audio_buffer: Vec<f32>,
+    target_sample_rate: u32,
+    estimated_pitch: Option<f32>,
+}
+
+impl MorseDecoder {
+    pub fn new(source_sample_rate: u32, target_sample_rate: u32) -> Result<Self> {
+        let resampler = if source_sample_rate != target_sample_rate {
+            let params = InterpolationParameters {
+                sinc_len: 256,
+                f_cutoff: 0.95,
+                interpolation: InterpolationType::Linear,
+                oversampling_factor: 256,
+                window: WindowFunction::BlackmanHarris,
+            };
+            SincFixedIn::new(
+                target_sample_rate as f64 / source_sample_rate as f64,
+                2.0,
+                params,
+                1024, // chunk size
+                1,    // channels
+            )?
+        } else {
+            // Create a dummy resampler if not needed
+            SincFixedIn::new(
+                1.0,
+                1.0,
+                InterpolationParameters {
+                    sinc_len: 2,
+                    ..Default::default()
+                },
+                1024,
+                1,
+            )?
+        };
+
+        Ok(Self {
+            resampler,
+            filter_hp: BiquadFilter::new(FilterType::HighPass, FREQ_MIN_HZ, target_sample_rate),
+            filter_lp: BiquadFilter::new(FilterType::LowPass, FREQ_MAX_HZ, target_sample_rate),
+            audio_buffer: Vec::new(),
+            target_sample_rate,
+            estimated_pitch: None,
+        })
+    }
+
+    pub fn process(&mut self, chunk: &[f32]) -> Result<()> {
+        let waves_in = vec![chunk.to_vec()];
+        let mut resampled = self.resampler.process(&waves_in, None)?;
+
+        let mut audio_chunk = resampled.remove(0);
+        self.filter_hp.process(&mut audio_chunk);
+        self.filter_lp.process(&mut audio_chunk);
+
+        self.audio_buffer.extend(audio_chunk);
+        Ok(())
+    }
+
+    pub fn finalize(&mut self) -> Result<String> {
+        if self.audio_buffer.is_empty() {
+            return Ok(String::new());
+        }
+
+        // 1. Detect Pitch using STFT
+        let pitch = self.detect_pitch_stft()?;
+        log::info!("Estimated pitch: {:.2} Hz", pitch);
+        self.estimated_pitch = Some(pitch);
+
+        // 2. Extract signal power using Goertzel filter
+        let goertzel_window_size = (self.target_sample_rate / 50) as usize; // ~20ms window
+        let mut goertzel_filter =
+            Goertzel::new(pitch, self.target_sample_rate, goertzel_window_size);
+        let power_signal = goertzel_filter.process_stream(&self.audio_buffer);
+
+        // 3. Find optimal WPM and Threshold
+        let (best_wpm, best_threshold) = self.find_best_params(&power_signal)?;
+        log::info!(
+            "Best fit: WPM = {}, Threshold = {:.4}",
+            best_wpm,
+            best_threshold
+        );
+
+        // 4. Decode with optimal parameters
+        let text = self.decode_with_params(&power_signal, best_wpm, best_threshold);
+
+        Ok(text)
+    }
+
+    fn detect_pitch_stft(&self) -> Result<f32> {
+        let fft_size = 2048;
+        let step_size = fft_size / 4;
+
+        let mut planner = FftPlanner::new();
+        let fft = planner.plan_fft_forward(fft_size);
+        let window: Vec<f32> = (0..fft_size)
+            .map(|i| 0.54 - 0.46 * (2.0 * std::f32::consts::PI * i as f32 / fft_size as f32).cos())
+            .collect();
+
+        let mut spectrum_sum = vec![0.0; fft_size / 2];
+        let mut count = 0;
+
+        for chunk in self.audio_buffer.windows(fft_size).step_by(step_size) {
+            let mut buffer: Vec<Complex<f32>> = chunk
+                .iter()
+                .zip(window.iter())
+                .map(|(s, w)| Complex::new(s * w, 0.0))
+                .collect();
+
+            fft.process(&mut buffer);
+
+            for (i, v) in buffer.iter().take(fft_size / 2).enumerate() {
+                spectrum_sum[i] += v.norm_sqr();
+            }
+            count += 1;
+        }
+
+        if count == 0 {
+            bail!("Not enough audio data to detect pitch.");
+        }
+
+        let df = self.target_sample_rate as f32 / fft_size as f32;
+        let mut max_power = 0.0;
+        let mut best_freq = 0.0;
+
+        for i in 0..fft_size / 2 {
+            let freq = i as f32 * df;
+            if freq >= FREQ_MIN_HZ && freq <= FREQ_MAX_HZ {
+                let power = spectrum_sum[i];
+                if power > max_power {
+                    max_power = power;
+                    best_freq = freq;
+                }
+            }
+        }
+
+        Ok(best_freq)
+    }
+
+    fn find_best_params(&self, power_signal: &[f32]) -> Result<(f32, f32)> {
+        let mut best_cost = f32::MAX;
+        let mut best_wpm = 20.0;
+        let mut best_threshold = 0.0;
+
+        let mean_power = power_signal.iter().sum::<f32>() / power_signal.len() as f32;
+
+        for wpm_int in 5..=40 {
+            let wpm = wpm_int as f32;
+            for l in (10..=90).step_by(5) {
+                let threshold = mean_power * (l as f32 / 100.0);
+                let cost = self.calculate_cost(power_signal, wpm, threshold);
+                if cost < best_cost {
+                    best_cost = cost;
+                    best_wpm = wpm;
+                    best_threshold = threshold;
+                }
+            }
+        }
+        Ok((best_wpm, best_threshold))
+    }
+
+    fn calculate_cost(&self, power_signal: &[f32], wpm: f32, threshold: f32) -> f32 {
+        let dot_len_ms = 1200.0 / wpm;
+        let dot_len_samples = (dot_len_ms / 1000.0)
+            * (self.target_sample_rate as f32 / ((self.target_sample_rate / 50) as f32));
+
+        let mut on_intervals = Vec::new();
+        let mut off_intervals = Vec::new();
+        let mut current_len = 0;
+        let mut is_on = power_signal[0] > threshold;
+
+        for &p in power_signal {
+            if (p > threshold) == is_on {
+                current_len += 1;
+            } else {
+                if is_on {
+                    on_intervals.push(current_len);
+                } else {
+                    off_intervals.push(current_len);
+                }
+                is_on = !is_on;
+                current_len = 1;
+            }
+        }
+
+        if on_intervals.is_empty() {
+            return f32::MAX;
+        }
+
+        let cost_on: f32 = on_intervals
+            .iter()
+            .map(|&len| {
+                let cost_dot = (len as f32 / dot_len_samples - 1.0).powi(2);
+                let cost_dash = (len as f32 / dot_len_samples - 3.0).powi(2);
+                cost_dot.min(cost_dash)
+            })
+            .sum();
+
+        cost_on / on_intervals.len() as f32
+    }
+
+    fn decode_with_params(&self, power_signal: &[f32], wpm: f32, threshold: f32) -> String {
+        let dot_len_ms = 1200.0 / wpm;
+        // The power signal has a lower sample rate because of the Goertzel windowing
+        let power_signal_rate =
+            self.target_sample_rate as f32 / (self.target_sample_rate as f32 / 50.0);
+        let dot_len_samples = (dot_len_ms / 1000.0) * power_signal_rate;
+
+        let mut result = String::new();
+        let mut current_letter = String::new();
+
+        let mut current_len = 0;
+        let mut is_on = power_signal[0] > threshold;
+
+        for &p in power_signal.iter().chain(std::iter::once(&0.0)) {
+            // Add sentinel
+            if (p > threshold) == is_on {
+                current_len += 1;
+            } else {
+                let len_norm = current_len as f32 / dot_len_samples;
+                if is_on {
+                    // end of a tone
+                    if (len_norm - 1.0).abs() < (len_norm - 3.0).abs() {
+                        current_letter.push('0'); // dot
+                    } else {
+                        current_letter.push('1'); // dash
+                    }
+                } else {
+                    // end of a space
+                    if len_norm > 2.0 {
+                        // inter-letter space
+                        if let Some(c) = morse_to_char(&current_letter) {
+                            result.push(c);
+                        } else if !current_letter.is_empty() {
+                            result.push('?'); // Unknown character
+                        }
+                        current_letter.clear();
+                        if len_norm > 5.0 {
+                            // word space
+                            result.push(' ');
+                        }
+                    }
+                    // else, it's an inter-element space, do nothing
+                }
+                is_on = !is_on;
+                current_len = 1;
+            }
+        }
+        result
+    }
+}
+
+fn morse_to_char(s: &str) -> Option<char> {
+    match s {
+        "01" => Some('A'),
+        "1000" => Some('B'),
+        "1010" => Some('C'),
+        "100" => Some('D'),
+        "0" => Some('E'),
+        "0010" => Some('F'),
+        "110" => Some('G'),
+        "0000" => Some('H'),
+        "00" => Some('I'),
+        "0111" => Some('J'),
+        "101" => Some('K'),
+        "0100" => Some('L'),
+        "11" => Some('M'),
+        "10" => Some('N'),
+        "111" => Some('O'),
+        "0110" => Some('P'),
+        "1101" => Some('Q'),
+        "010" => Some('R'),
+        "000" => Some('S'),
+        "1" => Some('T'),
+        "001" => Some('U'),
+        "0001" => Some('V'),
+        "011" => Some('W'),
+        "1001" => Some('X'),
+        "1011" => Some('Y'),
+        "1100" => Some('Z'),
+        "01111" => Some('1'),
+        "00111" => Some('2'),
+        "00011" => Some('3'),
+        "00001" => Some('4'),
+        "00000" => Some('5'),
+        "10000" => Some('6'),
+        "11000" => Some('7'),
+        "11100" => Some('8'),
+        "11110" => Some('9'),
+        "11111" => Some('0'),
+        _ => None,
+    }
+}
diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..d288331
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,72 @@
+use anyhow::{bail, Result};
+use clap::Parser;
+use hound::{SampleFormat, WavReader};
+use std::path::PathBuf;
+
+mod decoder;
+use decoder::MorseDecoder;
+
+const TARGET_SAMPLE_RATE: u32 = 12000; // Same as ggmorse's kBaseSampleRate
+const CHUNK_SIZE: usize = 4096;
+
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+struct Cli {
+    /// Path to the input WAV file
+    #[arg(value_name = "WAV_FILE")]
+    wav_file: PathBuf,
+}
+
+fn main() -> Result<()> {
+    // Set up logging. Use `RUST_LOG=info` or `RUST_LOG=debug` to see output.
+    env_logger::init();
+    let cli = Cli::parse();
+
+    log::info!("Opening WAV file: {:?}", cli.wav_file);
+    let mut reader = WavReader::open(cli.wav_file)?;
+    let spec = reader.spec();
+    log::info!("WAV spec: {:?}", spec);
+
+    if spec.sample_format != SampleFormat::Int && spec.sample_format != SampleFormat::Float {
+        bail!(
+            "Unsupported sample format: {:?}. Only 16-bit Int and 32-bit Float are supported.",
+            spec.sample_format
+        );
+    }
+
+    // --- Create the Morse Decoder ---
+    let mut decoder = MorseDecoder::new(spec.sample_rate, TARGET_SAMPLE_RATE)?;
+
+    // --- Read and process the audio in chunks ---
+    let samples_f32: Vec<f32> = if spec.sample_format == SampleFormat::Int {
+        reader
+            .samples::<i16>()
+            .map(|s| s.unwrap() as f32 / 32768.0)
+            .collect()
+    } else {
+        reader.samples::<f32>().map(|s| s.unwrap()).collect()
+    };
+
+    // Convert to mono by averaging channels if necessary
+    let mono_samples: Vec<f32> = if spec.channels > 1 {
+        samples_f32
+            .chunks_exact(spec.channels as usize)
+            .map(|chunk| chunk.iter().sum::<f32>() / spec.channels as f32)
+            .collect()
+    } else {
+        samples_f32
+    };
+
+    // Process the entire audio buffer
+    for chunk in mono_samples.chunks(CHUNK_SIZE) {
+        decoder.process(chunk)?;
+    }
+    
+    // Finalize decoding after all audio is processed
+    let decoded_text = decoder.finalize()?;
+
+    println!("\n--- Decoded Text ---");
+    println!("{}", decoded_text);
+
+    Ok(())
+}
\ No newline at end of file