use std::collections::VecDeque;
use std::env::args;
use std::f32::consts::PI;
use std::fs::File;
use std::io::Write;
use std::net::Ipv4Addr;
use std::net::UdpSocket;
use std::os::unix::thread;
use std::sync::mpsc::channel;
use std::sync::mpsc::sync_channel;
use std::time::Duration;

use cpal::traits::DeviceTrait;
use cpal::traits::HostTrait;
use eframe::NativeOptions;
use egui::Color32;
use egui_plot::Line;
use egui_plot::PlotPoints;
use egui_plot::Points;
use num::Complex;
use num::Integer;
use oxydsp_dsp::blocks::filtering::fir::FirFilter;
use oxydsp_dsp::blocks::iq::zero_if::ZeroIf;
use oxydsp_dsp::blocks::math::basic;
use oxydsp_dsp::blocks::math::basic::Multiplier;
use oxydsp_dsp::blocks::synthesis::OscillatorSource;
use oxydsp_dsp::blocks::ted::early_late::EarlyLateGate;
use oxydsp_dsp::blocks::utilities::adapters::FlatMap;
use oxydsp_dsp::blocks::utilities::adapters::Map;
use oxydsp_dsp::blocks::utilities::adapters::MapResultTagged;
use oxydsp_dsp::blocks::utilities::adapters::NullSink;
use oxydsp_dsp::blocks::utilities::adapters::Repeat;
use oxydsp_dsp::blocks::utilities::adapters::Scan;
use oxydsp_dsp::blocks::utilities::adapters::ScanTagged;
use oxydsp_dsp::blocks::utilities::adapters::Tee;
use oxydsp_dsp::blocks::utilities::channels::RxSource;
use oxydsp_dsp::blocks::utilities::channels::TxSink;
use oxydsp_dsp::blocks::utilities::iter::IterSource;
use oxydsp_dsp::filtering::fir::Fir;
use oxydsp_dsp::units::DigitalFrequency;
use oxydsp_flowgraph::block::BlockResult;
use oxydsp_flowgraph::flowgraph;
use oxydsp_flowgraph::io::AnonymousIn;
use oxydsp_flowgraph::io::In;
use oxydsp_flowgraph::tag::Tags;
use rand::random;

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

fn main()
{
    let args = std::env::args();
    if args.len() == 1
    {
        demodulator();
    }
    else
    {
        modulator();
    }
    println!("Hello, world!");
}

fn demodulator()
{
    let (signal_tx, signal_rx) = channel();
    // std::thread::spawn(move || {
    //     let udp_socket = UdpSocket::bind("0.0.0.0:25565").unwrap();
    //     let mut buffer = [0u8; 4096];
    //     while let Ok(size) = udp_socket.recv(&mut buffer)
    //     {
    //         let read = &mut buffer[..size];
    //         for bytes in read.chunks(4)
    //         {
    //             if bytes.len() == 4
    //             {
    //                 let _ = signal_tx
    //                     .send(f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]));
    //             }
    //         }
    //     }
    // });

    // Open output device
    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.iter()
                {
                    let _ = signal_tx.send(*x);
                }
            },
            move |err| panic!(),
            None, // None=blocking, Some(Duration)=timeout
        )
        .unwrap();

    let (signal_source, signal) = IterSource::new(signal_rx.into_iter());
    let (mut zero_if, baseband): (_, In<Complex<f32>>) = ZeroIf::new(
        signal,
        DigitalFrequency::from_time_frequency(CARRIER, SAMPLE_RATE as f64).into(),
    );
    zero_if.set_fir(Fir::lowpass(
        DigitalFrequency::from_time_frequency(CARRIER + 100., SAMPLE_RATE as f64),
        100,
    ));

    // Pulse shaping
    // gaussian fir
    let fir = Fir((0..SAMPLE_PER_SYMBOL)
        .map(|x| gaussian(0.3, x as f32 / SAMPLE_PER_SYMBOL as f32))
        .collect())
    .normalized();

    // RRC fir
    let rrc_symbol_count = 4;
    let rrc_fir_length = SAMPLE_PER_SYMBOL * rrc_symbol_count;
    let fir = Fir((0..rrc_fir_length)
        .map(|x| {
            let centered = oxydsp_dsp::map(
                x as f32,
                0.,
                rrc_fir_length as f32,
                -(rrc_symbol_count as f32) * 0.5,
                rrc_symbol_count as f32 * 0.5,
            );
            root_raised_cosine(centered, 1., 1.)
        })
        .collect())
    .normalized();

    let (tee, b_a, b_b) = Tee::new(baseband);

    let (delay, b_b) = Scan::new(
        b_b,
        VecDeque::from(vec![Complex::<f32>::new(0., 0.); SAMPLE_PER_SYMBOL]),
        |history, x| {
            history.push_back(x);
            history.pop_front().unwrap().conj()
        },
    );

    let (mult, baseband) = Multiplier::new(b_a, b_b);
    let (to_arg, baseband) = Map::new(baseband, |x| (x * Complex::<f32>::new(0., -1.)).arg());
    let (phase_filter, baseband) = FirFilter::new(baseband, fir);

    let (eye_tx, eye_rx) = channel::<Vec<f32>>();
    let (eye, baseband) = Scan::new(baseband, VecDeque::new(), move |state, x| {
        if state.len() >= SAMPLE_PER_SYMBOL * 2
        {
            //state.pop_front();
            eye_tx.send(Vec::from(state.clone()));
            state.clear();
        }
        state.push_back(x);
    });
    let null_sink = NullSink::new(baseband);

    let graph = flowgraph![
        signal_source,
        zero_if,
        null_sink,
        eye,
        phase_filter,
        tee,
        to_arg,
        delay,
        mult
    ];
    File::create("demodulator.dot")
        .unwrap()
        .write_all(graph.get_dot().as_bytes())
        .unwrap();
    graph.run();

    let mut constellation = VecDeque::new();
    eframe::run_simple_native("Plot", NativeOptions::default(), move |ctx, _frame| {
        while let Ok(sample) = eye_rx.try_recv()
        {
            if constellation.len() >= 300
            {
                let _ = constellation.pop_back();
            }

            constellation.push_front(sample);
        }

        egui::CentralPanel::default().show(ctx, |ui| {
            egui_plot::Plot::new("hello").show(ui, |plot_ui| {
                for eye in constellation.iter()
                {
                    plot_ui.line(
                        Line::new(
                            "eye",
                            eye.iter()
                                .enumerate()
                                .map(|(i, e)| [i as f64, *e as f64])
                                .collect::<PlotPoints>(),
                        )
                        .color(Color32::GREEN),
                    );
                }
                // plot_ui.points(
                //     Points::new(
                //         "constellation",
                //         constellation
                //             .iter()
                //             .map(|s| [s.re as f64, s.im as f64])
                //             .collect::<PlotPoints>(),
                //     )
                //     .id("constellation")
                //     .color(Color32::GREEN),
                // );
            });
            ctx.request_repaint();
        });
    })
    .unwrap();
}

fn modulator()
{
    let (data_tx, data_rx) = channel();

    std::thread::spawn(move || {
        loop
        {
            let mut str = String::new();
            let input = std::io::stdin().read_line(&mut str).unwrap();

            for bit in str.as_bytes().iter().copied().flat_map(to_bits)
            {
                let _ = data_tx.send(bit);
            }
        }
    });

    let mut tags = Tags::new();
    let (bit_source, bits): (_, In<bool>) = RxSource::new(data_rx);
    //let last_tag = tags.allocate_tag("finished");

    let (phase_map, phase) = Scan::new(bits, Complex::<f32>::new(1., 0.), |state, bit| {
        if bit
        {
            *state *= Complex::<f32>::new(-1., 0.);
        }
        *state
    });

    // Convert to pulse train
    let (repeater, phase) = FlatMap::new(phase, |x| {
        let mut v = vec![Complex::<f32>::new(0., 0.); SAMPLE_PER_SYMBOL - 1];
        //let mut v = vec![x; SAMPLE_PER_SYMBOL - 1];
        v.push(x);
        v
    });

    // RRC fir
    let rrc_symbol_count = 4;
    let rrc_fir_length = SAMPLE_PER_SYMBOL * rrc_symbol_count;
    let fir = Fir((0..rrc_fir_length)
        .map(|x| {
            let centered = oxydsp_dsp::map(
                x as f32,
                0.,
                rrc_fir_length as f32,
                -(rrc_symbol_count as f32) * 0.5,
                rrc_symbol_count as f32 * 0.5,
            );
            root_raised_cosine(centered, 1., 1.)
        })
        .collect());

    let (phase_filter, phase) = FirFilter::new(phase, fir);

    let (oscillator, passband) = OscillatorSource::<f32>::new(
        DigitalFrequency::from_time_frequency(CARRIER, SAMPLE_RATE as f64).into(),
    );
    let (multiplier, passband) = Multiplier::new(passband, phase);

    let reverb_length = 200;
    // let (reverb, passband) = FirFilter::new(
    //     passband,
    //     Fir((0..reverb_length)
    //         .map(|x| (-20. * (x as f32) / (reverb_length as f32)).exp())
    //         .collect())
    //     .normalized(),
    // );

    let (awgn, passband) = Map::new(passband, |x| {
        x + Complex::<f32>::new(2. * (random::<f32>() - 0.5), 2. * (random::<f32>() - 0.5)) * 0.0
    });

    let (udp_map, passband) = Scan::new(
        passband,
        UdpSocket::bind("0.0.0.0:0").unwrap(),
        |sckt, sample| {
            std::thread::sleep(Duration::from_micros(15));
            sckt.send_to(&(sample.re).to_le_bytes(), "127.0.0.1:25565")
                .unwrap();
            sample
        },
    );
    let (audio_tx, audio_rx) = channel();
    let tx_sink = TxSink::new(passband, audio_tx);

    let graph = flowgraph![
        bit_source,
        phase_map,
        repeater,
        oscillator,
        multiplier,
        tx_sink, //reverb,
        awgn,
        udp_map,
        phase_filter
    ];

    // Open output device
    let host = cpal::default_host();
    let device = host
        .default_output_device()
        .expect("no output device available");
    let mut supported_configs_range = device
        .supported_output_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_output_stream(
            &supported_config.into(),
            move |data: &mut [f32], _: &cpal::OutputCallbackInfo| {
                for x in data.iter_mut()
                {
                    if let Ok(y) = audio_rx.try_recv()
                    {
                        *x = y.re * 0.01;
                    }
                    else
                    {
                        *x = 0.;
                    }
                }
            },
            move |err| panic!(),
            None, // None=blocking, Some(Duration)=timeout
        )
        .unwrap();
    let _ = graph.run().join();
}

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)
}

pub fn gaussian(sigma: f32, t: f32) -> f32
{
    let sq = (t - 0.5) / sigma;
    (-sq * sq).exp()
}

pub fn root_raised_cosine(t: f32, beta: f32, ts: f32) -> f32
{
    let eps = 1e-8;

    if t.abs() < eps
    {
        // t = 0 special case
        return (1.0 / ts.sqrt()) * (1.0 + beta * (4.0 / PI - 1.0));
    }

    if beta > 0.0 && (t.abs() - ts / (4.0 * beta)).abs() < eps
    {
        // t = ±T / (4β) special case
        let term1 = (1.0 + 2.0 / PI) * (PI / (4.0 * beta)).sin();
        let term2 = (1.0 - 2.0 / PI) * (PI / (4.0 * beta)).cos();
        return (beta / (ts.sqrt() * 2.0_f32.sqrt())) * (term1 + term2);
    }

    // General case
    let numerator = (PI * t * (1.0 - beta) / ts).sin()
        + 4.0 * beta * t / ts * (PI * t * (1.0 + beta) / ts).cos();

    let denominator = PI * t * (1.0 - (4.0 * beta * t / ts).powi(2)) / ts;

    (1.0 / ts.sqrt()) * (numerator / denominator)
}
