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

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

/// Top-level simulation class.
pub struct Model {
    /// per-population grid (one for each population)
    population_grids: Vec<Grid>,

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

    /// Global grid diffusivity.
    diffusivity: usize,

    /// Current model iteration.
    iteration: usize,

    /// Color palette
    palette: Palette,

    time_per_agent_list: Vec<f64>,
    time_per_step_list: Vec<f64>,
}

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.population_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 pop in 0..n_populations {
            let agents = (0..particles_per_grid)
                .map(|i| Agent::new(width, height, pop, &mut rng, i))
                .collect();
            grids.push(Grid::new(width, height, &mut rng, agents));
        }

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

    pub fn step(&mut self) {
        combine(&mut self.population_grids, &self.attraction_table);

        let agents_tick_time = Instant::now();
        self.population_grids.par_iter_mut().for_each(|grid| {
            grid.tick();
            grid.diffuse(self.diffusivity);
        });
        let agents_tick_elapsed = agents_tick_time.elapsed().as_millis() as f64;
        let agents_num: usize = self.population_grids.iter().map(|g| g.agents.len()).sum();
        let ms_per_agent = agents_tick_elapsed / agents_num as f64;

        self.time_per_agent_list.push(ms_per_agent);
        self.time_per_step_list.push(agents_tick_elapsed);
        self.iteration += 1;
    }

    pub fn run(&mut self, steps: usize) {
        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("#>-"));

        for _ in 0..steps {
            self.step();
            pb.inc(1);
        }
        pb.finish();

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

    pub fn population_grids(&self) -> &[Grid] {
        &self.population_grids
    }

    pub fn palette(&self) -> Palette {
        self.palette
    }
}