use crate::{agent::Agent, blur::Blur, buffer::Buf};
use rand::Rng;
use rand_distr::Uniform;
use rayon::{iter::ParallelIterator, prelude::*};
use std::fmt::{Display, Formatter};

/// A population configuration.
#[derive(Debug, Clone, Copy)]
pub struct PopulationConfig {
    pub sensor_distance: f32,
    pub step_distance: f32,
    pub sensor_angle: f32,
    pub rotation_angle: f32,

    deposition_amount: f32,
}

impl Display for PopulationConfig {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self)
    }
}

impl PopulationConfig {
    /// Construct a random configuration.
    pub fn new<R: Rng + ?Sized>(rng: &mut R) -> Self {
        PopulationConfig {
            sensor_distance: rng.random_range(0.0..=64.0),
            step_distance: rng.random_range(0.2..=2.0),
            sensor_angle: rng.random_range(0.0_f32..=120.0).to_radians(),
            rotation_angle: rng.random_range(0.0_f32..=120.0).to_radians(),
            deposition_amount: rng.random_range(5.0..=5.0),
        }
    }
}

// A 2D grid with a scalar value per each grid block. Each grid is occupied by a single population, hence we store the population config inside the grid.
#[derive(Debug, Clone)]
pub struct Grid {
    pub config: PopulationConfig,
    pub width: usize,
    pub height: usize,

    pub data: Vec<f32>,

    // Scratch space for the blur operation.
    buf: Buf,

    blur: Blur,
    pub agents: Vec<Agent>,
}

impl Grid {
    /// Create a new grid filled with random floats in the [0.0..1.0) range.
    pub fn new<R: Rng + ?Sized>(
        width: usize,
        height: usize,
        rng: &mut R,
        agents: Vec<Agent>,
    ) -> Self {
        let range = Uniform::new(0.0, 1.0).expect("unable to create uniform distr");
        let data = rng.sample_iter(range).take(width * height).collect();

        Grid {
            width,
            height,
            data,
            config: PopulationConfig::new(rng),
            buf: Buf::new(width, height),
            blur: Blur::new(width),
            agents,
        }
    }

    /// Truncate x and y and return a corresponding index into the data slice.
    const fn index(&self, x: f32, y: f32) -> usize {
        crate::util::index(self.width, self.height, x, y)
    }

    /// Diffuse grid data and apply a decay multiplier.
    pub fn diffuse(&mut self, radius: usize) {
        self.blur.run(
            &mut self.data,
            &mut self.buf.buf,
            self.width,
            self.height,
            radius as f32,
            0.1, // decay is always 0.1
        );
    }

    pub fn tick(&mut self) {
        self.agents.par_iter_mut().for_each(|agent| {
            agent.tick(&self.buf, self.config, self.width, self.height);
        });
        self.deposit_all();
    }

    pub fn deposit_all(&mut self) {
        for agent in self.agents.iter() {
            let idx = self.index(agent.x, agent.y);
            self.data[idx] += self.config.deposition_amount;
        }
    }
}

pub fn combine<T>(grids: &mut [Grid], attraction_table: &[T])
where
    T: AsRef<[f32]> + Sync,
{
    let datas: Vec<_> = grids.iter().map(|grid| &grid.data).collect();
    let bufs: Vec<_> = grids.iter().map(|grid| &grid.buf.buf).collect();

    // We mutate grid buffers and read grid data. We use unsafe because we need shared/unique borrows on different fields of the same Grid struct.
    bufs.iter().enumerate().for_each(|(i, buf)| {
        let buf_ptr = *buf as *const Vec<f32> as *mut Vec<f32>;
        // SAFETY! we can take these are raw pointers because we are
        // getting it from a `&mut [Grid]`
        let buf_ptr_mut = unsafe { buf_ptr.as_mut().unwrap_unchecked() };

        buf_ptr_mut.fill(0.0);
        datas.iter().enumerate().for_each(|(j, other)| {
            let multiplier = attraction_table[i].as_ref()[j];
            buf_ptr_mut
                .iter_mut()
                .zip(*other)
                .for_each(|(to, from)| *to += from * multiplier)
        })
    });
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_grid_new() {
        let mut rng = rand::rng();
        let grid = Grid::new(8, 8, &mut rng, vec![]);
        assert_eq!(grid.index(0.5, 0.6), 0);
        assert_eq!(grid.index(1.5, 0.6), 1);
        assert_eq!(grid.index(0.5, 1.6), 8);
        assert_eq!(grid.index(2.5, 0.6), 2);
        assert_eq!(grid.index(2.5, 1.6), 10);
        assert_eq!(grid.index(7.9, 7.9), 63);
    }
}