use aoc_runner_derive::{aoc, aoc_generator};
use grid::Grid;

#[aoc_generator(day4)]
fn parse(input: &str) -> Grid<u8> {
    input.parse().unwrap()
}

#[aoc(day4, part1)]
fn part1(input: &Grid<u8>) -> u64 {
    (0..input.height() * input.width())
        .filter(|i| *input.get(&input.coord(*i as i64).unwrap()).unwrap() == b'@')
        .map(|i| input.neighbours_count(&input.coord(i as i64).unwrap(), |c| *c == b'@'))
        .filter(|n| *n < 4)
        .count() as u64
}
#[aoc(day4, part2)]
fn part2(input: &Grid<u8>) -> u64 {
    let mut grid = input.clone();
    let mut removed = 0;

    loop {
        let mut removed_iteration = 0;
        for i in 0..grid.width() * grid.height() {
            let pos = grid.coord(i as i64).unwrap();
            if grid.get(&pos).is_some_and(|c| *c == b'@')
                && grid.neighbours_count(&pos, |c| *c == b'@') < 4
            {
                // remove the roll
                grid.set(&pos, b'.');
                removed_iteration += 1;
            }
        }
        if removed_iteration == 0 {
            return removed;
        }
        removed += removed_iteration
    }
}

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

    const EXAMPLE: &str = "..@@.@@@@.
@@@.@.@.@@
@@@@@.@.@@
@.@@@@..@.
@@.@@@@.@@
.@@@@@@@.@
.@.@.@.@@@
@.@@@.@@@@
.@@@@@@@@.
@.@.@@@.@.";

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse(EXAMPLE)), 13);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse(EXAMPLE)), 43);
    }
}