use rpds::RedBlackTreeSet;
use std::io;
use std::io::prelude::*;
use std::iter;
use std::iter::FromIterator;
use std::process;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Day 24: Planet of Discord")]
/// Simulates the life and death of Eris bugs
///
/// See https://adventofcode.com/2019/day/24 for details.
struct Opt {
    /// How many iterations of the game of life should be run
    /// If not provided, runs until a state appears twice.
    #[structopt(short = "n")]
    depth: Option<usize>,
    /// Interprets the map as being one part of an infinite fractal map
    #[structopt(short = "f")]
    fractal: bool,
}

fn main() {
    let stdin = io::stdin();
    let opt = Opt::from_args();

    let initial_state: State = stdin
        .lock()
        .lines()
        .map(|x| exit_on_failed_assertion(x, "Error reading input"))
        .collect::<State>()
        .with_fractal_mode(opt.fractal);

    let final_state = match opt.depth {
        Some(depth) => initial_state.n_steps_forward(depth),
        None => initial_state.first_repeated_state(),
    };

    println!("Bugs: {}", final_state.alive.size());
    println!("Biodiversity: {}", final_state.biodiversity_rating());
}

fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message: &str) -> A {
    match data {
        Ok(data) => data,
        Err(e) => {
            eprintln!("{}: {}", message, e);
            process::exit(1);
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct State {
    alive: RedBlackTreeSet<Coordinate>,
    fractal_mode: bool,
}

impl FromIterator<String> for State {
    fn from_iter<T: IntoIterator<Item = String>>(iter: T) -> Self {
        State {
            alive: iter
                .into_iter()
                .enumerate()
                .flat_map(move |(y, line)| {
                    line.chars()
                        .enumerate()
                        .filter(move |(_x, c)| *c == '#')
                        .map(move |(x, _c)| Coordinate {
                            x: x as isize,
                            y: y as isize,
                            depth: 0,
                        })
                        .collect::<Vec<Coordinate>>()
                })
                .collect(),
            fractal_mode: false,
        }
    }
}

impl State {
    fn with_fractal_mode(&self, fractal_mode: bool) -> State {
        State {
            fractal_mode,
            ..self.clone()
        }
    }

    fn n_steps_forward(&self, n: usize) -> Self {
        iter::successors(Some(self.clone()), |state| Some(state.next()))
            .nth(n)
            .unwrap()
    }

    fn first_repeated_state(&self) -> State {
        iter::successors(
            Some((RedBlackTreeSet::new(), self.clone())),
            |(seen, state)| Some((seen.insert(state.clone()), state.next())),
        )
        .find(|(seen, state)| seen.contains(state))
        .unwrap()
        .1
    }

    fn next(&self) -> State {
        State {
            alive: self
                .still_alive_next_round()
                .chain(self.comes_alive_next_round())
                .collect(),
            ..self.clone()
        }
    }

    fn biodiversity_rating(&self) -> usize {
        self.alive
            .iter()
            .map(|coord| 2_usize.pow((coord.y * 5 + coord.x) as u32))
            .sum()
    }

    fn still_alive_next_round<'a>(&'a self) -> impl Iterator<Item = Coordinate> + 'a {
        self.alive
            .iter()
            .filter(move |coord| self.count_alive_neighbours(**coord) == 1)
            .cloned()
    }

    fn comes_alive_next_round<'a>(&'a self) -> impl Iterator<Item = Coordinate> + 'a {
        self.alive
            .iter()
            .flat_map(move |coord| self.neighbours(*coord))
            .filter(move |coord| !self.alive.contains(coord))
            .filter(move |coord| {
                self.count_alive_neighbours(*coord) == 1 || self.count_alive_neighbours(*coord) == 2
            })
    }

    fn count_alive_neighbours(&self, coordinate: Coordinate) -> usize {
        self.neighbours(coordinate)
            .into_iter()
            .filter(|coord| self.alive.contains(coord))
            .count()
    }

    fn neighbours(&self, coordinate: Coordinate) -> Vec<Coordinate> {
        if self.fractal_mode {
            [(-1, 0), (1, 0), (0, -1), (0, 1)]
                .into_iter()
                .map(|(dx, dy)| Coordinate {
                    x: coordinate.x + dx,
                    y: coordinate.y + dy,
                    depth: coordinate.depth,
                })
                .flat_map(move |coord| match (coord.x, coord.y) {
                    (x, _y) if x < 0 => vec![Coordinate {
                        x: 1,
                        y: 2,
                        depth: coord.depth - 1,
                    }]
                    .into_iter(),
                    (_x, y) if y < 0 => vec![Coordinate {
                        x: 2,
                        y: 1,
                        depth: coord.depth - 1,
                    }]
                    .into_iter(),
                    (x, _y) if x >= 5 => vec![Coordinate {
                        x: 3,
                        y: 2,
                        depth: coord.depth - 1,
                    }]
                    .into_iter(),
                    (_x, y) if y >= 5 => vec![Coordinate {
                        x: 2,
                        y: 3,
                        depth: coord.depth - 1,
                    }]
                    .into_iter(),
                    (2, 2) => match (coordinate.x, coordinate.y) {
                        (1, 2) => (0..5)
                            .map(|y| Coordinate {
                                x: 0,
                                y,
                                depth: coord.depth + 1,
                            })
                            .collect::<Vec<_>>()
                            .into_iter(),
                        (2, 1) => (0..5)
                            .map(|x| Coordinate {
                                x,
                                y: 0,
                                depth: coord.depth + 1,
                            })
                            .collect::<Vec<_>>()
                            .into_iter(),
                        (3, 2) => (0..5)
                            .map(|y| Coordinate {
                                x: 4,
                                y,
                                depth: coord.depth + 1,
                            })
                            .collect::<Vec<_>>()
                            .into_iter(),
                        (2, 3) => (0..5)
                            .map(|x| Coordinate {
                                x,
                                y: 4,
                                depth: coord.depth + 1,
                            })
                            .collect::<Vec<_>>()
                            .into_iter(),
                        (_, _) => vec![].into_iter(),
                    },
                    _ => vec![coord.clone()].into_iter(),
                })
                .collect()
        } else {
            [(-1, 0), (1, 0), (0, -1), (0, 1)]
                .into_iter()
                .map(|(dx, dy)| Coordinate {
                    x: coordinate.x + dx,
                    y: coordinate.y + dy,
                    depth: coordinate.depth,
                })
                .filter(|coord| coord.x >= 0 && coord.x < 5 && coord.y >= 0 && coord.y < 5)
                .collect()
        }
    }
}

#[derive(Debug, Clone, Copy, PartialOrd, Ord, PartialEq, Eq)]
struct Coordinate {
    x: isize,
    y: isize,
    depth: isize,
}