Day 24: Planet of Discord

use clap::{App, Arg};
use std::collections::{HashSet, VecDeque};
use std::error::Error;
use std::fmt;

fn main() -> Result<(), Box<dyn Error>> {
    let matches = App::new("day24")
        .arg(
            Arg::with_name("INPUT")
                .help("Sets the input file to use")
                .required(true),
        )
        .get_matches();

    let input_file = matches.value_of("INPUT").unwrap();
    let input = std::fs::read_to_string(input_file)?;

    let mut world = SimpleWorld::new_from_string(&input);
    world.run_until_repeat();
    println!("part 1: {}", world.biodiversity_rating());

    let mut rworld = RecursiveWorld::new(Level::new_from_string(&input));
    for _ in 0..200 {
        rworld.step();
    }
    println!("part 2: {}", rworld.count_live());

    Ok(())
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
enum Tile {
    Live,
    Dead,
}

impl fmt::Display for Tile {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Tile::Live => write!(f, "#"),
            Tile::Dead => write!(f, "."),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct SimpleWorld {
    tiles: Vec<Tile>,
    width: usize,
}

impl SimpleWorld {
    fn new_from_string(input: &str) -> SimpleWorld {
        let mut tiles = Vec::with_capacity(25);
        for c in input.chars() {
            let tile = match c {
                '#' => Tile::Live,
                '.' => Tile::Dead,
                ' ' | '\n' | '\r' => continue,
                _ => panic!("Unexpected character: {}", c),
            };
            tiles.push(tile);
        }
        assert_eq!(tiles.len(), 25);
        SimpleWorld { tiles, width: 5 }
    }

    fn step(&mut self) {
        let copy = self.tiles.clone();
        let copy_rows = copy.chunks(self.width).collect::<Vec<_>>();
        let mut rows = self.tiles.chunks_mut(self.width).collect::<Vec<_>>();

        for (y, &copy_row) in copy_rows.iter().enumerate() {
            for (x, tile) in copy_row.iter().enumerate() {
                let n = SimpleWorld::count_neighbors(x, y, &copy_rows);
                match tile {
                    Tile::Live => {
                        if n != 1 {
                            rows[y][x] = Tile::Dead;
                        }
                    }
                    Tile::Dead => {
                        if n == 1 || n == 2 {
                            rows[y][x] = Tile::Live;
                        }
                    }
                }
            }
        }
    }

    fn count_neighbors(x: usize, y: usize, copy_rows: &[&[Tile]]) -> u8 {
        let mut result = 0;
        if x != 0 && copy_rows[y][x - 1] == Tile::Live {
            result += 1;
        }
        if y != 0 && copy_rows[y - 1][x] == Tile::Live {
            result += 1;
        }
        if x != copy_rows[y].len() - 1 && copy_rows[y][x + 1] == Tile::Live {
            result += 1;
        }
        if y != copy_rows.len() - 1 && copy_rows[y + 1][x] == Tile::Live {
            result += 1;
        }
        result
    }

    fn run_until_repeat(&mut self) {
        let mut seen_states = HashSet::new();
        loop {
            if !seen_states.insert(self.clone()) {
                break;
            }
            self.step();
        }
    }

    fn biodiversity_rating(&self) -> u64 {
        self.tiles
            .iter()
            .enumerate()
            .filter(|(_, &tile)| tile == Tile::Live)
            .map(|(i, _)| 2_u64.pow(i as u32))
            .sum()
    }
}

impl fmt::Display for SimpleWorld {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for line in self.tiles.chunks(self.width) {
            for tile in line {
                tile.fmt(f)?;
            }
            writeln!(f, "")?;
        }
        Ok(())
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
struct Level {
    tiles: [Tile; 24],
}

impl Level {
    fn new() -> Level {
        Level {
            tiles: [Tile::Dead; 24],
        }
    }

    fn new_from_string(input: &str) -> Level {
        let mut level = Level::new();
        for (y, line) in input.trim().lines().enumerate() {
            for (x, c) in line.trim().chars().enumerate() {
                if (x, y) == (2, 2) {
                    continue;
                }
                level[(x as u8, y as u8)] = match c {
                    '#' => Tile::Live,
                    '.' => Tile::Dead,
                    _ => panic!("Unexpected character: '{}'", c),
                };
            }
        }
        level
    }

    fn is_empty(&self) -> bool {
        self.tiles == [Tile::Dead; 24]
    }

    fn coord_to_slot(loc: (u8, u8)) -> usize {
        let r = (loc.1 * 5 + loc.0) as usize;
        assert_ne!(r, 12);
        if r >= 12 {
            r - 1
        } else {
            r
        }
    }

    fn count_top(&self) -> u8 {
        self.tiles[0..5]
            .iter()
            .filter(|&&t| t == Tile::Live)
            .count() as u8
    }

    fn count_bottom(&self) -> u8 {
        self.tiles[19..24]
            .iter()
            .filter(|&&t| t == Tile::Live)
            .count() as u8
    }

    fn count_left(&self) -> u8 {
        [
            self.tiles[0],
            self.tiles[5],
            self.tiles[10],
            self.tiles[14],
            self.tiles[19],
        ]
        .iter()
        .filter(|&&t| t == Tile::Live)
        .count() as u8
    }

    fn count_right(&self) -> u8 {
        [
            self.tiles[4],
            self.tiles[9],
            self.tiles[13],
            self.tiles[18],
            self.tiles[23],
        ]
        .iter()
        .filter(|&&t| t == Tile::Live)
        .count() as u8
    }
}

impl std::ops::Index<(u8, u8)> for Level {
    type Output = Tile;

    fn index(&self, loc: (u8, u8)) -> &Self::Output {
        &self.tiles[Level::coord_to_slot(loc)]
    }
}

impl std::ops::IndexMut<(u8, u8)> for Level {
    fn index_mut(&mut self, loc: (u8, u8)) -> &mut Self::Output {
        &mut self.tiles[Level::coord_to_slot(loc)]
    }
}

impl fmt::Display for Level {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for y in 0..5 {
            for x in 0..5 {
                if (x, y) == (2, 2) {
                    write!(f, "?")?;
                } else {
                    self[(x, y)].fmt(f)?;
                }
            }
            writeln!(f, "")?;
        }
        Ok(())
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct RecursiveWorld {
    outermost_depth: i32,
    levels: VecDeque<Level>,
}

impl RecursiveWorld {
    fn new(initial_level: Level) -> RecursiveWorld {
        let mut levels = VecDeque::new();
        levels.push_back(initial_level);
        RecursiveWorld {
            outermost_depth: 0,
            levels,
        }
    }

    fn step(&mut self) {
        let snapshot = self.clone();

        fn is_one_or_two(value: u8) -> bool {
            value == 1 || value == 2
        }

        for (i, level) in snapshot.levels.iter().enumerate() {
            let outer_level = if i == 0 {
                None
            } else {
                Some(&snapshot.levels[i-1])
            };
            let inner_level = snapshot.levels.get(i + 1);

            let new_level = &mut self.levels[i];

            for y in 0..5 {
                for x in 0..5 {
                    if (x, y) == (2, 2) {
                        continue;
                    }
                    let tile = level[(x, y)];
                    let n = RecursiveWorld::count_neighbors(level, (x, y), inner_level, outer_level);
                    match tile {
                        Tile::Live => {
                            if n != 1 {
                                new_level[(x, y)] = Tile::Dead;
                            }
                        }
                        Tile::Dead => {
                            if n == 1 || n == 2 {
                                new_level[(x, y)] = Tile::Live;
                            }
                        }
                    }
                }
            }
        }

        // check if we need to add new levels on the ends
        {
            let outermost_level = &snapshot.levels[0];
            let mut new_level = Level::new();
            if is_one_or_two(outermost_level.count_top()) {
                new_level[(2, 1)] = Tile::Live;
            }
            if is_one_or_two(outermost_level.count_bottom()) {
                new_level[(2, 3)] = Tile::Live;
            }
            if is_one_or_two(outermost_level.count_left()) {
                new_level[(1, 2)] = Tile::Live;
            }
            if is_one_or_two(outermost_level.count_right()) {
                new_level[(3, 2)] = Tile::Live;
            }
            if !new_level.is_empty() {
                self.outermost_depth -= 1;
                self.levels.push_front(new_level);
            }
        }

        {
            let innermost_level = &snapshot.levels[snapshot.levels.len() - 1];
            let mut new_level = Level::new();
            if innermost_level[(2, 1)] == Tile::Live {
                new_level[(0, 0)] = Tile::Live;
                new_level[(1, 0)] = Tile::Live;
                new_level[(2, 0)] = Tile::Live;
                new_level[(3, 0)] = Tile::Live;
                new_level[(4, 0)] = Tile::Live;
            }
            if innermost_level[(1, 2)] == Tile::Live {
                new_level[(0, 0)] = Tile::Live;
                new_level[(0, 1)] = Tile::Live;
                new_level[(0, 2)] = Tile::Live;
                new_level[(0, 3)] = Tile::Live;
                new_level[(0, 4)] = Tile::Live;
            }
            if innermost_level[(3, 1)] == Tile::Live {
                new_level[(4, 0)] = Tile::Live;
                new_level[(4, 1)] = Tile::Live;
                new_level[(4, 2)] = Tile::Live;
                new_level[(4, 3)] = Tile::Live;
                new_level[(4, 4)] = Tile::Live;
            }
            if innermost_level[(2, 3)] == Tile::Live {
                new_level[(0, 4)] = Tile::Live;
                new_level[(1, 4)] = Tile::Live;
                new_level[(2, 4)] = Tile::Live;
                new_level[(3, 4)] = Tile::Live;
                new_level[(4, 4)] = Tile::Live;
            }
            if !new_level.is_empty() {
                self.levels.push_back(new_level);
            }
        }
    }

    fn count_neighbors(
        level: &Level,
        loc: (u8, u8),
        inner_level: Option<&Level>,
        outer_level: Option<&Level>,
    ) -> u8 {
        let (x, y) = loc;
        let mut result = 0;
        // check left
        if x != 0 {
            if (x, y) == (3, 2) {
                if let Some(inner_level) = inner_level {
                    result += inner_level.count_right();
                }
            } else if level[(x - 1, y)] == Tile::Live {
                result += 1;
            }
        } else if let Some(outer_level) = outer_level {
            if outer_level[(1, 2)] == Tile::Live {
                result += 1;
            }
        }
        // check up
        if y != 0 {
            if (x, y) == (2, 3) {
                if let Some(inner_level) = inner_level {
                    result += inner_level.count_bottom();
                }
            } else if level[(x, y - 1)] == Tile::Live {
                result += 1;
            }
        } else if let Some(outer_level) = outer_level {
            if outer_level[(2, 1)] == Tile::Live {
                result += 1;
            }
        }
        // check right
        if x != 4 {
            if (x, y) == (1, 2) {
                if let Some(inner_level) = inner_level {
                    result += inner_level.count_left();
                }
            } else if level[(x + 1, y)] == Tile::Live {
                result += 1;
            }
        } else if let Some(outer_level) = outer_level {
            if outer_level[(3, 2)] == Tile::Live {
                result += 1;
            }
        }
        // check down
        if y != 4 {
            if (x, y) == (2, 1) {
                if let Some(inner_level) = inner_level {
                    result += inner_level.count_top();
                }
            } else if level[(x, y + 1)] == Tile::Live {
                result += 1;
            }
        } else if let Some(outer_level) = outer_level {
            if outer_level[(2, 3)] == Tile::Live {
                result += 1;
            }
        }
        result
    }

    fn count_live(&self) -> u64 {
        self.levels
            .iter()
            .map(|l| l.tiles.iter().filter(|&&t| t == Tile::Live).count() as u64)
            .sum()
    }
}

impl fmt::Display for RecursiveWorld {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (i, level) in self.levels.iter().enumerate() {
            let depth = i as i32 + self.outermost_depth;
            writeln!(f, "Depth {}:\n{}", depth, level)?;
        }
        Ok(())
    }
}

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

    #[test]
    fn test_simple_steps() -> Result<(), Box<dyn Error>> {
        let mut world = SimpleWorld::new_from_string(
            "
            ....#
            #..#.
            #..##
            ..#..
            #....
        ",
        );
        world.step();
        assert_eq!(
            world,
            SimpleWorld::new_from_string(
                "
                #..#.
                ####.
                ###.#
                ##.##
                .##..
            "
            )
        );
        world.step();
        assert_eq!(
            world,
            SimpleWorld::new_from_string(
                "
                #####
                ....#
                ....#
                ...#.
                #.###
            "
            )
        );
        world.step();
        assert_eq!(
            world,
            SimpleWorld::new_from_string(
                "
                #....
                ####.
                ...##
                #.##.
                .##.#
            "
            )
        );
        world.step();
        assert_eq!(
            world,
            SimpleWorld::new_from_string(
                "
                ####.
                ....#
                ##..#
                .....
                ##...
            "
            )
        );
        Ok(())
    }

    #[test]
    fn test_simple_loop() -> Result<(), Box<dyn Error>> {
        let mut world = SimpleWorld::new_from_string(
            "
            ....#
            #..#.
            #..##
            ..#..
            #....
        ",
        );
        world.run_until_repeat();
        assert_eq!(
            world,
            SimpleWorld::new_from_string(
                "
                .....
                .....
                .....
                #....
                .#...
            "
            )
        );
        assert_eq!(world.biodiversity_rating(), 2129920);
        Ok(())
    }

    #[test]
    fn test_recursive_world() -> Result<(), Box<dyn Error>> {
        let mut rworld = RecursiveWorld::new(Level::new_from_string(
            "
            ....#
            #..#.
            #.?##
            ..#..
            #....
        ",
        ));
        for _ in 0..10 {
            rworld.step();
        }
        assert_eq!(rworld.count_live(), 99);
        Ok(())
    }
}