mod circle;
mod dummy;
mod mandelbrot;
mod rectangle;
mod wgpu;

use crate::color::set_alpha;
use crate::component::arrow::Arrow;
use crate::component::code::{pseudocode, Code};
use crate::component::grid::Grid;
use crate::egui;
use crate::egui::{Color32, Context, Frame, Pos2, Shape, Stroke, Ui};
use crate::fade_in::{fade_in, fade_in_manual, fade_out, fade_out_manual};
use crate::governor::Governor;
use crate::slide::s4_automata::{conways_game_of_life, randomize_grid};
use crate::slide::s6_fractals::circle::circle;
use crate::slide::s6_fractals::dummy::is_black;
use crate::slide::s6_fractals::mandelbrot::mandelbrot;
use crate::slide::s6_fractals::rectangle::rectangle;
use crate::slide::s6_fractals::wgpu::Accelerator;
use crate::slide::Slide;
use crate::window::WindowPosition;
use eframe::egui::style::Margin;
use eframe::egui::{Align2, ColorImage, Vec2};
use eframe::epaint::CircleShape;
use std::iter;
use std::ops::{Bound, RangeBounds};

/// Mandelbrot, etc.
pub struct Fractals {
    /// Where we are in the transitions.
    state: FractalsState,
    /// For showing more pixels.
    governor: Governor,
    /// Pseudocode display.
    code: Code,
    /// Persistent grid state.
    grid: Grid,
}

#[derive(Default)]
#[allow(unused)]
enum FractalsState {
    /// Nothing on screen (except header).
    #[default]
    Before,
    /// Fading in grid.
    Grid {
        /// When we started fading in the grid.
        fade_start: f64,
    },
    /// Erasing the automata cells.
    Erase {
        /// When we started fading out the dots.
        fade_start: f64,
    },
    /// Fade in axes.
    Axes {
        /// When we started fading in the axes.
        fade_start: f64,
    },
    Rectangle {
        /// How many pixels of the grid we've filled in.
        pixels: usize,
    },
    Circle {
        /// How many pixels of the grid we've filled in.
        pixels: usize,
    },
    Mandelbrot {
        /// When we started fading in the quote about fractals.
        quote: Option<f64>,
        /// How many pixels of the grid we've filled in.
        pixels: usize,
        /// Whether fractal traversal mode is enabled.
        traversal: bool,
    },
    MandelbrotZoom {
        /// When we started fading in.
        fade_start: f64,
        /// When we started smoothing the gradient.
        smooth_start: Option<f64>,
        /// When we started zooming in.
        zoom_start: Option<f64>,
        /// Fractal zoom background image.
        image: ColorImage,
        /// GPU accelerator.
        accelerator: Accelerator,
        zoom: f32,
    },
}

impl Default for Fractals {
    fn default() -> Self {
        const HORIZONTAL_OFFSET: f32 = 275.0;

        // Code goes on the left.
        let mut code = Code::default();
        code.name = "fractal_code";
        code.position = WindowPosition::FromCenter(Vec2::new(-HORIZONTAL_OFFSET, 0.0));
        code.code = pseudocode(include_str!("s6_fractals/dummy.rs"));

        // Grid goes on the right.
        let mut grid = Grid::default();
        grid.name = "fractal_grid";
        grid.position = WindowPosition::FromCenter(Vec2::new(HORIZONTAL_OFFSET, 0.0));

        grid.size_cells = 128;

        randomize_grid(&mut grid);

        // Make grid more life-like.
        for _ in 0..10 {
            grid = conways_game_of_life(&grid, 0, Color32::BLACK, false);
        }

        // Grid cell stroke aliases too much at this resolution.
        grid.stroke_width = 0.0;

        Self {
            state: FractalsState::default(),
            governor: Governor::default(),
            code,
            grid,
        }
    }
}

impl Slide for Fractals {
    fn transition(&mut self, ctx: &Context) -> bool {
        // Increment state by one.
        match &mut self.state {
            FractalsState::Before => {
                self.state = FractalsState::Grid {
                    fade_start: ctx.input().time,
                }
            }
            FractalsState::Grid { .. } => {
                // Make sure we finished fading in.
                self.code.alpha = 1.0;
                self.grid.alpha = 1.0;

                self.state = FractalsState::Erase {
                    fade_start: ctx.input().time,
                };
            }
            FractalsState::Erase { .. } => {
                self.state = FractalsState::Axes {
                    fade_start: ctx.input().time,
                }
            }
            FractalsState::Axes { .. } => self.state = FractalsState::Rectangle { pixels: 0 },
            FractalsState::Rectangle { .. } => self.state = FractalsState::Circle { pixels: 0 },
            FractalsState::Circle { .. } => {
                self.state = FractalsState::Mandelbrot {
                    quote: None,
                    pixels: 0,
                    traversal: false,
                }
            }
            FractalsState::Mandelbrot {
                quote, traversal, ..
            } => {
                if quote.is_none() {
                    *quote = Some(ctx.input().time);
                } else if !*traversal {
                    *traversal = true;
                } else {
                    self.state = FractalsState::MandelbrotZoom {
                        fade_start: ctx.input().time,
                        smooth_start: None,
                        zoom_start: None,
                        image: ColorImage::new([1280, 720], Color32::RED),
                        accelerator: Accelerator::new(1280, 720),
                        zoom: 1.0,
                    }
                }
            }
            FractalsState::MandelbrotZoom {
                smooth_start,
                zoom_start,
                ..
            } => {
                if smooth_start.is_none() {
                    *smooth_start = Some(ctx.input().time);
                } else if zoom_start.is_none() {
                    *zoom_start = Some(ctx.input().time);
                } else {
                    return true;
                }
            }
        }
        false
    }

    fn show(&mut self, ui: &mut Ui) {
        let all_fade_out_start = match &self.state {
            FractalsState::MandelbrotZoom { fade_start, .. } => *fade_start,
            _ => f64::INFINITY,
        };

        Frame::none().margin(Margin::same(20.0)).show(ui, |ui| {
            ui.vertical_centered(|ui| {
                fade_out(ui, all_fade_out_start, |ui| {
                    ui.heading("Fractals");
                    let quote_fade_start = match &self.state {
                        FractalsState::Before | FractalsState::Grid { .. } | FractalsState::Erase { .. } | FractalsState::Axes { .. } | FractalsState::Rectangle { .. } | FractalsState::Circle { .. } => None,
                        FractalsState::Mandelbrot { quote, .. } => *quote,
                        _ => Some(0.0),
                    };
                    if let Some(quote_fade_start) = quote_fade_start {
                        fade_in(ui, quote_fade_start, |ui| {
                            ui.add_space(ui.available_height() - 50.0);
                            ui.centered_and_justified(|ui| {
                                ui.label(r#""Beautiful, damn hard, increasingly useful. That's fractals." -Benoit Mandelbrot"#);
                            });
                        });
                    }
                });
            });
        });

        ui.ctx().request_repaint();

        // The function that will be used to color the grid.
        //
        // It is dynamically typed (so that multiple distinct functions can be used).
        //
        // Note that we use indirection via a reference but not a heap-allocated [`Box`]. We can get
        // away with that since the function doesn't leave our stack frame.
        let mut algo: Box<dyn Fn(f32, f32) -> Color32> = Box::new(wrap_bool_algo(is_black));

        // Don't necessarily need to understand what is going on here.
        fn wrap_bool_algo<A: Fn(f32, f32) -> bool>(a: A) -> impl Fn(f32, f32) -> Color32 {
            move |x: f32, y: f32| {
                if a(x, y) {
                    Color32::BLACK
                } else {
                    Color32::WHITE
                }
            }
        }

        // Again, don't necessarily need to understand what is going on here.
        #[allow(unused)]
        fn wrap_f32_algo<A: Fn(f32, f32) -> f32>(a: A) -> impl Fn(f32, f32) -> Color32 {
            move |x: f32, y: f32| {
                let c = (a(x, y) * 255.0) as u8;
                // Slight bluish tint.
                Color32::from_rgb((c as u16 * 9 / 10) as u8, (c as u16 * 19 / 20) as u8, c)
            }
        }

        // Limit to this many pixels. If [`None`], will render all pixels.
        // We use a mutable reference so we can increment in one place.
        let mut limit_pixels: Option<&mut usize> = None;

        // Whether to skip rendering and clearing entirely.
        let mut skip = false;

        match &mut self.state {
            FractalsState::Before => {
                // Don't proceed to render the code/grid.
                return;
            }
            &mut FractalsState::Grid { fade_start } => {
                // Fade in the code and grid.
                fade_in_manual(ui, fade_start, |_, alpha| {
                    self.code.alpha = alpha;
                    self.grid.alpha = alpha;
                });
                skip = true;
            }
            &mut FractalsState::Erase { fade_start } => {
                fade_out_manual(ui, fade_start, |_, alpha| {
                    // Fade out the automata cells.
                    for (_, _, cell) in self.grid.iter_mut() {
                        if cell.a() > 0 {
                            *cell = set_alpha(*cell, alpha);
                        }
                    }
                });
                skip = true;
            }
            &mut FractalsState::Axes { .. } => {
                skip = true;
            }
            FractalsState::Rectangle { pixels } => {
                self.code.code = pseudocode(include_str!("s6_fractals/rectangle.rs"));
                algo = Box::new(wrap_bool_algo(rectangle));
                limit_pixels = Some(pixels)
            }
            FractalsState::Circle { pixels } => {
                self.code.code = pseudocode(include_str!("s6_fractals/circle.rs"));
                algo = Box::new(wrap_bool_algo(circle));
                limit_pixels = Some(pixels);
            }
            FractalsState::Mandelbrot { pixels, .. } => {
                self.code.code = pseudocode(include_str!("s6_fractals/mandelbrot.rs"));
                algo = Box::new(wrap_bool_algo(mandelbrot));
                limit_pixels = Some(pixels);
            }
            FractalsState::MandelbrotZoom {
                fade_start,
                smooth_start,
                zoom_start,
                image,
                accelerator,
                zoom,
            } => {
                fade_out_manual(ui, *fade_start, |_, alpha| {
                    self.grid.alpha = alpha;
                    self.code.alpha = alpha;
                });

                let inverse_fade_time =
                    1.0 - ((ui.ctx().input().time - *fade_start) as f32 * 0.25).clamp(0.0, 1.0);

                let smooth_time = smooth_start
                    .map(|s| (ui.ctx().input().time - s) as f32)
                    .unwrap_or(0.0);
                let smoothing = (smooth_time * 0.75).min(1.0);

                let fade_target = Pos2::new(-3.75, 0.0);
                let zoom_target = Pos2::new(-0.613287065, -0.44526086);

                let zoom_time = zoom_start
                    .map(|s| (ui.ctx().input().time - s) as f32)
                    .unwrap_or(0.0);
                let screen = ui.ctx().input().screen_rect();
                *zoom =
                    (3.0 * 0.95f32.powf(zoom_time * 10.0).max(0.0001)) + inverse_fade_time * 7.0;
                let lerp = (zoom_time * 0.4).sqrt().min(1.0);
                let center = Pos2::new(
                    zoom_target.x * lerp + fade_target.x * inverse_fade_time,
                    zoom_target.y * lerp + fade_target.y * inverse_fade_time,
                );
                let cutoff = 270f32.min(32.0 + (18.0 / *zoom).sqrt());
                accelerator.run(center, *zoom, cutoff, smoothing, image);
                egui::Image::new(&ui.ctx().load_texture("zoom", image.clone()), Vec2::ZERO)
                    .paint_at(ui, screen);

                fade_in(ui, all_fade_out_start, |ui| {
                    ui.style_mut()
                        .visuals
                        .widgets
                        .noninteractive
                        .fg_stroke
                        .color = Color32::WHITE;
                    ui.add_space(ui.available_height() - 50.0);
                    ui.centered_and_justified(|ui| {
                        ui.label(format!(
                            "x={:.5}, y={:.5}, max={:.1}, scale={:.5}",
                            center.x, center.y, cutoff, *zoom
                        ))
                    });
                });
            }
        }

        // Which pixels the algorithm should render.
        let mut render_range: (Bound<usize>, Bound<usize>) = (Bound::Unbounded, Bound::Unbounded);

        // Which pixels should be set to transparent.
        let mut clear_range: (Bound<usize>, Bound<usize>) = (Bound::Unbounded, Bound::Unbounded);

        // Double mutable borrow to avoid moving limit_pixels.
        if let Some(limit_pixels) = &mut limit_pixels {
            if **limit_pixels < self.grid.size_cells.pow(2)
                && self.governor.ready(ui.ctx().input().time, 0.02)
            {
                let new_limit_pixels = **limit_pixels + self.grid.size_cells.pow(2) * 33 / 32 / 32;

                // This is the range we just added, and therefore must render.
                render_range = (
                    Bound::Included(**limit_pixels),
                    Bound::Excluded(new_limit_pixels),
                );

                // This is all the pixels after (that may be obsolete).
                clear_range = (Bound::Included(new_limit_pixels), Bound::Unbounded);

                **limit_pixels = new_limit_pixels;
            } else {
                // Waiting on timer, don't touch the pixels.
                skip = true;
            }
        }

        // Paint the grid.
        if !skip {
            // Store the resolution before mutably borrowing self.grid.
            let size_cells = self.grid.size_cells;

            for (i, (gx, gy, c)) in self.grid.iter_mut().enumerate() {
                if render_range.contains(&i) {
                    // 0 to 1.
                    let nx = (gx as f32 + 0.5) / size_cells as f32;
                    let ny = (gy as f32 + 0.5) / size_cells as f32;

                    // -2 to 2.
                    let x = nx * 4.0 - 2.0;
                    let y = ny * 4.0 - 2.0;

                    *c = algo(x, y);
                } else if clear_range.contains(&i) {
                    *c = Color32::TRANSPARENT;
                }
            }
        }

        self.code.show(ui.ctx());
        self.grid.show_with_overlays(ui.ctx(), |ui, to_screen| {
            match &self.state {
                &FractalsState::Axes { fade_start } => {
                    fade_in_manual(ui, fade_start, |ui, alpha| {
                        let color = set_alpha(Color32::BLACK, alpha);
                        let middle = self.grid.size_cells / 2;
                        Arrow {
                            origin: to_screen * self.grid.center(middle, middle),
                            tip: to_screen * self.grid.center(self.grid.size_cells - 1, middle),
                            stroke: color,
                            label: "x".into(),
                            label_at_tip: true,
                            label_align: Some(Align2::CENTER_CENTER),
                            label_offset: Vec2::new(-20.0, 10.0),
                            stroke_width: 4.0,
                            font_size: 25.0,
                            ..Arrow::default()
                        }
                        .show(ui);
                        Arrow {
                            origin: to_screen * self.grid.center(middle, middle),
                            tip: to_screen * self.grid.center(middle, 0),
                            stroke: color,
                            label: "y".into(),
                            label_at_tip: true,
                            label_align: Some(Align2::CENTER_CENTER),
                            label_offset: Vec2::new(-10.0, 20.0),
                            stroke_width: 4.0,
                            font_size: 25.0,
                            ..Arrow::default()
                        }
                        .show(ui);
                    });
                }
                &FractalsState::Mandelbrot { traversal, .. } if traversal => {
                    //const TIME_SCALE: f64 = 0.4;
                    let mouse_position = ui.ctx().input().pointer.hover_pos();
                    if let Some(mouse_position) = mouse_position {
                        //let mut last_pos = Pos2::new(0.5 * (ui.ctx().input().time * TIME_SCALE).cos() as f32 - 1.0, 0.5 * (ui.ctx().input().time * TIME_SCALE).sin() as f32);
                        let mut last_pos = ((to_screen.inverse() * mouse_position).to_vec2() * 4.0
                            - Vec2::splat(2.0))
                        .to_pos2();
                        ui.painter().add(Shape::Circle(CircleShape {
                            center: to_screen
                                * ((last_pos.to_vec2() + Vec2::splat(2f32)) * 0.25).to_pos2(),
                            radius: 10.0,
                            fill: Color32::RED,
                            stroke: Stroke::default(),
                        }));
                        for (i, pos) in mandelbrot_iter(last_pos).take(16).enumerate() {
                            Arrow {
                                origin: to_screen
                                    * ((last_pos.to_vec2() + Vec2::splat(2f32)) * 0.25).to_pos2(),
                                tip: to_screen
                                    * ((pos.to_vec2() + Vec2::splat(2f32)) * 0.25).to_pos2(),
                                stroke: set_alpha(
                                    Color32::LIGHT_BLUE,
                                    1.0 / (1.0 + i as f32 * 0.2),
                                ),
                                stroke_width: 4.0,
                                ..Arrow::default()
                            }
                            .show(ui);
                            last_pos = pos;
                        }
                    }
                }
                _ => {}
            }
        });
    }
}

/// Iterates points in a traversal of the mandelbrot set.
///
/// Never terminates, so caller is responsible for setting limits.
pub fn mandelbrot_iter(Pos2 { x, y }: Pos2) -> impl Iterator<Item = Pos2> + 'static {
    let mut x1 = 0.0;
    let mut y1 = 0.0;

    // Rust iterator magic ;)
    iter::repeat(()).map(move |_| {
        let x_tmp = x1 * x1 - y1 * y1 + x;
        y1 = 2.0 * x1 * y1 + y;
        x1 = x_tmp;
        Pos2::new(x1, y1)
    })
}