use rpds::rbt_set;
use rpds::vector::Vector;
use rpds::RedBlackTreeMap;
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 20: Donut Maze")]
/// Finds the shortest path through a maze with portals.
///
/// See https://adventofcode.com/2019/day/20 for details.
struct Opt {
    /// include the rule that going through portals changes your depth
    #[structopt(short = "d")]
    include_depth: bool,
}

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

    let maze = stdin
        .lock()
        .lines()
        .map(|x| exit_on_failed_assertion(x, "Error reading input"))
        .collect::<MazeBuilder>()
        .build();

    println!("{}", maze.shortest_route(opt.include_depth));
}

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);
        }
    }
}

struct MazeBuilder {
    map: Vector<Vector<char>>,
}

impl FromIterator<String> for MazeBuilder {
    fn from_iter<T: IntoIterator<Item = String>>(iter: T) -> Self {
        MazeBuilder {
            map: iter
                .into_iter()
                .map(|line| line.chars().collect())
                .collect(),
        }
    }
}

impl MazeBuilder {
    fn build(self) -> Maze {
        Maze {
            walls: self
                .map
                .iter()
                .map(|line| line.iter().map(|ch| *ch != '.').collect())
                .collect(),
            portals: self.grouped_portals(),
            entrance: self
                .all_portals()
                .find(|(id, _)| *id == ['A', 'A'])
                .unwrap()
                .1,
            exit: self
                .all_portals()
                .find(|(id, _)| *id == ['Z', 'Z'])
                .unwrap()
                .1,
        }
    }

    fn grouped_portals(&self) -> Vector<(Point, Point)> {
        self.all_portals()
            .fold(
                (Vector::new(), RedBlackTreeMap::new()),
                |(matched, unmatched): (
                    Vector<(Point, Point)>,
                    RedBlackTreeMap<[char; 2], Point>,
                ),
                 (next_id, next_p)| match unmatched.get(&next_id) {
                    Some(pair) => (
                        matched.push_back(pair.clone().inside_out(
                            next_p,
                            self.map[0].len(),
                            self.map.len(),
                        )),
                        unmatched.remove(&next_id),
                    ),
                    None => (matched, unmatched.insert(next_id, next_p)),
                },
            )
            .0
    }

    fn all_portals(&self) -> impl Iterator<Item = ([char; 2], Point)> + '_ {
        self.horizontal_trailing_portals()
            .chain(self.horizontal_leading_portals())
            .chain(self.vertical_trailing_portals())
            .chain(self.vertical_leading_portals())
    }

    fn horizontal_trailing_portals(&self) -> impl Iterator<Item = ([char; 2], Point)> + '_ {
        // .XX
        (0..self.map.len())
            .flat_map(move |y| (0..self.map[y].len() - 2).map(move |x| Point { x, y }))
            .filter(move |p| {
                self.map[p.y][p.x] == '.'
                    && self.map[p.y][p.x + 1].is_alphabetic()
                    && self.map[p.y][p.x + 2].is_alphabetic()
            })
            .map(move |p| ([self.map[p.y][p.x + 1], self.map[p.y][p.x + 2]], p))
    }

    fn horizontal_leading_portals(&self) -> impl Iterator<Item = ([char; 2], Point)> + '_ {
        // XX.
        (0..self.map.len())
            .flat_map(move |y| (0..self.map[y].len() - 2).map(move |x| Point { x, y }))
            .filter(move |p| {
                self.map[p.y][p.x + 2] == '.'
                    && self.map[p.y][p.x + 1].is_alphabetic()
                    && self.map[p.y][p.x].is_alphabetic()
            })
            .map(move |p| {
                (
                    [self.map[p.y][p.x], self.map[p.y][p.x + 1]],
                    Point { x: p.x + 2, y: p.y },
                )
            })
    }

    fn vertical_trailing_portals(&self) -> impl Iterator<Item = ([char; 2], Point)> + '_ {
        // .XX
        (0..self.map[0].len())
            .flat_map(move |x| (0..self.map.len() - 2).map(move |y| Point { x, y }))
            .filter(move |p| {
                self.map[p.y][p.x] == '.'
                    && self.map[p.y + 1][p.x].is_alphabetic()
                    && self.map[p.y + 2][p.x].is_alphabetic()
            })
            .map(move |p| ([self.map[p.y + 1][p.x], self.map[p.y + 2][p.x]], p))
    }

    fn vertical_leading_portals(&self) -> impl Iterator<Item = ([char; 2], Point)> + '_ {
        // XX.
        (0..self.map[0].len())
            .flat_map(move |x| (0..self.map.len() - 2).map(move |y| Point { x, y }))
            .filter(move |p| {
                self.map[p.y + 2][p.x] == '.'
                    && self.map[p.y + 1][p.x].is_alphabetic()
                    && self.map[p.y][p.x].is_alphabetic()
            })
            .map(move |p| {
                (
                    [self.map[p.y][p.x], self.map[p.y + 1][p.x]],
                    Point { x: p.x, y: p.y + 2 },
                )
            })
    }
}

#[derive(Debug)]
struct Maze {
    walls: Vector<Vector<bool>>,
    portals: Vector<(Point, Point)>,
    entrance: Point, // AA
    exit: Point,     // ZZ
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct Point {
    x: usize,
    y: usize,
}

impl Point {
    fn add(&self, x: isize, y: isize) -> Point {
        Point {
            x: (self.x as isize + x) as usize,
            y: (self.y as isize + y) as usize,
        }
    }

    fn inside_out(self, other: Point, width: usize, height: usize) -> (Point, Point) {
        if self.closest_side(width, height) > other.closest_side(width, height) {
            (self, other)
        } else {
            (other, self)
        }
    }

    fn closest_side(&self, width: usize, height: usize) -> usize {
        self.x.min(width - self.x).min(self.y).min(height - self.y)
    }
}

impl Maze {
    fn shortest_route(&self, include_depth: bool) -> usize {
        iter::successors(
            Some((
                rbt_set![(self.entrance.clone(), 0)],
                rbt_set![(self.entrance.clone(), 0)],
            )),
            |(explored, locations)| {
                Some(Maze::next_depth_states(
                    explored.clone(),
                    self.next_locations(locations, explored, include_depth),
                ))
            },
        )
        // .inspect(|(explored, states)| eprintln!("{:?}", states))
        .take_while(|(_explored, states)| states.size() > 0)
        .enumerate()
        .find(|(_i, (_explored, states))| {
            states
                .iter()
                .any(|(p_state, depth)| *p_state == self.exit && (!include_depth || *depth == 0))
        })
        .unwrap()
        .0
    }

    fn next_depth_states(
        explored: RedBlackTreeSet<(Point, isize)>,
        locations: RedBlackTreeSet<(Point, isize)>,
    ) -> (
        RedBlackTreeSet<(Point, isize)>,
        RedBlackTreeSet<(Point, isize)>,
    ) {
        (
            locations
                .iter()
                .fold(explored, |acc, next| acc.insert(next.clone())),
            locations,
        )
    }

    fn next_locations(
        &self,
        locations: &RedBlackTreeSet<(Point, isize)>,
        explored: &RedBlackTreeSet<(Point, isize)>,
        include_depth: bool,
    ) -> RedBlackTreeSet<(Point, isize)> {
        locations
            .iter()
            .flat_map(|(p, depth)| {
                [(-1, 0), (1, 0), (0, -1), (0, 1)]
                    .iter()
                    .map(move |(dx, dy)| (p.add(*dx, *dy), *depth))
                    .chain(
                        self.portals
                            .iter()
                            .filter(move |(from, _to)| p == from)
                            .map(move |(_from, to)| (to.clone(), depth + 1)),
                    )
                    .chain(
                        self.portals
                            .iter()
                            .filter(move |(_to, from)| p == from)
                            .map(move |(to, _from)| (to.clone(), depth - 1)),
                    )
            })
            .filter(|(p_next, depth)| {
                !self.walls[p_next.y][p_next.x] && (!include_depth || *depth >= 0)
            })
            .filter(|state| !explored.contains(state))
            .collect()
    }
}

#[test]
fn portal_maze_example_1() {
    let maze: Maze = r"         A           
         A           
  #######.#########  
  #######.........#  
  #######.#######.#  
  #######.#######.#  
  #######.#######.#  
  #####  B    ###.#  
BC...##  C    ###.#  
  ##.##       ###.#  
  ##...DE  F  ###.#  
  #####    G  ###.#  
  #########.#####.#  
DE..#######...###.#  
  #.#########.###.#  
FG..#########.....#  
  ###########.#####  
             Z       
             Z       "
        .split('\n')
        .map(|s| s.to_string())
        .collect::<MazeBuilder>()
        .build();

    assert_eq!(maze.shortest_route(false), 23);
    assert_eq!(maze.shortest_route(true), 26);
}