use std::iter::Sum;

use num::Float;
use num::NumCast;
use oxydsp_flowgraph::BlockIO;
use oxydsp_flowgraph::block::Block;
use oxydsp_flowgraph::io::In;
use oxydsp_flowgraph::io::Out;
use oxydsp_flowgraph::tag::Tag;
use oxydsp_flowgraph::tag::TagKey;

use crate::filtering::fir::Fir;
use crate::filtering::fir::FirFilter;
use crate::filtering::history_buf::HistoryBuf;

#[derive(BlockIO)]
pub struct EarlyLateGate<T: Float + Send + Sync + Sum + Clone + NumCast + 'static>
{
    #[input]
    input: In<T>,

    #[output]
    output: Out<T>,

    symbol_length: usize,

    // Window looking at symbol_length samples at a time
    window: HistoryBuf<T>,

    // The current location of the window, in relation to the last sample
    window_location: usize,

    window_center: usize,

    // The next window location, in relation to the last sample such that the window is centered on
    // a symbol center (hopefully)
    next_sample: f32,
    loop_filter: FirFilter<T, T, T>,

    symbol_tag_key: TagKey<T>,
}

impl<T> EarlyLateGate<T>
where
    T: Float + Sum + Clone + 'static + Send + Sync + NumCast + Default,
{
    pub fn new(
        input: In<T>,
        loop_filter: Fir<T>,
        symbol_length: usize,
        symbol_tag_key: TagKey<T>,
    ) -> (Self, In<T>)
    {
        let (output, samples) = oxydsp_flowgraph::io::stream();
        (
            Self {
                input,
                output,
                window: HistoryBuf::new(Default::default(), symbol_length),
                symbol_length,
                window_location: 0,
                window_center: symbol_length / 2,
                next_sample: symbol_length as f32, // We assume that the first symbol is 1.5 windows into
                // the stream
                loop_filter: FirFilter::new(loop_filter),
                symbol_tag_key,
            },
            samples,
        )
    }
}

impl<'view, T> Block for EarlyLateGate<T>
where
    T: Float + Sum + Clone + 'static + Send + Sync + NumCast,
{
    fn work(&mut self) -> oxydsp_flowgraph::block::BlockResult
    {
        self.output.push_iter(self.input.iter().map(|input| {
            // Bring new sample in
            self.window.push(input.0);
            self.window_location += 1;

            let sample = self.window.as_slice()[self.window_center];
            let mut tag = None;
            if self.window_location >= self.next_sample as usize
            {
                let early_index = self.window_center - (0.25 * self.symbol_length as f32) as usize;
                let late_index = self.window_center + (0.25 * self.symbol_length as f32) as usize;

                let early_sample = self.window.as_slice()[early_index];
                let late_sample = self.window.as_slice()[late_index];

                let error = (late_sample - early_sample) * sample;
                let correction = self.loop_filter.next(error);

                // Figure out next sample location
                self.next_sample +=
                    (self.symbol_length as f32 + correction.to_f32().unwrap()).max(0.);

                // Turn everything back relative to current sample
                self.next_sample -= self.window_location as f32;
                self.window_location = 0;

                tag = Some(Tag::with_entry(self.symbol_tag_key.clone(), error));
            }

            (sample, tag).into()
        }));

        oxydsp_flowgraph::block::BlockResult::Ok
    }
}
