use nannou::{geom::Range, prelude::*};

use crate::audio::AudioSample;

mod audio;

fn main() {
    nannou::app(Model::new).update(update).run();
}

struct Complex {
    real: f32,
    imag: f32,
}

impl Complex {
    fn new(real: f32, imag: f32) -> Self {
        Self { real, imag }
    }

    fn length(&self) -> f32 {
        self.real.hypot(self.imag)
    }
}

struct Settings {
    width: f32,
    offset: u32,
}

struct Model {
    settings: Settings,
    window: Entity,
    target_samples: Vec<f32>,
    result_samples: Vec<f32>,
    dft_result: Vec<Complex>,
}

const MAX_SAMPLES: usize = 1024;

fn generate_sin(frequency: usize) -> Vec<f32> {
    let factor = (frequency as f32 / MAX_SAMPLES as f32) * (2.0 * std::f32::consts::PI);
    (0..MAX_SAMPLES)
        .map(|i| f32::sin((i as f32) * factor))
        .collect()
}

fn generate_cos(frequency: usize) -> Vec<f32> {
    let factor = (frequency as f32 / MAX_SAMPLES as f32) * (2.0 * std::f32::consts::PI);
    (0..MAX_SAMPLES)
        .map(|i| f32::cos((i as f32) * factor))
        .collect()
}

fn dot_product(a: &[f32], b: &[f32]) -> f32 {
    a.iter().zip(b.iter()).map(|(a, b)| a * b).sum()
}

fn naive_dft(samples: &[f32]) -> Vec<Complex> {
    let mut result = Vec::new();

    for i in 0..(samples.len() / 2) {
        let sin = generate_sin(i);
        let cos = generate_cos(i);

        let dot = Complex::new(
            dot_product(samples, &sin) * (2. / MAX_SAMPLES as f32),
            dot_product(samples, &cos) * (2. / MAX_SAMPLES as f32),
        );

        result.push(dot);
    }

    result
}

fn inverse_dft(samples: &[Complex]) -> Vec<f32> {
    let mut result = Vec::new();

    for i in 0..MAX_SAMPLES {
        let mut sum = 0.0;
        for (freq, amplitude) in samples.iter().enumerate() {
            let sample =
                (i as f32) * (freq as f32 / MAX_SAMPLES as f32) * (2.0 * std::f32::consts::PI);

            sum += f32::sin(sample) * amplitude.real;
            sum += f32::cos(sample) * amplitude.imag;
        }

        result.push(sum);
    }

    result
}

impl Model {
    fn new(app: &App) -> Self {
        let sample = AudioSample::new("./resources/sample-3s.wav");

        let window = app
            .new_window()
            .size(512, 512)
            .title("nannou")
            .view(view)
            .build();

        Model {
            window,
            settings: Settings {
                width: 1.0,
                offset: 0,
            },
            target_samples: sample.mono_data(),
            result_samples: vec![],
            dft_result: vec![],
        }
    }

    pub fn current_samples(&self) -> Vec<f32> {
        let len = self.target_samples.len();
        let width = (self.settings.width * len as f32) as usize;
        let offset = self.settings.offset as usize;

        let window: Vec<_> = self.target_samples[offset..(offset + width).min(len)]
            .iter()
            .collect();

        let mut samples = Vec::with_capacity(MAX_SAMPLES);

        let sample_compression = window.len() as f32 / MAX_SAMPLES as f32;
        let mut next_pick = 0.;
        for (i, sample) in window.into_iter().enumerate() {
            if next_pick > i as f32 {
                continue;
            }

            samples.push(*sample);
            next_pick += sample_compression;
        }

        samples
    }
}

fn update(app: &App, model: &mut Model) {
    let current_samples = model.current_samples();

    let settings = &mut model.settings;
    let result_samples = &mut model.result_samples;
    let dft_result = &mut model.dft_result;

    let ctx = app.egui_for_window(model.window);

    egui::Window::new("Settings").show(&ctx, |ui| {
        ui.label("Offset:");
        ui.add(egui::Slider::new(
            &mut settings.offset,
            0..=model.target_samples.len() as u32,
        ));

        ui.label("Width:");
        ui.add(egui::Slider::new(&mut settings.width, 0.0..=1.0));

        if ui.button("Apply").clicked() {
            *dft_result = naive_dft(&current_samples);
            *result_samples = inverse_dft(&dft_result);
        }
    });
}

fn view(app: &App, model: &Model) {
    let draw = app.draw();
    draw.background().color(WHITE);

    let window = app.window(model.window);
    let window_rect = window.rect();

    let height = window_rect.h() / 3.0;
    let center_rect = Rect {
        x: window_rect.x,
        y: Range::new(height / 2., -height / 2.),
    };

    draw_samples_point(&draw, &model.current_samples(), center_rect.shift_y(height));
    draw_samples_point(&draw, &model.result_samples, center_rect);
    draw_samples_bar(&draw, &model.dft_result, center_rect.shift_y(-height));
}

fn draw_samples_point(draw: &Draw, samples: &[f32], area: Rect) {
    let sample_width = area.w() / samples.len() as f32;

    for (i, &sample) in samples.iter().enumerate() {
        draw.ellipse()
            .x_y(
                i as f32 * sample_width + area.left(),
                sample * area.h() / 2.0 + area.mid_left().y,
            )
            .radius(1.)
            .color(BLACK);
    }
}

fn draw_samples_bar(draw: &Draw, samples: &[Complex], area: Rect) {
    let sample_width = area.w() / samples.len() as f32;

    for (i, sample) in samples.iter().enumerate() {
        let left = i as f32 * sample_width + area.left();

        draw.rect()
            .x_y(left, area.bottom() + area.h() / 2.0)
            .w_h(sample_width, sample.length().abs() * area.h())
            .color(BLACK);
    }
}