// 2-FSK Modulator

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

use crate::complex::{Complex, Complex32};
use crate::fft::{self, FFT, FFTDirection};
use crate::{units, 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())
    }
}

impl<'a, T: Iterator<Item = bool>> Iterator for BFSKMod<'a, T> {
    type Item = Complex32;

    fn next(&mut self) -> Option<Self::Item> {
        self.step_modulate()
    }
}

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

    // Correlators
    pos_correlator: Box<[Complex32]>,
    neg_correlator: Box<[Complex32]>,
}

impl BFSKDem {
    pub fn new(samples_per_bit: u32, deviation: f32) -> Self {

        let mut nco_pos = Nco::new(deviation);
        let mut nco_neg = Nco::new(-deviation);

        BFSKDem {
            samples_per_bit,
            deviation,
            sample_index: 0,

            pos_correlator: (0..samples_per_bit).map(|_| {nco_pos.step(); nco_pos.cexp()}).collect(),
            neg_correlator: (0..samples_per_bit).map(|_| {nco_neg.step(); nco_neg.cexp()}).collect()
        }
    }

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

        
        //perform correlation
        let positive_energy: Complex32 = self.pos_correlator.iter().zip(baseband.iter())
            .enumerate()
            .map(|(i, (a, b))| *a * *b * windows::blackmann(i as f32 / self.samples_per_bit as f32)).sum();
        let negative_energy: Complex32 = self.neg_correlator.iter().zip(baseband.iter())
            .enumerate()
            .map(|(i, (a, b))| *a * *b * windows::blackmann(i as f32 / self.samples_per_bit as f32)).sum();

        positive_energy.mag().powi(2) < negative_energy.mag().powi(2)
    }
}

/*
// 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::blackmann),
            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()
    }
}
*/
