use aoc_runner_derive::{aoc, aoc_generator};
use std::collections::HashSet;

#[derive(Clone)]
pub struct Data {
    active: HashSet<(i64, i64, i64, i64)>,
}

impl Data {
    pub fn count_neighbors(&self, x: i64, y: i64, z: i64, w: i64) -> usize {
        let mut ret = 0;
        for dx in &[-1, 0, 1] {
            for dy in &[-1, 0, 1] {
                for dz in &[-1, 0, 1] {
                    for dw in &[-1, 0, 1] {
                        if (*dx == 0) && (*dy == 0) && (*dz == 0) && (*dw == 0) {
                            continue;
                        }
                        let is_active = self.active.contains(&(x + dx, y + dy, z + dz, w + dw));
                        if is_active {
                            ret += 1;
                        }
                    }
                }
            }
        }
        ret
    }
    pub fn update(&mut self) {
        let mut checked = HashSet::new();
        let mut next_active = HashSet::new();
        for &(x, y, z, w) in &self.active {
            for &dx in &[-1, 0, 1] {
                for &dy in &[-1, 0, 1] {
                    for &dz in &[-1, 0, 1] {
                        for &dw in &[-1, 0, 1] {
                            let (nx, ny, nz, nw) = (x + dx, y + dy, z + dz, w + dw);
                            if !checked.insert((nx, ny, nz, nw)) {
                                continue;
                            }
                            let nb = self.count_neighbors(nx, ny, nz, nw);
                            let is_active = self.active.contains(&(nx, ny, nz, nw));
                            if nb == 3 || is_active && nb == 2 {
                                next_active.insert((nx, ny, nz, nw));
                            }
                        }
                    }
                }
            }
        }
        self.active = next_active;
    }

    pub fn count(&self) -> usize {
        return self.active.len();
    }
}

#[aoc_generator(day17, part2)]
pub fn input_generator(input: &str) -> Data {
    let mut active = HashSet::new();
    let z = 0;
    let w = 0;
    for (y, line) in input.trim().lines().enumerate() {
        for (x, c) in line.trim().chars().enumerate() {
            if c == '#' {
                active.insert((x as i64, y as i64, z as i64, w as i64));
            }
        }
    }
    return Data { active };
}

#[aoc(day17, part2)]
pub fn solve_part1(input: &Data) -> usize {
    let mut data = input.clone();
    for _ in 0..6 {
        data.update()
    }
    return data.count();
}