#![allow(dead_code)]

use std::{
    collections::VecDeque,
    f32::consts::PI,
    fs::File,
    i16,
    io::{self, Read, Write},
    ops::{Add, Div, Mul, Sub},
    os::unix::thread,
    sync::{
        Arc,
        atomic::{AtomicU32, Ordering},
        mpsc::{self, Receiver, Sender, TryRecvError, channel, sync_channel},
    },
    time::Duration,
};

mod bfsk;
mod complex;
pub mod fft;
mod filtering;
mod iq;
mod nco;
mod signal;
mod units;
mod windows;

use bfsk::BFSKMod;
use complex::Complex;
use complex::Complex32;
use cpal::{
    SampleRate,
    traits::{DeviceTrait, HostTrait, StreamTrait},
};
use fft::DFTAlgorithm;
use nco::Nco;

use eframe::egui::{self, Color32, Context, Vec2b, debug_text::print, decode_animated_image_uri};
use egui_plot::{self, Bar, BarChart, Legend, Line, LineStyle, Plot, PlotPoints, VLine};
use plotters::style::Color;
use rand::{Rng, seq::index::sample};

use crate::{
    bfsk::BFSKDem,
    fft::{FFT, dft::NaiveDFT},
    filtering::{
        fir::FIRFilter,
        impulse_response::design::{self, frequency_response, ir_from_transfer_function},
    },
    iq::IQSampler,
    units::frequency::{self, hz_to_rad_per_sample},
};

// Utilities
fn map<T>(input: T, in_min: T, in_max: T, out_min: T, out_max: T) -> T
where
    T: Clone + Add<Output = T> + Mul<Output = T> + Sub<Output = T> + Div<Output = T>,
{
    ((input - in_min.clone()) / (in_max - in_min)) * (out_max - out_min.clone()) + out_min
}
const BAUD_RATE: u32 = 1200;

fn main() {
    modulate();
    //demodulate();
    //return;
    // Set up CPAL
    let host = cpal::default_host();

    let device = host
        .default_input_device()
        .expect("No input device available");
    let mut config = device
        .supported_input_configs()
        .unwrap()
        .next()
        .unwrap()
        .with_sample_rate(SampleRate(48000));

    // Channel to move samples from callback to main thread
    let (tx, rx) = sync_channel::<f32>(4096);

    // Build input stream
    /*
    let stream = device
        .build_input_stream(
            &config.into(),
            move |data: &[f32], _| {
                for x in data.iter() {
                    let _ = tx.send(*x * 10.); // non-blocking send
                }
            },
            move |err| eprintln!("Stream error: {}", err),
            None,
        )
        .unwrap();
    stream.play().unwrap();
    */

    let (eye_tx, eye_rx) = mpsc::channel::<Vec<f32>>();
    let (ctx_tx, ctx_rx) = mpsc::channel::<Arc<egui::Context>>();
    std::thread::spawn(move || demodulator(rx, eye_tx, ctx_rx));

    std::thread::spawn(move || {
        let spec = hound::WavSpec {
            channels: 1,
            sample_rate: 48000,
            bits_per_sample: 16,
            sample_format: hound::SampleFormat::Int,
        };

        let mut writer = hound::WavWriter::create("audio/noised.wav", spec).unwrap();

        let mut reader = hound::WavReader::open("audio/modulated.wav").unwrap();
        let mut rand = rand::rng();
        let samples = reader.samples::<i16>();
        for x in samples {
            let noise = rand.random::<f32>() * 2. - 1.;
            let sample = x.unwrap() as f32 / i16::MAX as f32 + noise * 0.9;
            let _ = tx.send(sample);
            writer
                .write_sample((sample * i16::MAX as f32) as i16)
                .unwrap();
            std::thread::sleep(Duration::from_micros(21));
        }
        writer.finalize().unwrap();
    });

    let native_options = eframe::NativeOptions::default();
    let _ = eframe::run_native(
        "Egui",
        native_options,
        Box::new(|cc| Ok(Box::new(EguiApp::new(cc, eye_rx, ctx_tx)))),
    );
}

// Early late gate
struct ELGate {
    samples_per_symbol: f32,
    buffer: VecDeque<f32>, // Store baseband, matched filtered samples,

    loop_filter: FIRFilter,
    delta: f32,
    next_sample: f32,
    current_position: f32,
}

impl ELGate {
    pub fn new(samples_per_symbol: f32, loop_filter: FIRFilter) -> Self {
        Self {
            samples_per_symbol,
            loop_filter,
            buffer: VecDeque::with_capacity(2 * samples_per_symbol.ceil() as usize),
            delta: 0.5,
            next_sample: samples_per_symbol,
            current_position: 0.,
        }
    }

    pub fn next(&mut self, sample: f32) -> Option<bool> {
        Some(self.next_eye(sample)?.0) // Ignore eye
    }

    pub fn next_eye(&mut self, sample: f32) -> Option<(bool, Vec<f32>)> {
        self.buffer.push_front(sample);
        self.current_position += 1.;
        if self.current_position >= self.next_sample {
            // Sample center, early late
            let early_id = (self.samples_per_symbol / 2. + self.samples_per_symbol * self.delta)
                .floor()
                .min(self.buffer.len() as f32 - 1.) as usize;
            let late_id = (self.samples_per_symbol / 2. - self.samples_per_symbol * self.delta)
                .floor()
                .max(0.) as usize;
            let sample_id = (self.samples_per_symbol / 2.) as usize;

            let early_sample = self.buffer[early_id];
            let late_sample = self.buffer[late_id];
            let sample = self.buffer[sample_id];

            let error = (early_sample - late_sample) * sample;

            // Remove until only current sample is in buffer (the next sample might need data from
            // current sample if error advances sample position)
            while self.buffer.len() > self.samples_per_symbol as usize {
                let _ = self.buffer.pop_back();
            }

            // Predict next length based on error
            self.current_position = self.next_sample - self.next_sample.floor();
            self.next_sample = self.samples_per_symbol - self.loop_filter.next_real(error);

            Some((sample > 0., Vec::from(self.buffer.clone())))
        } else {
            None
        }
    }
}

fn demodulator(
    rx: Receiver<f32>,
    eye_sender: Sender<Vec<f32>>,
    ctx_rx: Receiver<Arc<egui::Context>>,
) {
    // Wait for egui context to request redraw
    let ctx = ctx_rx.recv().unwrap();

    // Modulation parameters
    let frequency = 1700.;
    let deviation = 500.;

    // Data parameters
    let sample_rate = 48000;
    let baud_rate = BAUD_RATE;
    let sample_per_symbol = sample_rate / baud_rate;

    let mut iq_sampler = IQSampler::new(hz_to_rad_per_sample(frequency, sample_rate as f32));

    // Corellators
    let mut nco_pos = Nco::new(hz_to_rad_per_sample(deviation, sample_rate as f32));
    let mut nco_neg = Nco::new(hz_to_rad_per_sample(-deviation, sample_rate as f32));
    let corellator_length = (sample_per_symbol as f32 * 1.5) as usize;
    let corellator_pos = (0..corellator_length)
        .map(|i| {
            nco_pos.step();
            nco_pos.cexp() * windows::blackmann(i as f32 / (corellator_length as f32))
            //nco_pos.cexp()
        })
        .collect::<Vec<_>>();
    let corellator_neg = (0..corellator_length)
        .map(|i| {
            nco_neg.step();
            nco_neg.cexp() * windows::blackmann(i as f32 / (corellator_length as f32))
            //nco_neg.cexp()
        })
        .collect::<Vec<_>>();
    let mut matched_filter_pos = FIRFilter::new(&corellator_pos);
    let mut matched_filter_neg = FIRFilter::new(&corellator_neg);
    matched_filter_pos.normalize_sum();
    matched_filter_neg.normalize_sum();

    let loop_p = 0.3;
    let loop_i = 0.1;
    let mut matched_filter = FIRFilter::new(&[Complex32::new(1., 0.); 40]);
    matched_filter.normalize_sum();

    let mut loop_filter_ir = vec![Complex32::new(loop_i, 0.); 30];
    let len = loop_filter_ir.len();
    //loop_filter_ir[0] = Complex32::new(loop_p + loop_i, 0.);
    loop_filter_ir[len - 1] = Complex32::new(loop_p + loop_i, 0.);

    let mut loop_filter = FIRFilter::new(&loop_filter_ir);
    //loop_filter.normalize_sum();
    let mut elg = ELGate::new(sample_per_symbol as f32, loop_filter);

    // Timing recovery
    let mut preamble_count = 0;
    let mut bit_index = 0;
    let mut last = 0u8;
    while let Ok(real_sample) = rx.recv() {
        let iq = iq_sampler.sample(real_sample);

        // Perform corellation
        let neg_energy = matched_filter_neg.next(iq);
        let pos_energy = matched_filter_pos.next(iq);

        let matched = matched_filter.next_real(pos_energy.mag() - neg_energy.mag());

        if let Some((bit, eye)) = elg.next_eye(matched) {
            let _ = eye_sender.send(eye);
            ctx.request_repaint();

            last <<= 1;
            last |= bit as u8;

            if preamble_count >= 2 {
                bit_index += 1;
                if bit_index == 8 {
                    if last == 4 {
                        println!();
                        preamble_count = 0;
                    } else {
                        print!("{}", last as char);
                        bit_index = 0;
                    }
                }
            }
            if last == 0xD8 {
                preamble_count += 1;
            }
        }
    }
}

//#[derive(Default)]
struct EguiApp {
    eye_rx: Receiver<Vec<f32>>,
    eyes: VecDeque<Vec<f32>>,
}

impl EguiApp {
    fn new(
        cc: &eframe::CreationContext<'_>,
        eye_rx: Receiver<Vec<f32>>,
        ctx_tx: Sender<Arc<egui::Context>>,
    ) -> Self {
        ctx_tx.send(Arc::new(cc.egui_ctx.clone())).unwrap();

        Self {
            eye_rx,
            eyes: VecDeque::new(),
        }
    }
}

impl eframe::App for EguiApp {
    fn update(&mut self, ctx: &egui::Context, frame: &mut eframe::Frame) {
        egui::CentralPanel::default().show(ctx, |ui| {
            let max_eyes = 100;

            while let Ok(eye) = self.eye_rx.try_recv() {
                self.eyes.push_back(eye);
            }

            while self.eyes.len() > max_eyes {
                self.eyes.pop_front();
            }

            Plot::new("Eye")
                .legend(Legend::default())
                .show(ui, |plot_ui| {
                    //plot_ui.set_auto_bounds(Vec2b { x: false, y: false });
                    for eye in self.eyes.iter() {
                        let line = Line::new(
                            "Eye",
                            eye.iter()
                                .enumerate()
                                .map(|(i, x)| [i as f64, *x as f64])
                                .collect::<Vec<_>>(),
                        )
                        .color(Color32::LIGHT_GREEN);
                        plot_ui.line(line);
                    }
                })
        });
    }
}

fn demodulate() {
    let mut reader = hound::WavReader::open("audio/modulated.wav").unwrap();
    let samples = reader.samples::<i16>();

    // Modulation parameters
    let frequency = 1700.;
    let deviation = 500.;

    // Data parameters
    let sample_rate = 48000;
    let baud_rate = BAUD_RATE;
    let sample_per_symbol = sample_rate / baud_rate;

    let mut iq_sampler = IQSampler::new(hz_to_rad_per_sample(frequency, sample_rate as f32));
    let iq = samples
        .map(|x| iq_sampler.sample(x.unwrap() as f32 / i16::MAX as f32))
        .collect::<Vec<_>>();

    let mut nco_pos = Nco::new(hz_to_rad_per_sample(deviation, sample_rate as f32));
    let mut nco_neg = Nco::new(hz_to_rad_per_sample(-deviation, sample_rate as f32));
    let corellator_pos = (0..(sample_per_symbol * 1))
        .map(|_| {
            nco_pos.step();
            nco_pos.cexp()
        })
        .collect::<Vec<_>>();
    let corellator_neg = (0..(sample_per_symbol * 1))
        .map(|_| {
            nco_neg.step();
            nco_neg.cexp()
        })
        .collect::<Vec<_>>();

    let mut matched_filter_pos = FIRFilter::new(&corellator_pos);
    let mut matched_filter_neg = FIRFilter::new(&corellator_neg);

    let mut dem = vec![];

    for x in &iq {
        let pos = matched_filter_pos.next(x.clone());
        let neg = matched_filter_neg.next(*x);
        dem.push(pos.mag() - neg.mag());
    }

    // Symbol recovery

    let mut bits = vec![];
    let delta = 0.5;
    let alpha = 0.01;
    let mut current_sps = sample_per_symbol as f32;
    let mut current_position = current_sps / 2.;

    while current_position < dem.len() as f32 {
        // Sample before after
        let early_id = (current_position - (delta * current_sps)).max(0.).floor() as u32;
        let late_id = (current_position + (delta * current_sps)).max(0.).floor() as u32;
        if late_id as usize >= dem.len() {
            break;
        }
        let early = dem[early_id as usize];
        let late = dem[late_id as usize];
        let error = early * early - late * late;
        current_sps -= alpha * error;

        bits.push(dem[current_position.floor() as usize] > 0.);

        current_position += current_sps;
    }

    //assert!(bits.len() % 8 == 0);

    let mut out_file = File::create("out.txt").unwrap();
    let mut strip = 0;

    let bit_slice = bits.as_slice();
    for i in 0..100 {
        let byte = bits_to_byte(&bit_slice[(i as usize)..(i as usize + 8)]);
        if byte == 0b01010111u8 {
            strip = i + 8;
        }
    }

    for i in 0..strip {
        bits.remove(i as usize);
    }
    for x in bits.chunks(8) {
        if x.len() != 8 {
            break;
        }
        out_file.write_all(&[bits_to_byte(x)]).unwrap();
    }
}

fn modulate() {
    // Modulation parameters
    let frequency = 1700.;
    let deviation = 500.;

    // Data parameter
    let sample_rate = 48000;
    let baud_rate = BAUD_RATE;

    let host = cpal::default_host();
    let device = host.default_output_device().unwrap();
    let mut supported_configs_range = device.supported_output_configs().unwrap();
    let supported_config = supported_configs_range
        .find(|config| {
            config.sample_format() == cpal::SampleFormat::F32
                && config.min_sample_rate().0 <= 48000
                && config.max_sample_rate().0 >= 48000
        })
        .expect("Device does not support 48kHz f32 output");
    let config = supported_config
        .with_sample_rate(cpal::SampleRate(48_000))
        .config();

    loop {
        let mut buffer = String::new();
        let stdin = io::stdin(); // We get `Stdin` here.
        stdin.read_line(&mut buffer).unwrap();

        for c in buffer.bytes() {
            print!("{}", c as char);
        }
        println!();

        // Construct payload
        let mut bitstream = std::iter::repeat_n(0b01010101u8, 64)
            .chain(std::iter::repeat_n(0xD8, 2))
            .chain(buffer.bytes())
            .chain(std::iter::repeat_n(4u8, 32))
            .flat_map(byte_to_bits);

        let mut modulator = BFSKMod::new(
            sample_rate / baud_rate,
            units::frequency::hz_to_rad_per_sample(deviation, sample_rate as f32),
            &mut bitstream,
        );

        let mut lo = Nco::new(units::frequency::hz_to_rad_per_sample(
            frequency,
            sample_rate as f32,
        ));

        // To send
        let stream = modulator
            .zip(lo)
            .map(|(s, up)| (s * up).re)
            .collect::<Vec<_>>();

        let sample_clock = Arc::new(AtomicU32::new(0));
        let (tx, rx) = channel::<()>();

        let stream = device
            .build_output_stream(
                &config,
                move |data: &mut [f32], _: &cpal::OutputCallbackInfo| {
                    for d in data.iter_mut() {
                        if sample_clock.load(Ordering::Relaxed) as usize == stream.len() {
                            tx.send(()).unwrap();
                            break;
                        }
                        *d = stream[sample_clock.fetch_add(1, Ordering::Relaxed) as usize]
                    }
                },
                move |err| {
                    eprintln!("Stream error: {}", err);
                },
                None,
            )
            .unwrap();

        stream.play().unwrap();
        let _ = rx.recv();
        stream.pause().unwrap();
    }
}

fn byte_to_bits(byte: u8) -> Vec<bool> {
    vec![
        byte & 1 == 1,
        (byte >> 1) & 1 == 1,
        (byte >> 2) & 1 == 1,
        (byte >> 3) & 1 == 1,
        (byte >> 4) & 1 == 1,
        (byte >> 5) & 1 == 1,
        (byte >> 6) & 1 == 1,
        (byte >> 7) & 1 == 1,
    ]
}

/*
fn bits_to_byte(bits: &[bool]) -> u8 {
    bits[7] as u8 |
    bits[6] as u8 >> 1 |
    bits[5] as u8 >> 2 |
    bits[4] as u8 >> 3 |
    bits[3] as u8 >> 4 |
    bits[2] as u8 >> 5 |
    bits[1] as u8 >> 6 |
    bits[0] as u8 >> 7
}
*/

fn bits_to_byte(bits: &[bool]) -> u8 {
    bits[0] as u8
        | (bits[1] as u8) << 1
        | (bits[2] as u8) << 2
        | (bits[3] as u8) << 3
        | (bits[4] as u8) << 4
        | (bits[5] as u8) << 5
        | (bits[6] as u8) << 6
        | (bits[7] as u8) << 7
}
