type Grid a = S.Set a

class Ord a => GridN a where
  neighbours :: S.Set a -> M.Map a Int

instance GridN (Int,Int,Int) where
    neighbours s = 
        let m = M.fromSet (const 1) s
            idxs = Prelude.filter (\(x,y,z) -> x /= 0 || y /= 0 || z /= 0) $ 
                      [(-1,,), (0,,), (1,,)] <*> [-1, 0, 1] <*> [-1, 0, 1]
            keys = map (\(dx,dy,dz) -> M.mapKeys (\(x,y,z) -> (x+dx, y+dy, z+dz)) m) idxs
        in M.unionsWith (+) keys

instance GridN (Int,Int,Int,Int) where
      neighbours s = 
        let m = M.fromSet (const 1) s
            idxs = Prelude.filter (\(x,y,z,w) -> x /= 0 || y /= 0 || z /= 0 || w /= 0) $ 
                      ((,,,) <$> [-1, 0, 1]) <*> [-1, 0, 1] <*> [-1, 0, 1] <*> [-1, 0, 1]
            keys = map (\(dx,dy,dz,dw) -> M.mapKeys (\(x,y,z,w) -> (x+dx, y+dy, z+dz, w+dw)) m) idxs
        in M.unionsWith (+) keys

parse :: String -> Grid (Int,Int,Int)
parse = parse' 0 0 S.empty
    where
        parse' x y s (c:cs) = case c of
            '#'  -> parse' (x+1) y (S.insert (x, y, 0::Int) s) cs
            '.'  -> parse' (x+1) y s cs
            '\n' -> parse' 0 (y+1) s cs
        parse' _ _ s [] = s

step nbours s  = let ns = nbours s
                 in  S.union (M.keysSet $ M.filter (`elem` [2, 3]) (ns `M.restrictKeys` s))
                             (M.keysSet $ M.filter (== 3) (ns `M.withoutKeys` s))

task1 :: Grid (Int,Int,Int) -> Int
task1 = S.size . (!! 6) . iterate (step neighbours)

task2 :: Grid (Int,Int,Int,Int) -> Int
task2 = S.size . (!! 6) . iterate (step neighbours)

main  = do
  is <- readFile "day17_input" <&> parse
  print (task1 is)
  print (task2 $ S.map (\(x, y, z) -> (x, y, z, 0)) is)

from copy import deepcopy

class Grid3D:
    def __init__(self, start_width, cycles):
        self.grid = []
        self.total_width = start_width + 2 * cycles
        self.total_depth = 1 + 2 * cycles
        self.load_start = self.total_width // 2 - start_width // 2
        self.cycles = cycles
        for z in range(self.total_depth):
            self.grid.append([])
            for y in range(self.total_width):
                self.grid[-1].append([])
                for x in range(self.total_width):
                    self.grid[-1][-1].append(0)

    def load(self, plane):
        for y, row in enumerate(plane, start=self.load_start):
            for x, value in enumerate(row, start=self.load_start):
                if value == "#":
                    self.grid[self.load_start][y][x] = 1

    def tick(self):
        next_grid = deepcopy(self.grid)

        for z, plane in enumerate(self.grid):
            for y, row in enumerate(plane):
                for x, value in enumerate(row):
                    count = self.count_neighbors(x, y, z)
                    if value == 1 and count not in [2, 3]:
                        next_grid[z][y][x] = 0
                    elif value == 0 and count == 3:
                        next_grid[z][y][x] = 1

        self.grid = next_grid

    def tick_all(self):
        for t in range(self.cycles):
            self.tick()

    def count_neighbors(self, x, y, z):
        total = 0
        for zi in range(max(z - 1, 0), min(z + 2, self.total_depth)):
            for yi in range(max(y - 1, 0), min(y + 2, self.total_width)):
                for xi in range(max(x - 1, 0), min(x + 2, self.total_width)):
                    if x == xi and y == yi and z == zi:
                        continue

                    total += self.grid[zi][yi][xi]
        return total

    def count_alive(self):
        return sum(cell for plane in self.grid
                        for row in plane
                        for cell in row)


def part1(plane):
    grid = Grid3D(len(plane), 6)
    grid.load(plane)
    grid.tick_all()
    print(grid.count_alive())


class Grid4D:
    def __init__(self, start_width, cycles):
        self.grid = []
        self.total_width = start_width + 2 * cycles
        self.total_depth = 1 + 2 * cycles
        self.load_start = self.total_width // 2 - start_width // 2
        self.cycles = cycles
        for w in range(self.total_depth):
            self.grid.append([])
            for z in range(self.total_depth):
                self.grid[-1].append([])
                for y in range(self.total_width):
                    self.grid[-1][-1].append([])
                    for x in range(self.total_width):
                        self.grid[-1][-1][-1].append(0)

    def load(self, plane):
        for y, row in enumerate(plane, start=self.load_start):
            for x, value in enumerate(row, start=self.load_start):
                if value == "#":
                    self.grid[self.load_start][self.load_start][y][x] = 1

    def tick(self):
        next_grid = deepcopy(self.grid)

        for w, cube in enumerate(self.grid):
            for z, plane in enumerate(cube):
                for y, row in enumerate(plane):
                    for x, value in enumerate(row):
                        count = self.count_neighbors(x, y, z, w)
                        if value == 1 and count not in [2, 3]:
                            next_grid[w][z][y][x] = 0
                        elif value == 0 and count == 3:
                            next_grid[w][z][y][x] = 1

        self.grid = next_grid

    def tick_all(self):
        for t in range(self.cycles):
            self.tick()

    def count_neighbors(self, x, y, z, w):
        total = 0
        for wi in range(max(w - 1, 0), min(w + 2, self.total_depth)):
            for zi in range(max(z - 1, 0), min(z + 2, self.total_depth)):
                for yi in range(max(y - 1, 0), min(y + 2, self.total_width)):
                    for xi in range(max(x - 1, 0), min(x + 2, self.total_width)):
                        if x == xi and y == yi and z == zi and wi == w:
                            continue

                        total += self.grid[wi][zi][yi][xi]
        return total

    def count_alive(self):
        return sum(cell for cube in self.grid
                        for plane in cube
                        for row in plane
                        for cell in row)


def part2(plane):
    grid = Grid4D(len(plane), 6)
    grid.load(plane)
    grid.tick_all()
    print(grid.count_alive())

if __name__ == "__main__":
    with open("data/day17.txt") as f:
        plane = f.read().splitlines()
    part1(plane)
    part2(plane)

import numpy as np
import tensorflow as tf

data = np.loadtxt("input.txt", 'U', comments=None)
bool_data = data.view('U1') == '#'
bool_data = bool_data.reshape((1, 1, data.size, -1))

winit, zinit, yinit, xinit = bool_data.shape
wdim, xdim, ydim, zdim = 40, 40, 40, 40

padding = tf.constant([[1, 1], [1, 1], [1, 1], [1, 1]])
neighbors_np = np.ones([3, 3, 3, 3])
neighbors_np[1, 1, 1, 1] = 0
conv_filt = tf.reshape(tf.cast(neighbors_np, tf.float32),  [3, 3, 3, 3, 1, 1])

init_data_np = np.zeros([wdim, zdim, xdim, ydim], dtype=np.bool)
ystart = xstart = (xdim-xinit)//2
yend = xend = xstart+xinit
wstart = zstart = zdim//2
wend = zend = zstart+1
init_data_np[wstart:wend, zstart:zend, ystart:yend, xstart:xend] = bool_data
init_data = tf.convert_to_tensor(init_data_np)

@tf.function
def conv4d(data, conv_filt):
    # input, kernel, and output sizes
    (b, wi, zi, yi, xi, c) = data.shape.as_list()
    (wk, zk, yk, xk, ik, ok) = conv_filt.shape.as_list()

    # output size and tensor
    wo = wi - wk + 1
    results = [ None ]*wo

    # convolve each kernel frame i with each input frame j
    for i in range(wk):
        for j in range(wi):

            # add results to this output frame
            out_frame = j - (i - wk//2) - (wi - wo)//2
            if out_frame < 0 or out_frame >= wo:
                continue

            # convolve input frame j with kernel frame i
            frame_conv3d = tf.nn.convolution(tf.reshape(data[:,:,j,:,:], (b, zi, yi, xi, c)), conv_filt[:,:,:,i])

            if results[out_frame] is None:
                results[out_frame] = frame_conv3d
            else:
                results[out_frame] += frame_conv3d

    return tf.stack(results, axis=2)

@tf.function
def generate(this_gen):
    padded = tf.pad(this_gen, padding)
    padded = tf.reshape(padded, [1, wdim+2, zdim+2, ydim+2, xdim+2, 1]) # need for tf.nn.convolution

    convolved = conv4d(tf.cast(padded, tf.float32), conv_filt)
    neighbors = tf.reshape(convolved, [wdim, zdim, xdim, ydim])

    three_neighbors = tf.math.equal(neighbors, 3)
    two_or_three_neighbors = tf.math.logical_or(tf.math.equal(neighbors, 2), three_neighbors)

    next_gen = tf.math.logical_or(this_gen, three_neighbors)
    next_gen = tf.math.logical_and(next_gen, two_or_three_neighbors)
    return next_gen

generation = init_data
for _ in range(6):
    generation = generate(generation)

print("total", tf.math.reduce_sum(tf.cast(generation, tf.int32)))

use std::cmp::{max, min};
use std::collections::HashMap;
use std::fmt;
use std::hash::Hash;
use std::ops::{AddAssign, RangeInclusive};

// --- model

#[derive(Eq, PartialEq, Copy, Clone)]
enum Cube {
    Inactive,
    Active
}

impl From<char> for Cube {
    fn from(c: char) -> Self {
        match c {
            '#' => Cube::Active,
            _ => Cube::Inactive
        }
    }
}

trait Position: Eq + Hash {
    fn neighbours(&self) -> Box<dyn Iterator<Item = Self> + '_>;
}

#[derive(Debug, Eq, PartialEq, Copy, Clone, Hash)]
struct Pos3(i64, i64, i64);

#[derive(Debug, Eq, PartialEq, Copy, Clone, Hash)]
struct Pos4(i64, i64, i64, i64);

impl Position for Pos3 {
    fn neighbours(&self) -> Box<dyn Iterator<Item = Self> + '_> {
        let it = (-1..=1).flat_map(
            move |z| (-1..=1).flat_map(
                move |y| (-1..=1).map(
                    move |x| Pos3(self.0+x, self.1+y, self.2+z)
                )
            )
        ).filter(move |p| p != self);

        Box::new(it)
    }
}

impl Position for Pos4 {
    fn neighbours(&self) -> Box<dyn Iterator<Item = Self> + '_> {
        let it = (-1..=1).flat_map(
            move |w| (-1..=1).flat_map(
                move |z| (-1..=1).flat_map(
                    move |y| (-1..=1).map(
                        move |x| Pos4(self.0+x, self.1+y, self.2+z, self.3+w)
                    )
                )
            )
        ).filter(move |p| p != self);

        Box::new(it)
    }
}

struct Bounds3 {
    x: RangeInclusive<i64>,
    y: RangeInclusive<i64>,
    z: RangeInclusive<i64>
}

struct Bounds4 {
    x: RangeInclusive<i64>,
    y: RangeInclusive<i64>,
    z: RangeInclusive<i64>,
    w: RangeInclusive<i64>
}

impl Default for Bounds3 {
    fn default() -> Self {
        Bounds3 {
            x: 0..=0,
            y: 0..=0,
            z: 0..=0
        }
    }
}

impl Default for Bounds4 {
    fn default() -> Self {
        Bounds4 {
            x: 0..=0,
            y: 0..=0,
            z: 0..=0,
            w: 0..=0
        }
    }
}

impl AddAssign<Pos3> for Bounds3 {
    fn add_assign(&mut self, pos: Pos3) {
        self.x = min(*self.x.start(), pos.0) ..= max(*self.x.end(), pos.0);
        self.y = min(*self.y.start(), pos.1) ..= max(*self.y.end(), pos.1);
        self.z = min(*self.z.start(), pos.2) ..= max(*self.z.end(), pos.2);
    }
}

impl AddAssign<Pos4> for Bounds4 {
    fn add_assign(&mut self, pos: Pos4) {
        self.x = min(*self.x.start(), pos.0) ..= max(*self.x.end(), pos.0);
        self.y = min(*self.y.start(), pos.1) ..= max(*self.y.end(), pos.1);
        self.z = min(*self.z.start(), pos.2) ..= max(*self.z.end(), pos.2);
        self.w = min(*self.w.start(), pos.3) ..= max(*self.w.end(), pos.3);
    }
}

trait Dimension<Pos: Position + Copy> where Self: Sized {
    fn grid(&self) -> &HashMap<Pos, Cube>;

    fn iter(&self) -> Box<dyn Iterator<Item = Pos> + '_>;

    fn at(&self, p: &Pos) -> &Cube;

    fn next_generation(&self) -> Self;

    fn occupied_neighbours(&self, p: &Pos) -> usize {
        p.neighbours()
            .filter(|p| 
                self.at(p) == &Cube::Active
            ).count()
    }

    fn bounds<Bounds: Default + AddAssign<Pos>>(&self) -> Bounds {
        let mut bounds = Bounds::default();
        for pos in self.grid().keys() {
            bounds += *pos;
        }
        bounds
    }

    fn active_cubes(&self) -> usize {
        self.grid().values().filter(|c| *c == &Cube::Active).count()
    }

    fn next_generation_grid(&self) -> HashMap<Pos, Cube> {
        self.iter().map(|pos| {
            let occupied = self.occupied_neighbours(&pos);
            let new_state = match self.at(&pos) {
                Cube::Active => 
                    if occupied == 2 || occupied == 3 {
                        Cube::Active
                    } else {
                        Cube::Inactive
                    }

                Cube::Inactive => 
                    if occupied == 3 {
                        Cube::Active
                    } else {
                        Cube::Inactive
                    }
            };
            (pos, new_state)
        }).collect()
    }
}

#[derive(Clone)]
struct PocketDimension<Pos: Position> {
    grid: HashMap<Pos, Cube>
}


impl PartialEq for PocketDimension<Pos3> {
    fn eq(&self, other: &Self) -> bool {
        let mut bounds: Bounds3 = self.bounds();
        for pos in other.iter() { 
            bounds += pos;
        }
        for z in bounds.z {
            for y in (&bounds.y).clone() {
                for x in (&bounds.x).clone() {
                    let pos = Pos3(x, y, z);
                    if self.at(&pos) != other.at(&pos) {
                        return false;
                    }
                }
            }
        }
        true
    }
}

impl PartialEq for PocketDimension<Pos4> {
    fn eq(&self, other: &Self) -> bool {
        let mut bounds: Bounds4 = self.bounds();
        for pos in other.iter() { 
            bounds += pos;
        }
        for w in bounds.w {
            for z in bounds.z.clone() {
                for y in (&bounds.y).clone() {
                    for x in (&bounds.x).clone() {
                        let pos = Pos4(x, y, z, w);
                        if self.at(&pos) != other.at(&pos) {
                            return false;
                        }
                    }
                }
            }
        }
        true
    }
}

impl PocketDimension<Pos3> {
    fn new3(origin: &Pos3, s: &str) -> Self {
        let mut grid = HashMap::new();

        for (z, zs) in s.split("\n\n").enumerate() {
            for (y, ys) in zs.lines().enumerate() {
                for (x, xs) in ys.trim().chars().enumerate() {
                    grid.insert(Pos3(origin.0 + x as i64, origin.1 + y as i64, origin.2 + z as i64), Cube::from(xs));
                }
            }
        }

        PocketDimension { grid }
    }
}

impl Dimension<Pos3> for PocketDimension<Pos3> {
    fn grid(&self) -> &HashMap<Pos3, Cube> {
        &self.grid
    }

    fn at(&self, p: &Pos3) -> &Cube {
        self.grid.get(p).unwrap_or(&Cube::Inactive)
    }

    fn iter(&self) -> Box<dyn Iterator<Item = Pos3> + '_> {
        let bounds: Bounds3 = self.bounds();
        let (xmin, xmax) = (*bounds.x.start() - 1, *bounds.x.end() + 1);
        let (ymin, ymax) = (*bounds.y.start() - 1, *bounds.y.end() + 1);
        let (zmin, zmax) = (*bounds.z.start() - 1, *bounds.z.end() + 1);

        let it = (zmin..=zmax).flat_map(move |z|
            (ymin..=ymax).flat_map(move |y|
                (xmin..=xmax).map(move |x| Pos3(x, y, z) )
            )
        );

        Box::new(it)
    }

    fn next_generation(&self) -> Self {
        PocketDimension { grid: self.next_generation_grid() }
    }
}

impl PocketDimension<Pos4> {
    fn new4(origin: &Pos4, s: &str) -> Self {
        let mut grid = HashMap::new();

        for (z, zs) in s.split("\n\n").enumerate() {
            for (y, ys) in zs.lines().enumerate() {
                for (x, xs) in ys.trim().chars().enumerate() {
                    grid.insert(Pos4(origin.0 + x as i64, origin.1 + y as i64, origin.2 + z as i64, 0), Cube::from(xs));
                }
            }
        }

        PocketDimension { grid }
    }
}

impl Dimension<Pos4> for PocketDimension<Pos4> {
    fn grid(&self) -> &HashMap<Pos4, Cube> {
        &self.grid
    }

    fn at(&self, p: &Pos4) -> &Cube {
        self.grid.get(p).unwrap_or(&Cube::Inactive)
    }

    fn iter(&self) -> Box<dyn Iterator<Item = Pos4> + '_> {
        let bounds: Bounds4 = self.bounds();
        let (xmin, xmax) = (*bounds.x.start() - 1, *bounds.x.end() + 1);
        let (ymin, ymax) = (*bounds.y.start() - 1, *bounds.y.end() + 1);
        let (zmin, zmax) = (*bounds.z.start() - 1, *bounds.z.end() + 1);
        let (wmin, wmax) = (*bounds.w.start() - 1, *bounds.w.end() + 1);

        let it = (wmin..=wmax).flat_map(move |w|
            (zmin..=zmax).flat_map(move |z|
                (ymin..=ymax).flat_map(move |y|
                    (xmin..=xmax).map(move |x| Pos4(x, y, z, w) )
                )
            )
        );

        Box::new(it)
    }

    fn next_generation(&self) -> Self {
        PocketDimension { grid: self.next_generation_grid() }
    }

}

impl fmt::Debug for Cube {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Cube::Inactive => write!(f, "."),
            Cube::Active => write!(f, "#")
        }    
    }
}

impl fmt::Debug for PocketDimension<Pos3> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let bounds: Bounds3 = self.bounds();
        write!(f, "zs={:?} ys={:?} xs={:?}\n", bounds.z, bounds.y, bounds.x)?;
        for z in bounds.z {
            write!(f, "z={:?}\n", z)?;
            for y in (&bounds.y).clone() {
                for x in (&bounds.x).clone() {
                    write!(f, "{:?}", self.at(&Pos3(x,y,z)))?;
                }
                write!(f, " {}\n", y)?;
            }
        }
        Ok(())
    }
}

// --- problems

fn part1(input: &str) -> usize {  
    let mut p = PocketDimension::new3(&Pos3(0,0,0), input);
    for _ in 0..6 {
        p = p.next_generation();
    }
    p.active_cubes()
}

fn part2(input: &str) -> usize {
    let mut p = PocketDimension::new4(&Pos4(0,0,0,0), input);
    for _ in 0..6 {
        p = p.next_generation();
    }
    p.active_cubes()
}


fn main() {
    let input = std::fs::read_to_string("./input.txt").unwrap();
    println!("part1 {:?}", part1(&input));
    println!("part2 {:?}", part2(&input));
}


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

    fn test_grid() -> &'static str {
        ".#.
         ..#
         ###"
    }

    #[test]
    fn test_init() {
        let pd = PocketDimension::new3(&Pos3(0,0,0), test_grid());
        assert_eq!(pd.at(&Pos3(0,0,0)), &Cube::Inactive);
        assert_eq!(pd.at(&Pos3(1,0,0)), &Cube::Active);
        assert_eq!(pd.at(&Pos3(3,6,9)), &Cube::Inactive);
        assert_eq!(pd.at(&Pos3(2,1,0)), &Cube::Active);
    }

    #[test]
    fn test_neighbours_3d() {
        assert_eq!(Pos3(0,0,0).neighbours().count(), 26);
    }

    #[test]
    fn test_neighbours_4d() {
        assert_eq!(Pos4(0,0,0,0).neighbours().count(), 80);
    }

    #[test]
    fn test_occupied_neighbours() {
        let pd = PocketDimension::new3(&Pos3(0,0,0), test_grid());
        assert_eq!(pd.occupied_neighbours(&Pos3(0,0,0)), 1);        
        assert_eq!(pd.occupied_neighbours(&Pos3(1,2,0)), 3);        
    }

    #[test]
    fn test_generations() {
        let pd = PocketDimension::new3(&Pos3(0,0,0), test_grid());

        let gen1 = pd.next_generation();
        assert_eq!(gen1, PocketDimension::new3(&Pos3(0,1,-1),
            "#..
             ..#
             .#.

             #.#
             .##
             .#.

             #..
             ..#
             .#."
        ));

        let gen2 = gen1.next_generation();
        assert_eq!(gen2, PocketDimension::new3(&Pos3(-1,0,-2),
            ".....
             .....
             ..#..
             .....
             .....

             ..#..
             .#..#
             ....#
             .#...
             .....

             ##...
             ##...
             #....
             ....#
             .###.

             ..#..
             .#..#
             ....#
             .#...
             .....

             .....
             .....
             ..#..
             .....
             ....."
        ));
    }

    #[test]
    fn test_six_generations_v1() {
        let mut p = PocketDimension::new3(&Pos3(0,0,0), test_grid());
        for _ in 0..6 {
             p = p.next_generation();
        }
        assert_eq!(p.active_cubes(), 112);
    }
}

#!/usr/bin/perl
use warnings;
use strict;
use feature qw{ say };

use ARGV::OrDATA;

{   package Grid;
    use Moo;
    has grid        => (is => 'ro');
    has width       => (is => 'lazy');
    has height      => (is => 'lazy');
    has depth       => (is => 'lazy');
    has time        => (is => 'lazy');
    has _neighbours => (is => 'lazy');

    sub at {
        my ($self, $x, $y, $z, $w) = @_;
        return 0 if $x < 0 || $y < 0 || $z < 0 || $w < 0;

        return ($self->grid->[$w]
                && $self->grid->[$w][$z]
                && $self->grid->[$w][$z][$y]
                && $self->grid->[$w][$z][$y][$x]) ? 1 : 0
    }

    sub iter {
        my ($self, $code, $m) = @_;
        $m //= 0;
        for my $w (0 .. $self->time - 1 + $m) {
            for my $z (0 .. $self->depth - 1 + $m) {
                for my $y (0 .. $self->height - 1 + $m) {
                    for my $x (0 .. $self->width - 1 + $m) {
                        $code->($x, $y, $z, $w);
                    }
                }
            }
        }
    }

    sub count {
        my ($self) = @_;
        my $c = 0;
        $self->iter(sub {
            my ($x, $y, $z, $w) = @_;
            ++$c if $self->at($x, $y, $z, $w);
        });
        return $c
    }

    sub neighbours {
        my ($self, $x, $y, $z, $w) = @_;
        $self->_neighbours->[ $w + 1 ][ $z + 1 ][ $y + 1 ][ $x + 1 ]
    }

    sub next {
        my ($self) = @_;
        my @next;
        $self->iter(sub {
            my ($x, $y, $z, $w) = @_;
            my $neighbours = $self->neighbours($x - 1, $y - 1, $z - 1, $w - 1);
            my $is_active = $self->at($x - 1, $y - 1, $z - 1, $w - 1);
            $next[$w][$z][$y][$x]
                = $is_active
                ? ($neighbours == 2 || $neighbours == 3)
                : ($neighbours == 3);
        }, 2);
        return \@next
    }

    sub _build_time   { scalar @{ $_[0]->grid } }
    sub _build_depth  { scalar @{ $_[0]->grid->[0] } }
    sub _build_height { scalar @{ $_[0]->grid->[0][0] } }
    sub _build_width  { scalar @{ $_[0]->grid->[0][0][0] } }

    sub _build__neighbours {
        my ($self) = @_;
        my @n;
        $self->iter(sub {
            my ($x, $y, $z, $w) = @_;
            for my $i (-1 .. 1) {
                for my $j (-1 .. 1) {
                    for my $k (-1 .. 1) {
                        for my $l (-1 .. 1) {
                            next unless $i || $j || $k || $l;

                            $n[ $w + $l + 1 ]
                                [ $z + $i + 1 ]
                                [ $y + $j + 1 ]
                                [ $x + $k + 1 ]
                                += $self->at($x, $y, $z, $w);
                        }
                    }
                }
            }
        });
        return \@n
    }
}

my @in;
while (<>) {
    chomp;
    push @in, [map $_ eq '#', split //];
}
my $grid = 'Grid'->new(grid => [[\@in]]);

for (1 .. 6) {
    $grid = 'Grid'->new(grid => $grid->next);
}

say $grid->count;

__DATA__
.#.
..#
###

import Data.List
import Data.Array
import Data.Array.MArray
import Data.Array.IO
import Data.Foldable
import Control.Monad

type Point = (Int, Int, Int)

sumTuple :: Point -> Point -> Point
sumTuple (z1,y1,x1) (z2,y2,x2) = (z1+z2, y1+y2, x1+x2)

dirs :: [Point]
dirs = [(z,y,x) | z <- [-1..1], y <- [-1..1], x <- [-1..1], (z,y,x) /= (0,0,0)]

neighbours :: Point -> [Point]
neighbours p = map (sumTuple p) dirs

count :: Char -> String -> Int
count x = length . filter (==x)

newState :: Array Point Char -> Point -> Maybe Char
newState arr pos = 
  let alive = count '#' $ map (arr !) $ neighbours pos
      curr   = arr ! pos
  in 
    if curr == '#' && (alive > 3 || alive < 2) then Just '.'
    else if curr == '.' && alive == 3 then Just '#'
    else Nothing

update :: Array Point Char -> IOUArray Point Char -> Point -> IO ()
update immutarr mutarr p = 
  forM_ (newState immutarr p) (writeArray mutarr p)

pointsBetween :: Point -> Point -> [Point]
pointsBetween (lz,ly,lx) (uz,uy,ux) = 
  [(z,y,x) | z <- [lz..uz], y <- [ly..uy], x <- [lx..ux]]

generation :: IOUArray Point Char -> IO ()
generation mutarr = do 
  immutarr <- freeze mutarr 

  let (l,u) = bounds immutarr

  mapM_ (update immutarr mutarr) $ pointsBetween (sumTuple (1,1,1) l) (sumTuple (-1,-1,-1) u)

countAlive :: Array Point Char -> Int
countAlive immutarr = 
  let (l,u) = bounds immutarr

  in foldl (\acc p -> if immutarr ! p == '#' then acc+1 else acc) 0 $ pointsBetween (sumTuple (1,1,1) l) (sumTuple (-1,-1,-1) u)

main :: IO ()
main = do 
  input <- lines <$> readFile "input.txt" -- Part 1: 448
  let c = 7
      width  = length (head input) + 1
      height = length input + 1

  -- Using (z,y,x) format since that makes it print correctly
  mutarr <- newArray ((-c,-c,-c), (c,height+c,width+c)) '.' :: IO (IOUArray Point Char)

  let init = [(e, (x,y)) | (y,es) <- zip [0..] input, (x,e) <- zip [0..] (input !! y)]
  mapM_ (\(e, (x,y)) -> writeArray mutarr (0,y, x) e) init


  replicateM_ 6 (generation mutarr)

  immutarr <- freeze mutarr :: IO (Array Point Char)
  -- print immutarr
  print $ countAlive immutarr

$actives = {0 => {0 => {}}}
x_bounds = {max: 0, min: 0}
y_bounds = {max: 0, min: 0}
z_bounds = {max: 0, min: 0}
w_bounds = {max: 0, min: 0}

File.readlines("input.txt").each_with_index do |line, y|
  line.each_char.with_index do |char, x|
    if char == "#"
      $actives[0][0][y] ||= {}
      $actives[0][0][y][x] = true

      x_bounds[:max] = x if x > x_bounds[:max]
      x_bounds[:min] = x if x < x_bounds[:min]
    end
  end
  y_bounds[:max] = y if y > y_bounds[:max]
  y_bounds[:min] = y if y < y_bounds[:min]
end

def get_cube_status(x, y, z, w)
  if ($actives[w][z][y][x] rescue false) == true
    :active
  else
    :inactive
  end
end

def get_active_neighbours(x, y, z, w)
  active_neighbours = []
  (w - 1..w + 1).each do |current_w|
    (z - 1..z + 1).each do |current_z|
      (y - 1..y + 1).each do |current_y|
        (x - 1..x + 1).each do |current_x|
          next if [x, y, z, w] == [current_x, current_y, current_z, current_w]

          if get_cube_status(current_x, current_y, current_z, current_w) == :active
            active_neighbours << [current_x, current_y, current_z, current_w]
          end
        end
      end
    end
  end
  active_neighbours
end


6.times do
  w_bounds[:max] += 1
  w_bounds[:min] -= 1
  z_bounds[:max] += 1
  z_bounds[:min] -= 1
  x_bounds[:max] += 1
  x_bounds[:min] -= 1
  y_bounds[:max] += 1
  y_bounds[:min] -= 1

  next_state = Marshal.load(Marshal.dump($actives))
  (w_bounds[:min]..w_bounds[:max]).each do |w|
    (z_bounds[:min]..z_bounds[:max]).each do |z|
      (y_bounds[:min]..y_bounds[:max]).each do |y|
        (x_bounds[:min]..x_bounds[:max]).each do |x|
          cube_status = get_cube_status(x, y, z, w)
          active_neighbours = get_active_neighbours(x, y, z, w)
          case cube_status
          when :active
            next_state[w][z][y][x] = nil unless (active_neighbours.size == 2 || active_neighbours.size == 3)
          when :inactive
            if active_neighbours.size == 3
              next_state[w] ||= {}
              next_state[w][z] ||= {}
              next_state[w][z][y] ||= {}
              next_state[w][z][y][x] = true
            end
          end
        end
      end
    end
  end
  $actives = next_state
end


def count_active(actives, x_bounds, y_bounds, z_bounds, w_bounds)
  active_count = 0
  (w_bounds[:min]..w_bounds[:max]).each do |w|
    (z_bounds[:min]..z_bounds[:max]).each do |z|
      (y_bounds[:min]..y_bounds[:max]).each do |y|
        (x_bounds[:min]..x_bounds[:max]).each do |x|
          active = (actives[w][z][y][x] rescue false) == true
          active_count += 1 if active
        end
      end
    end
  end
  active_count
end

# Helper function for when you need to visualise the layout
def print_grid(actives, x_bounds, y_bounds, z_bounds, w_bounds)
  (w_bounds[:min]..w_bounds[:max]).each do |w|
    (z_bounds[:min]..z_bounds[:max]).each do |z|
      print "z = #{z}, w = #{w}\n"
      (y_bounds[:min]..y_bounds[:max]).each do |y|
        (x_bounds[:min]..x_bounds[:max]).each do |x|
          active = (actives[z][y][x] rescue false) == true
          print(active ? '#' : '.')
        end
        print "\n"
      end
      print "\n"
    end
    print "\n\n"
  end
end

# print_grid($actives, x_bounds, y_bounds, z_bounds, w_bounds)
p count_active($actives, x_bounds, y_bounds, z_bounds, w_bounds)

Before any cycles:

z=0
.#.
..#
###

Before any cycles:

z=0
  ABC

0 .#.
1 ..#
2 ###


After 1 cycle:

z=-1
  ABC

0 #..
1 ..#
2 .#.

z=0
  ABC

0 #.#
1 .##
2 .#.

z=1
  ABC

0 #..
1 ..#
2 .#.

Before any cycles:

z=0
  ABC

0 .#.
1 ..#
2 ##X