use crate::{
    grid::{combine, Grid},
    imgdata::{ImgData, ThinGridData},
    palette::{random_palette, Palette},
    agent::Agent,
};

use indicatif::{ParallelProgressIterator, ProgressBar, ProgressStyle};
// use rand::Rng;
use rand_distr::{Distribution, Normal};
use rayon::{iter::ParallelIterator, prelude::*};
use std::{path::Path, time::Instant};

// Top-level simulation class.
pub struct Model {
    // The grid they move on.
    grids: Vec<Grid>,

    // Attraction table governs interaction across populations
    attraction_table: Vec<Vec<f32>>,

    // Global grid diffusivity.
    diffusivity: usize,

    // Current model iteration.
    iteration: i32,

    // Color palette
    palette: Palette,

    // List of ImgData to be processed post-simulation into images
    img_data_vec: Vec<ImgData>,
}

impl Model {
    const ATTRACTION_FACTOR_MEAN: f32 = 1.0;
    const ATTRACTION_FACTOR_STD: f32 = 0.1;
    const REPULSION_FACTOR_MEAN: f32 = -1.0;
    const REPULSION_FACTOR_STD: f32 = 0.1;

    pub fn print_configurations(&self) {
        for (i, grid) in self.grids.iter().enumerate() {
            println!("Grid {}: {}", i, grid.config);
        }
        println!("Attraction table: {:#?}", self.attraction_table);
    }

    // Construct a new model with random initial conditions and random configuration.
    pub fn new(
        width: usize,
        height: usize,
        n_particles: usize,
        n_populations: usize,
        diffusivity: usize,
    ) -> Self {
        let particles_per_grid = (n_particles as f64 / n_populations as f64).ceil() as usize;
        let _n_particles = particles_per_grid * n_populations;

        let mut rng = rand::thread_rng();

        let attraction_distr =
            Normal::new(Self::ATTRACTION_FACTOR_MEAN, Self::ATTRACTION_FACTOR_STD).unwrap();
        let repulstion_distr =
            Normal::new(Self::REPULSION_FACTOR_MEAN, Self::REPULSION_FACTOR_STD).unwrap();

        let mut attraction_table = Vec::with_capacity(n_populations);
        for i in 0..n_populations {
            attraction_table.push(Vec::with_capacity(n_populations));
            for j in 0..n_populations {
                attraction_table[i].push(if i == j {
                    attraction_distr.sample(&mut rng)
                } else {
                    repulstion_distr.sample(&mut rng)
                });
            }
        }

        let mut grids: Vec<Grid> = Vec::new();
        for _ in 0..n_populations {
            let agents = (0..particles_per_grid)
            .map(|i| Agent::new(width, height, i / particles_per_grid, &mut rng, i)).collect();
            grids.push(Grid::new(width, height, &mut rng, agents));

        }
        Model {
            grids,
            attraction_table,
            diffusivity,
            iteration: 0,
            palette: random_palette(),
            img_data_vec: Vec::new(),
        }
    }

    // Simulates `steps` # of steps
    #[inline]
    pub fn run(&mut self, steps: usize) {
        let debug: bool = false;

        let pb = ProgressBar::new(steps as u64);
        pb.set_style(
            ProgressStyle::default_bar()
                .template(
                    "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta} {percent}%, {per_sec})",
                )
                .progress_chars("#>-"),
        );

        let mut time_per_agent_list: Vec<f64> = Vec::new();
        let mut time_per_step_list: Vec<f64> = Vec::new();

        let agents_num: usize = self.grids.iter().map(|grid| grid.agents.len()).sum();
        for i in 0..steps {
            if debug {
                println!("Starting tick for all agents...")
            };

            // Combine grids
            let grids = &mut self.grids;
            combine(grids, &self.attraction_table);
            let grids_immutable_1 = &grids.clone();

            let agents_tick_time = Instant::now();

            // Tick agents
            for grid in grids.iter_mut() {
                grid.agents.par_iter_mut().for_each(|agent| {
                    agent.tick(&grids_immutable_1[agent.population_id]);
                });
            }

            // Deposit
            let grids_immutable_2 = &grids.clone();
            for grid in grids_immutable_2.iter() {
                for agent in grid.agents.iter() {
                    self.grids[agent.population_id].deposit(agent.x, agent.y);
                }
            }

            // Diffuse + Decay
            let diffusivity = self.diffusivity;
            self.grids.par_iter_mut().for_each(|grid| {
                grid.diffuse(diffusivity);
            });

            self.save_image_data();

            let agents_tick_elapsed: f64 = agents_tick_time.elapsed().as_millis() as f64;
            let ms_per_agent: f64 = (agents_tick_elapsed as f64) / (agents_num as f64);
            time_per_agent_list.push(ms_per_agent);
            time_per_step_list.push(agents_tick_elapsed);

            if debug {
                println!(
                    "Finished tick for all agents. took {}ms\nTime per agent: {}ms\n",
                    agents_tick_elapsed, ms_per_agent
                )
            };

            self.iteration += 1;
            pb.set_position(i as u64);
        }
        pb.finish();

        let avg_per_step: f64 =
            time_per_step_list.iter().sum::<f64>() as f64 / time_per_step_list.len() as f64;
        let avg_per_agent: f64 =
            time_per_agent_list.iter().sum::<f64>() as f64 / time_per_agent_list.len() as f64;
        println!(
            "Average time per step: {}ms\nAverage time per agent: {}ms",
            avg_per_step, avg_per_agent
        );
    }

    fn strip_grid_data(grids: Vec<Grid>) -> Vec<ThinGridData> {
        return grids.iter().map(|grid| {
            ThinGridData::from_grid(grid)
        }).collect();
    }

    fn save_image_data(&mut self) {
        let grids = Self::strip_grid_data(self.grids.clone());
        let img_data = ImgData::new(grids, self.palette, self.iteration);
        self.img_data_vec.push(img_data);
        if self.grids[0].width > 1024 && self.grids[0].height > 1024 && self.img_data_vec.len() > 100 {
            self.render_all_imgdata();
            self.flush_image_data();
        }
    }

    pub fn flush_image_data(&mut self) {
        self.img_data_vec.clear();
    }

    pub fn render_all_imgdata(&self) {
        if !Path::new("./tmp").exists() {
            std::fs::create_dir("./tmp").expect("could create directory");
        }

        let pb = ProgressBar::new(self.img_data_vec.len() as u64);
        pb.set_style(ProgressStyle::default_bar().template(
            "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] ({pos}/{len}, {percent}%, {per_sec})",
        ));

        /*
        for img in &self.img_data_vec {
            Self::save_to_image(img.to_owned());
            pb.inc(1);
        }
        pb.finish();
        */

        (&self.img_data_vec)
            .par_iter()
            .progress_with(pb)
            .for_each(|img| {
                // Self::save_to_image(img.to_owned());
                img.save_to_image();
            });
    }

    /*
    pub fn save_to_image(imgdata: ImgData) {
        let (width, height) = (imgdata.grids[0].width, imgdata.grids[0].height);
        let mut img = image::RgbImage::new(width as u32, height as u32);

        let max_values: Vec<_> = imgdata
            .grids
            .iter()
            .map(|grid| grid.quantile(0.999) * 1.5)
            .collect();

        for y in 0..height {
            for x in 0..width {
                let i = y * width + x;
                let (mut r, mut g, mut b) = (0.0_f32, 0.0_f32, 0.0_f32);
                for (grid, max_value, color) in
                    multizip((&imgdata.grids, &max_values, &imgdata.palette.colors))
                {
                    let mut t = (grid.data()[i] / max_value).clamp(0.0, 1.0);
                    t = t.powf(1.0 / 2.2); // gamma correction
                    r += color.0[0] as f32 * t;
                    g += color.0[1] as f32 * t;
                    b += color.0[2] as f32 * t;
                }
                r = r.clamp(0.0, 255.0);
                g = g.clamp(0.0, 255.0);
                b = b.clamp(0.0, 255.0);
                img.put_pixel(x as u32, y as u32, image::Rgb([r as u8, g as u8, b as u8]));
            }
        }

        img.save(format!("./tmp/out_{}.png", imgdata.iteration).as_str())
            .unwrap();
    }
    */
}
