use std::cell::RefCell;
use std::os::unix::thread;
use std::time::Duration;

use cpal::traits::DeviceTrait;
use cpal::traits::HostTrait;
use cpal::traits::StreamTrait;
use egui::Color32;
use egui_plot::Line;
use egui_plot::PlotPoints;
use egui_plot::Points;
use num::Complex;
use num::complex::ComplexFloat;
use num::traits::sign;
use oxydsp_dsp::blocks::filtering::fir::FirFilter;
use oxydsp_dsp::blocks::filtering::pulse_shaping::PulseShaper;
use oxydsp_dsp::blocks::iq::zero_if::ZeroIf;
use oxydsp_dsp::blocks::math::basic::Multiplier;
use oxydsp_dsp::blocks::synthesis::OscillatorSource;
use oxydsp_dsp::blocks::utilities::adapters::Map;
use oxydsp_dsp::blocks::utilities::adapters::NullSink;
use oxydsp_dsp::blocks::utilities::adapters::Scan;
use oxydsp_dsp::blocks::utilities::channels::RxSource;
use oxydsp_dsp::blocks::utilities::channels::TxSink;
use oxydsp_dsp::blocks::utilities::graph_control::GraphKiller;
use oxydsp_dsp::blocks::utilities::iter::IterSource;
use oxydsp_dsp::filtering::fir::Fir;
use oxydsp_dsp::filtering::history_buf::HistoryBuf;
use oxydsp_dsp::units::DigitalFrequency;
use oxydsp_flowgraph::BlockIO;
use oxydsp_flowgraph::block::Block;
use oxydsp_flowgraph::flowgraph;
use oxydsp_flowgraph::io::In;
use oxydsp_flowgraph::io::Out;
use oxydsp_flowgraph::stream;
use oxydsp_flowgraph::tag::Tags;
use rand::random;

const CARRRIER_FREQ: f64 = 1000.;
const SAMPLE_RATE: usize = 48_000;
const SAMPLE_PER_SYMBOL: usize = 100;

#[derive(BlockIO)]
pub struct CostasLoop
{
    #[input]
    input: In<f32>,

    #[output]
    output: Out<Complex<f32>>,

    center_freq: DigitalFrequency,
    nco: oxydsp_dsp::synthesis::oscillator::Nco<f32>,

    loop_filter: oxydsp_dsp::filtering::fir::FirFilter<f32, f32, f32>,
    low_pass: oxydsp_dsp::filtering::fir::FirFilter<Complex<f32>, Complex<f32>, Complex<f32>>,
}

impl CostasLoop
{
    pub fn new(
        input: In<f32>,
        start_frequency: DigitalFrequency,
        loop_filter: Fir<f32>,
    ) -> (Self, In<Complex<f32>>)
    {
        let (output, iq) = oxydsp_flowgraph::io::stream();
        (
            CostasLoop {
                input,
                output,
                center_freq: start_frequency,
                nco: start_frequency.into(),
                loop_filter: oxydsp_dsp::filtering::fir::FirFilter::new(loop_filter),
                low_pass: oxydsp_dsp::filtering::fir::FirFilter::new(
                    Fir::lowpass(start_frequency, 100).normalized_len(),
                ),
            },
            iq,
        )
    }
}

impl Block for CostasLoop
{
    fn work(&mut self) -> oxydsp_flowgraph::block::BlockResult
    {
        self.output.push_iter(self.input.iter().map(|x| {
            let (signal, tag) = x.into();
            let lo = self.nco.next().unwrap();
            let iq = self
                .low_pass
                .next(Complex::new(lo.re * signal, lo.im * signal));

            let error = iq.re * iq.im.signum() - iq.im * iq.re.signum();
            let correction = self.loop_filter.next(error);
            self.nco.set_frequency(DigitalFrequency::from_rad(
                self.center_freq.as_rad() + correction as f64,
            ));

            (iq, tag).into()
        }));
        oxydsp_flowgraph::block::BlockResult::Ok
    }
}

pub fn main()
{
    demodulator();
    //modulator();
}

pub fn modulator()
{
    let bit_source = (0..8192).map(|_| [random::<bool>(), random::<bool>()]);

    let mut tags = Tags::default();

    let (mut iter_source, bits) = IterSource::new(bit_source);
    let kill_tag = tags.allocate_tag("Kill tag");
    iter_source.tag_last_with(kill_tag.clone());

    let (to_iq, iq) = Map::new(bits, |x| match x
    {
        [true, true] => Complex::new(1., 1.),
        [true, false] => Complex::new(1., -1.),
        [false, true] => Complex::new(-1., 1.),
        [false, false] => Complex::new(-1., -1.),
    });

    let (pulse_shaper, iq) = PulseShaper::new(
        iq,
        Fir::root_raised_cosine(4 * SAMPLE_PER_SYMBOL, 0.5, SAMPLE_PER_SYMBOL),
        SAMPLE_PER_SYMBOL,
    );
    let (carrier_oscillator, carrier) = OscillatorSource::new(
        DigitalFrequency::from_time_frequency(CARRRIER_FREQ, SAMPLE_RATE as f64).into(),
    );
    let (mixer, passband) = Multiplier::new(iq, carrier);

    let (tx, rx) = std::sync::mpsc::channel::<Complex<f32>>();
    let (graph_killer, passband) = GraphKiller::new(passband, kill_tag);
    let tx_sink = TxSink::new(passband, tx);

    let graph = flowgraph![
        iter_source,
        to_iq,
        pulse_shaper,
        carrier_oscillator,
        mixer,
        graph_killer,
        tx_sink
    ];
    graph.run(1);

    let spec = hound::WavSpec {
        channels: 1,
        sample_rate: SAMPLE_RATE as u32,
        bits_per_sample: 16,
        sample_format: hound::SampleFormat::Int,
    };
    let mut writer = hound::WavWriter::create("mod.wav", spec).unwrap();
    for iq in rx.iter()
    {
        let amplitude = 0.5 * i16::MAX as f32;
        writer.write_sample((iq.re * amplitude) as i16).unwrap();
    }
    writer.finalize().unwrap();
}

pub fn demodulator()
{
    let (audio_tx, audio_rx) = std::sync::mpsc::channel();
    let (rx_source, signal) = RxSource::new(audio_rx);

    let (downconverter, iq) = CostasLoop::new(
        signal,
        DigitalFrequency::from_time_frequency(CARRRIER_FREQ, SAMPLE_RATE as f64),
        Fir::proportional_integral(100, 0.000, 0.0000),
    );

    //let agc_filter = Fir::proportional_integral(100, 0.1, 0.001);
    // let (agc, iq) = Scan::new(iq, 1., |gain, iq|
    // {
    //     let mu = 0.1;
    //     let mag = iq.abs();
    //     *gain += mu * (1. - mag * *gain);
    //
    //     iq * *gain
    // });

    let (matched_filter, iq) = FirFilter::new(
        iq,
        Fir::<f32>::root_raised_cosine(4 * SAMPLE_PER_SYMBOL, 0.5, SAMPLE_PER_SYMBOL)
            .normalized_sqr(),
    );

    let (eye_i_tx, eye_i_rx) = std::sync::mpsc::channel::<Vec<f32>>();
    let (eye_q_tx, eye_q_rx) = std::sync::mpsc::channel::<Vec<f32>>();
    let (constellation_tx, constellation_rx) = std::sync::mpsc::channel();

    // let (debug, iq) = Scan::new(
    //     iq,
    //     (
    //         HistoryBuf::new(0., SAMPLE_PER_SYMBOL * 2),
    //         HistoryBuf::new(0., SAMPLE_PER_SYMBOL * 2),
    //         0usize,
    //     ),
    //     move |(buf_i, buf_q, counter), x| {
    //         buf_i.push(x.re);
    //         buf_q.push(x.im);
    //         let _ = constellation_tx.send(x);
    //         *counter += 1;
    //         if *counter >= SAMPLE_PER_SYMBOL * 2
    //         {
    //             let _ = eye_i_tx.send(buf_i.as_slice().iter().copied().collect::<Vec<_>>());
    //             let _ = eye_q_tx.send(buf_q.as_slice().iter().copied().collect::<Vec<_>>());
    //             *counter = 0;
    //         }
    //         x
    //     },
    // );
    let tx_sink = TxSink::new(iq, constellation_tx);

    let host = cpal::default_host();
    let device = host
        .default_input_device()
        .expect("no output device available");
    let mut supported_configs_range = device
        .supported_input_configs()
        .expect("error while querying configs");
    let supported_config = supported_configs_range
        .next()
        .expect("no supported config?!")
        .with_sample_rate(SAMPLE_RATE as u32);
    let _stream = device
        .build_input_stream(
            &supported_config.into(),
            move |data: &[f32], _: &cpal::InputCallbackInfo| {
                for x in data
                {
                    //let _ = audio_tx.send(*x * 10.);
                    let _ = audio_tx.send(random::<f32>());
                    //let _ = audio_tx.send(Complex::new(random::<f32>(), random::<f32>()));
                }
            },
            move |_err| {},
            None, // None=blocking, Some(Duration)=timeout
        )
        .unwrap();

    let graph = flowgraph![
        rx_source,
        downconverter,
        // agc,
        matched_filter,
        //debug,
        //null_sink
        tx_sink
    ];
    graph.run(6);

    let mut eye_i_history = HistoryBuf::new(vec![], 200);
    let mut eye_q_history = HistoryBuf::new(vec![], 200);
    let mut constellation = HistoryBuf::new(Complex::new(0., 0.), 5000);
    eframe::run_simple_native("Window", Default::default(), move |ctx, _frame| {
        for eye in eye_i_rx.try_iter().take(200)
        {
            eye_i_history.push(eye);
        }
        for eye in eye_q_rx.try_iter().take(200)
        {
            eye_q_history.push(eye);
        }
        for point in constellation_rx.try_iter().take(5000)
        {
            constellation.push(point);
        }

        egui::CentralPanel::default().show(ctx, |ui| {
            egui_plot::Plot::new("plot")
                .data_aspect(1.)
                .show(ui, |plot_ui| {
                    plot_ui.points(
                        Points::new(
                            "Constellation",
                            constellation
                                .as_slice()
                                .iter()
                                .map(|point| [point.re as f64, point.im as f64])
                                .collect::<PlotPoints>(),
                        )
                        .color(Color32::YELLOW.gamma_multiply_u8(70)),
                    );

                    for (eye_i, eye_q) in eye_i_history
                        .as_slice()
                        .iter()
                        .zip(eye_q_history.as_slice().iter())
                    {
                        plot_ui.line(
                            Line::new(
                                "In-Phase",
                                eye_i
                                    .iter()
                                    .enumerate()
                                    .map(|(i, x)| {
                                        [i as f64 / (SAMPLE_PER_SYMBOL as f64 * 2.) + 1., *x as f64]
                                    })
                                    .collect::<Vec<_>>(),
                            )
                            .color(Color32::RED),
                        );

                        plot_ui.line(
                            Line::new(
                                "Quadrature",
                                eye_q
                                    .iter()
                                    .enumerate()
                                    .map(|(i, x)| {
                                        [*x as f64, i as f64 / (SAMPLE_PER_SYMBOL as f64 * 2.) - 2.]
                                    })
                                    .collect::<Vec<_>>(),
                            )
                            .color(Color32::GREEN),
                        );

                        plot_ui.points(
                            Points::new(
                                "Constellation",
                                eye_i
                                    .iter()
                                    .zip(eye_q.iter())
                                    .skip(SAMPLE_PER_SYMBOL / 2)
                                    .step_by(SAMPLE_PER_SYMBOL)
                                    .map(|(i, q)| [*i as f64, *q as f64])
                                    .collect::<PlotPoints>(),
                            )
                            .color(Color32::GREEN)
                            .radius(1.5),
                        );
                    }
                });
        });
        ctx.request_repaint();
    })
    .unwrap();
}

pub fn to_bits(n: u8) -> [bool; 8]
{
    [
        (n & 1) == 1,
        (n >> 1) & 1 == 1,
        (n >> 2) & 1 == 1,
        (n >> 3) & 1 == 1,
        (n >> 4) & 1 == 1,
        (n >> 5) & 1 == 1,
        (n >> 6) & 1 == 1,
        (n >> 7) & 1 == 1,
    ]
}

pub fn from_bits(n: [bool; 8]) -> u8
{
    (n[0] as u8)
        | ((n[1] as u8) << 1)
        | ((n[2] as u8) << 2)
        | ((n[3] as u8) << 3)
        | ((n[4] as u8) << 4)
        | ((n[5] as u8) << 5)
        | ((n[6] as u8) << 6)
        | ((n[7] as u8) << 7)
}
