// 2-FSK Modulator

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

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

pub struct BFSKMod<'a, T: Iterator<Item = bool>> {
    samples_per_bit: u32,
    bandwidth: 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, bandwidth: f32, bit_stream: &'a mut T) -> Self {
        BFSKMod {
            samples_per_bit,
            bandwidth,
            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.bandwidth / 2.0
            } else {
                -self.bandwidth / 2.0
            };
            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: Box<dyn DFT>,
}

impl BFSKDem {
    pub fn new(samples_per_bit: u32, deviation: f32) -> Self {
        BFSKDem {
            samples_per_bit,
            deviation,
            sample_index: 0,
            fft: fft::create_fft(samples_per_bit as usize, fft::FFTDirection::Forward),
        }
    }

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

        self.fft
            .get_input()
            .iter_mut()
            .enumerate()
            .for_each(|(i, x)| *x = baseband[i]);
        self.fft.execute(windows::rectanguar);

        let bin_id = map(
            self.deviation,
            0.,
            PI,
            0.,
            (self.samples_per_bit / 2) as f32,
        )
        .floor() as i32;

        let bin_width = 5;

        let mut positive_energy = 0.0;
        for i in (bin_id - bin_width)..(bin_id + bin_width) {
            if i >= 0 && i < self.samples_per_bit as i32 {
                positive_energy += self.fft.get_output()[i as usize].mag();
            }
        }

        let mut negative_energy = 0.0;
        for i in (self.samples_per_bit as i32 - bin_id - bin_width)
            ..(self.samples_per_bit as i32 - bin_id + bin_width)
        {
            if i >= 0 && i < self.samples_per_bit as i32 {
                negative_energy += self.fft.get_output()[i as usize].mag();
            }
        }

        return positive_energy < negative_energy;
    }
}
