lat 0.1.0

use std::{
    array,
    ops::{Index, IndexMut},
};

pub struct Grid<T> {
    width: usize,
    height: usize,
    buffer: Vec<T>,
}

impl<T: Default + Clone> Grid<T> {
    pub fn from_defaults(width: usize, height: usize) -> Self {
        Self {
            width,
            height,
            buffer: vec![Default::default(); width * height],
        }
    }
}

impl<T> Index<(usize, usize)> for Grid<T> {
    type Output = T;

    fn index(&self, (x, y): (usize, usize)) -> &Self::Output {
        assert!(x < self.width);
        assert!(y < self.height);
        &self.buffer[y * self.width + x]
    }
}
impl<T> IndexMut<(usize, usize)> for Grid<T> {
    fn index_mut(&mut self, (x, y): (usize, usize)) -> &mut Self::Output {
        assert!(x < self.width);
        assert!(y < self.height);
        &mut self.buffer[y * self.width + x]
    }
}

impl<T> Grid<T> {
    pub fn from_fn(width: usize, height: usize, fun: &mut impl FnMut(usize, usize) -> T) -> Self {
        let mut buffer = Vec::with_capacity(width * height);

        for y in 0..height {
            for x in 0..width {
                buffer.push(fun(x, y));
            }
        }

        Self {
            width,
            height,
            buffer,
        }
    }
    pub fn width(&self) -> usize {
        self.width
    }
    pub fn height(&self) -> usize {
        self.height
    }
    pub fn map_once<U, F>(self, mapper: F) -> Grid<U>
    where
        F: FnMut(T) -> U,
    {
        let buffer = self.buffer.into_iter().map(mapper).collect();
        Grid {
            width: self.width,
            height: self.height,
            buffer,
        }
    }
    pub fn map<U, F>(&self, mapper: F) -> Grid<U>
    where
        F: FnMut(&T) -> U,
    {
        let buffer = self.buffer.iter().map(mapper).collect();
        Grid {
            width: self.width,
            height: self.height,
            buffer,
        }
    }
}
impl<T: Clone> Grid<T> {
    pub fn map_mut<F>(&mut self, mut mapper: F) -> &mut Grid<T>
    where
        F: FnMut(T) -> T,
    {
        self.buffer.iter_mut().for_each(|v| {
            let old = v.clone();
            let new = mapper(old);
            *v = new;
        });
        self
    }
    pub fn map_with_neighbors<U, F>(&self, mut mapper: F, default: T) -> Grid<U>
    where
        F: FnMut(&[[T; 3]; 3]) -> U,
        U: Default,
    {
        let (width, height) = (self.width, self.height);
        let mut buffer = Vec::with_capacity(width * height);
        let mut neighborhood = array::from_fn(|_| array::from_fn(|_| default.clone()));
        for y in 0..self.height {
            for x in 0..self.width {
                for oy in 0..3 {
                    for ox in 0..3 {
                        let x = x as isize + ox as isize - 1;
                        let y = y as isize + oy as isize - 1;
                        neighborhood[ox][oy] =
                            if x > 0 && x < width as _ && y > 0 && y < height as _ {
                                self[(x as _, y as _)].clone()
                            } else {
                                default.clone()
                            };
                    }
                }
                buffer.push(mapper(&neighborhood));
            }
        }
        Grid {
            width,
            height,
            buffer,
        }
    }
}

// /////////////////////////////////////////////////////////////////////////////
// Operator Implementation below
// /////////////////////////////////////////////////////////////////////////////

use core::ops::*;
macro_rules! impl_binop {
    ($($Trait:ident => $fun:ident),*) => {$(
        impl<'a, R: 'static, T: $Trait<&'a R>> $Trait<&'a Grid<R>> for Grid<T> {
            type Output = Grid<T::Output>;
            fn $fun(self, rhs: &'a Grid<R>) -> Self::Output {
                assert_eq!(self.width, rhs.width);
                assert_eq!(self.height, rhs.height);
                let buffer = self.buffer.into_iter().zip(rhs.buffer.iter()).map(|(l, r)| l.$fun(r)).collect::<Vec<_>>();
                Grid {
                    width: self.width,
                    height: self.height,
                    buffer
                }
            }
        }
    )*};
}
macro_rules! impl_assign {
    ($($Trait:ident => $fun:ident),*) => {$(
        impl<'a, R, T: $Trait<&'a R>> $Trait<&'a Grid<R>> for Grid<T> {
            fn $fun(&mut self, rhs: &'a Grid<R>) {
                assert_eq!(self.width, rhs.width);
                assert_eq!(self.height, rhs.height);
                self.buffer.iter_mut().zip(rhs.buffer.iter()).for_each(|(l, r)| l.$fun(r));
            }
        }
    )*};
}
macro_rules! impl_unop {
    ($($Trait:ident => $fun:ident),*) => {$(
        impl<'a, T: $Trait> $Trait for Grid<T> {
            type Output = Grid<T::Output>;
            fn $fun(self) -> Self::Output {
                let buffer = self.buffer.into_iter().map($Trait::$fun).collect::<Vec<_>>();
                Grid {
                    width: self.width,
                    height: self.height,
                    buffer
                }
            }
        }
    )*};
}

impl_assign!(
    AddAssign => add_assign,
    BitAndAssign => bitand_assign,
    BitOrAssign => bitor_assign,
    BitXorAssign => bitxor_assign,
    DivAssign => div_assign,
    MulAssign => mul_assign,
    RemAssign => rem_assign,
    ShlAssign => shl_assign,
    ShrAssign => shr_assign,
    SubAssign => sub_assign
);
impl_binop!(
    Add => add,
    Sub => sub,
    Mul => mul,
    Div => div,
    Rem => rem,
    Shr => shr,
    Shl => shl,
    BitAnd => bitand,
    BitOr => bitor,
    BitXor => bitxor
);
impl_unop!(
    Neg => neg,
    Not => not
);