use std::collections::HashMap;

fn main() {
    let input = advent::read_lines(18);
    println!("18a: {}", count_lights(&input, 100, false));
    println!("18b: {}", count_lights(&input, 100, true));
}

fn count_neighbors(grid: &HashMap<(isize, isize), bool>, pos: (isize, isize)) -> usize {
    let mut count = 0;
    for x in pos.0 - 1 ..= pos.0 + 1 {
        for y in pos.1 - 1 ..= pos.1 + 1 {
            if (x, y) == pos {
                continue;
            }
            if let Some(true) = grid.get(&(x, y)) {
                count += 1;
            }
        }
    }
    count
}

fn stuck_grid(grid: &mut HashMap<(isize, isize), bool>, dim: isize) {
    grid.insert((0, 0), true);
    grid.insert((dim-1, dim-1), true);
    grid.insert((0, dim-1), true);
    grid.insert((dim-1, 0), true);
}

fn step(grid: HashMap<(isize, isize), bool>, dim: usize, stuck: bool) -> HashMap<(isize, isize), bool> {
    let mut next_grid = HashMap::new();

    for y in 0 .. dim {
        for x in 0 .. dim {
            let pos = (x as isize, y as isize);
            let val = *grid.get(&pos).unwrap();
            let neighbors = count_neighbors(&grid, pos);

            let new_val = match val {
                true => [2, 3].contains(&neighbors),
                false => neighbors == 3,
            };
            next_grid.insert(pos, new_val);
        }
    }
    if stuck {
        stuck_grid(&mut next_grid, dim as isize);
    }

    next_grid
}

fn count_lights(input: &[String], steps: u16, stuck: bool) -> usize {
    let dim = input.len();

    let mut grid = HashMap::new();

    for (y, line) in input.iter().enumerate() {
        for (x, c) in line.chars().enumerate() {
            let pos = (x as isize, y as isize);
            match c {
                '#' => grid.insert(pos, true),
                '.' => grid.insert(pos, false),
                _ => unimplemented!(),
            };
        }
    }
    if stuck {
        stuck_grid(&mut grid, dim as isize);
    }

    for _ in 0 .. steps {
        grid = step(grid, dim, stuck);
    }
    grid.values().filter(|&val| *val).count()
}

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

    #[test]
    fn test() {
        let input = [
            ".#.#.#".to_string(),
            "...##.".to_string(),
            "#....#".to_string(),
            "..#...".to_string(),
            "#.#..#".to_string(),
            "####..".to_string(),
        ];
        assert_eq!(count_lights(&input, 4, false), 4);
        assert_eq!(count_lights(&input, 5, true), 17);
    }
}