use itertools::Itertools;

#[derive(Clone)]
pub struct OctopusGrid {
    pub grid: [[u8; 10]; 10],
}

impl OctopusGrid {
    pub fn step(&mut self) -> usize {
        let mut flash_grid: [[bool; 10]; 10] = Default::default();
        grid_iter().for_each(|(x, y)| self.grid[y][x] += 1);
        let mut flash_count = 0;
        loop {
            grid_iter().for_each(|(x, y)| {
                if self.grid[y][x] > 9 && !flash_grid[y][x] {
                    flash_grid[y][x] = true;
                    surrounding((x, y)).for_each(|(x, y)| self.grid[y][x] += 1);
                }
            });
            let new_flash_count = flash_grid
                .iter()
                .flat_map(|col| col.iter())
                .filter(|val| **val)
                .count();
            if new_flash_count == flash_count {
                break;
            } else {
                flash_count = new_flash_count;
            }
        }
        grid_iter().for_each(|(x, y)| {
            if flash_grid[y][x] {
                self.grid[y][x] = 0;
            }
        });
        flash_count
    }
}

impl ToString for OctopusGrid {
    fn to_string(&self) -> String {
        self.grid
            .iter()
            .map(|line| line.iter().map(u8::to_string).collect::<String>())
            .join("\n")
    }
}

fn surrounding((x, y): (usize, usize)) -> impl Iterator<Item = (usize, usize)> {
    let (x, y) = (x as isize, y as isize);
    (-1..=1)
        .cartesian_product(-1..=1)
        .filter_map(move |(x_off, y_off)| {
            if x + x_off >= 0
                && x + x_off < 10
                && y + y_off >= 0
                && y + y_off < 10
                && (x_off != 0 || y_off != 0)
            {
                Some(((x + x_off) as usize, (y + y_off) as usize))
            } else {
                None
            }
        })
}

fn grid_iter() -> impl Iterator<Item = (usize, usize)> {
    (0..10).cartesian_product(0..10)
}

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

    #[test]
    fn flash() {
        let mut grid = OctopusGrid {
            grid: Default::default(),
        };
        for _ in 0..9 {
            assert_eq!(grid.step(), 0);
        }
        assert_eq!(grid.step(), 100);
    }
}