#![allow(dead_code)]

use std::{
    f32::consts::PI,
    fs::File,
    i16,
    io::{Read, Write},
    ops::{Add, Div, Mul, Sub},
};

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 fft::DFTAlgorithm;
use nco::Nco;

use eframe::egui::{self, Color32, Context, debug_text::print};
use egui_plot::{self, Bar, BarChart, Legend, Line, LineStyle, Plot, PlotPoints, VLine};
use plotters::style::Color;

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
}

fn main() {
    modulate();
    println!("Demodulating");
    demodulate();

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

#[derive(Default)]
struct EguiApp {
    samples: Vec<f32>,

    iq: Vec<Complex32>,
    dem: Vec<f32>,
    vlines: Vec<f32>,
    elg_sampling: Vec<f32>,
    eye_diagram: Vec<Vec<f32>>,
}
const BAUD_RATE: u32 = 3200;

impl EguiApp {
    fn new(cc: &eframe::CreationContext<'_>) -> Self {
        let mut reader = hound::WavReader::open("audio/modulated.wav").unwrap();
        let samples = reader.samples::<i16>();
        let sample_count = 10_000;
        let input_test = samples
            .take(sample_count)
            .map(|x| x.unwrap() as f32 / i16::MAX as f32)
            .collect::<Vec<_>>();

        // 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 vlines = (0..sample_count)
            .filter(|i| i % sample_per_symbol as usize == 0)
            .map(|x| x as f32 / sample_count as f32)
            .collect();

        let mut iq_sampler = IQSampler::new(hz_to_rad_per_sample(frequency, sample_rate as f32));
        let iq = input_test
            .iter()
            .map(|x| iq_sampler.sample(*x))
            .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![];

        let mut loop_filter = FIRFilter::new(&[Complex32::new(1., 0.); 1]);
        for x in &iq {
            let pos = matched_filter_pos.next(x.clone());
            let neg = matched_filter_neg.next(*x);
            dem.push(
                loop_filter
                    .next(Complex32::new(pos.mag() - neg.mag(), 0.))
                    .re,
            );
        }

        // Symbol recovery

        let mut sample_ids = 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.;

        let mut eye_diagram = vec![];
        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;

            sample_ids.push(current_position);

            let eye = ((current_position - current_sps).max(0.) as usize
                ..(current_position + current_sps).min(sample_count as f32) as usize)
                .map(|i| dem[i])
                .collect();

            eye_diagram.push(eye);

            current_position += current_sps;
        }

        let elg_sampling = sample_ids
            .iter()
            .map(|x| *x / sample_count as f32)
            .collect();

        Self {
            samples: input_test.clone(),
            iq,
            dem,
            vlines,
            elg_sampling,
            eye_diagram,
        }
    }
}

impl eframe::App for EguiApp {
    fn update(&mut self, ctx: &egui::Context, frame: &mut eframe::Frame) {
        egui::CentralPanel::default().show(ctx, |ui| {
            ui.heading("Hello World!");
            Plot::new("Fourrier transform")
                .legend(Legend::default())
                .show(ui, |plot_ui| {
                    for (i, eye) in self.eye_diagram.iter().enumerate().skip(1) {
                        plot_ui.line(
                            Line::new(
                                "eye",
                                eye.iter()
                                    .enumerate()
                                    .map(|(i, x)| [i as f64 / eye.len() as f64, *x as f64])
                                    .collect::<Vec<_>>(),
                            )
                            .width(1.)
                            .color(Color32::RED),
                        );
                    }

                    self.vlines.iter().for_each(|l| {
                        plot_ui.vline(
                            VLine::new("Boundaries", *l as f64)
                                .color(Color32::LIGHT_BLUE)
                                .width(0.5)
                                .style(LineStyle::dashed_dense()),
                        );
                    });

                    self.elg_sampling.iter().for_each(|l| {
                        plot_ui.vline(
                            VLine::new("ELG", *l as f64)
                                .color(Color32::RED)
                                .width(0.5)
                                .style(LineStyle::dotted_dense()),
                        );
                    });

                    plot_ui.line(
                        Line::new(
                            "Passband",
                            self.samples
                                .iter()
                                .enumerate()
                                .map(|(i, x)| [i as f64 / self.samples.len() as f64, *x as f64])
                                .collect::<Vec<_>>(),
                        )
                        .width(2.),
                    );

                    plot_ui.line(
                        Line::new(
                            "Demodulated",
                            self.dem
                                .iter()
                                .enumerate()
                                .map(|(i, x)| [i as f64 / self.iq.len() as f64, *x as f64])
                                .collect::<Vec<_>>(),
                        )
                        .width(2.),
                    );
                })
        });
    }
}

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;

    // File to modulate
    let f = File::open("s.txt").unwrap();
    let mut bitstream = std::iter::repeat_n(0b01010101u8, 1)
        .chain(std::iter::repeat_n(0b01010111u8, 1))
        .chain(f.bytes().map(|x| x.unwrap()))
        .chain(std::iter::repeat_n(0u8, 1))
        .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,
    ));

    let spec = hound::WavSpec {
        channels: 1,
        sample_rate,
        bits_per_sample: 16,
        sample_format: hound::SampleFormat::Int,
    };

    let mut writer = hound::WavWriter::create("audio/modulated.wav", spec).unwrap();
    for (s, up) in modulator.zip(lo) {
        let sample = (s * up).re; // Project to I coords
        let amplitude = i16::MAX as f32;
        writer.write_sample((sample * amplitude) as i16).unwrap();
    }
    writer.finalize().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
}