// 2-FSK Modulator

use std::f32::consts::PI;

use crate::complex::{Complex, Complex32};
use crate::fft::{self, FFT, FFTDirection, windows};
use crate::map;
use crate::nco::Nco;

pub struct BFSKMod<'a, T: Iterator<Item = bool>> {
    samples_per_bit: u32,
    deviation: f32,
    bit_stream: &'a mut T,

    // State
    oscillator: Nco,
    sample_index: u32,
}

impl<'a, T> BFSKMod<'a, T>
where
    T: Iterator<Item = bool>,
{
    pub fn new(samples_per_bit: u32, deviation: f32, bit_stream: &'a mut T) -> Self {
        BFSKMod {
            samples_per_bit,
            deviation,
            oscillator: Nco::new(0.0),
            bit_stream,
            sample_index: samples_per_bit,
        }
    }

    pub fn step_modulate(&mut self) -> Option<Complex<f32>> {
        if self.sample_index == self.samples_per_bit {
            self.sample_index = 0;
            let bit = self.bit_stream.next()?;

            let frequency = if bit { self.deviation } else { -self.deviation };
            self.oscillator.set_frequency(frequency);
        }

        self.sample_index += 1;
        self.oscillator.step();
        Some(self.oscillator.cexp())
    }
}

// FSK Demodulator (dumb non coherent + no symbol timing recovery)
pub struct BFSKDem {
    samples_per_bit: u32,
    deviation: f32,
    // State
    sample_index: u32,
    fft: FFT,

    bin_pos: usize,
    bin_neg: usize,
}

impl BFSKDem {
    pub fn new(samples_per_bit: u32, deviation: f32) -> Self {
        // Calculate bin locations :
        let bin_index = map(deviation, 0., 2. * PI, 0., samples_per_bit as f32).round() as u32;

        println!("bin_index: {bin_index}");

        BFSKDem {
            samples_per_bit,
            deviation,
            sample_index: 0,
            fft: FFT::new(samples_per_bit as usize, windows::bartlett),
            bin_pos: bin_index as usize,
            bin_neg: (samples_per_bit - bin_index - 1) as usize, // -deviation = negative frequency = upper half
        }
    }

    pub fn demod(&mut self, baseband: &[Complex32]) -> bool {
        assert!(baseband.len() >= self.samples_per_bit as usize);
        self.fft.execute(baseband);

        let positive_energy = self.fft.get_output()[self.bin_pos];
        let negative_energy = self.fft.get_output()[self.bin_neg];

        positive_energy.mag() > negative_energy.mag()
    }
}