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 18: Many-Worlds Interpretation")]
/// Finds the shortest path through a maze with keys
///
/// See https://adventofcode.com/2019/day/18 for details.
struct Opt {}

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

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

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

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 Maze {
    blocks: Vec<Vec<char>>,
    start: BoardState,
    keys: KeySet,
}

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

impl Maze {
    fn from_blocks(blocks: Vec<Vec<char>>) -> Maze {
        Maze {
            start: BoardState {
                robots: blocks
                    .iter()
                    .enumerate()
                    .flat_map(|(y, line)| {
                        line.iter()
                            .enumerate()
                            .filter(|(_x, ch)| **ch == '@')
                            .map(move |(x, _ch)| Location { x, y })
                    })
                    .collect(),
                keys: KeySet::default(),
            },
            keys: blocks
                .iter()
                .flat_map(|line| line.iter())
                .fold(KeySet::default(), |acc, next| acc.add_key(*next)),
            blocks,
        }
    }

    fn block_at(&self, location: &Location) -> Option<char> {
        self.blocks
            .get(location.y)
            .and_then(|line| line.get(location.x))
            .cloned()
    }

    fn shortest_route(&self) -> usize {
        iter::successors(
            Some((
                ExploredStates::default().insert(&self.start),
                rbt_set![self.start.clone()],
            )),
            |(explored, locations)| {
                Some(Maze::next_depth_states(
                    explored.clone(),
                    self.next_locations(locations, explored),
                ))
            },
        )
        // .inspect(|(explored, states)| eprintln!("{:?}", states))
        .take_while(|(_explored, states)| states.size() > 0)
        .enumerate()
        .find(|(_i, (_explored, states))| states.iter().any(|state| state.keys == self.keys))
        .unwrap()
        .0
    }

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

    fn next_locations(
        &self,
        locations: &RedBlackTreeSet<BoardState>,
        explored: &ExploredStates,
    ) -> RedBlackTreeSet<BoardState> {
        locations
            .iter()
            .flat_map(|current| {
                (0..current.robots.len()).flat_map(move |i_robot| {
                    [(-1, 0), (1, 0), (0, -1), (0, 1)]
                        .iter()
                        .map(move |(dx, dy)| current.next(i_robot, *dx, *dy))
                        .map(move |(state, new_location)| {
                            self.block_at(&new_location)
                                .map(|c| state.with_key(c))
                                .unwrap_or(state)
                        })
                })
            })
            .filter(|state| self.is_open(state))
            .filter(|state| !explored.contains(state))
            .collect()
    }

    fn is_open(&self, state: &BoardState) -> bool {
        state.robots.iter().all(|location| {
            self.block_at(location)
                .map(|c| Maze::char_is_open(c, state.keys))
                .unwrap_or(false)
        })
    }

    fn char_is_open(c: char, keys: KeySet) -> bool {
        c != '#' && !keys.door_is_locked(c)
    }
}

#[derive(Default, Debug, Clone)]
struct ExploredStates {
    for_key: RedBlackTreeMap<KeySet, ExploredStatesForKey>,
}

#[derive(Default, Debug, Clone)]
struct ExploredStatesForKey {
    robots: Vector<RedBlackTreeSet<Location>>,
}

impl ExploredStates {
    fn insert(&self, new: &BoardState) -> ExploredStates {
        ExploredStates {
            for_key: self.for_key.insert(
                new.keys.clone(),
                self.for_key
                    .get(&new.keys)
                    .map(|current| ExploredStatesForKey {
                        robots: current
                            .robots
                            .iter()
                            .zip(new.robots.iter())
                            .map(|(current_explored, robot)| current_explored.insert(robot.clone()))
                            .collect(),
                    })
                    .unwrap_or_else(|| ExploredStatesForKey {
                        robots: new
                            .robots
                            .iter()
                            .map(|robot| RedBlackTreeSet::new().insert(robot.clone()))
                            .collect(),
                    }),
            ),
        }
    }

    fn contains(&self, state: &BoardState) -> bool {
        self.for_key
            .get(&state.keys)
            .filter(|for_key| {
                state
                    .robots
                    .iter()
                    .zip(for_key.robots.iter())
                    .all(|(r_state, r_explored)| r_explored.contains(r_state))
            })
            .is_some()
    }
}

#[derive(Default, Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct BoardState {
    robots: Vector<Location>,
    keys: KeySet,
}

impl BoardState {
    fn next(&self, i_robot: usize, dx: isize, dy: isize) -> (BoardState, Location) {
        (
            BoardState {
                robots: self
                    .robots
                    .set(i_robot, self.robots[i_robot].next(dx, dy))
                    .unwrap(),
                ..self.clone()
            },
            self.robots[i_robot].next(dx, dy),
        )
    }

    fn with_key(&self, c: char) -> BoardState {
        BoardState {
            keys: self.keys.add_key(c),
            ..self.clone()
        }
    }
}

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

impl Location {
    fn next(&self, dx: isize, dy: isize) -> Location {
        Location {
            x: (self.x as isize + dx) as usize,
            y: (self.y as isize + dy) as usize,
        }
    }
}

#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct KeySet {
    bitset: u32,
}

impl KeySet {
    fn add_key(self, key: char) -> KeySet {
        KeySet {
            bitset: self.bitset | KeySet::key_to_bitset(key),
        }
    }

    fn door_is_locked(self, door: char) -> bool {
        KeySet::door_to_bitset(door) & (!self.bitset) > 0
    }

    fn key_to_bitset(key: char) -> u32 {
        if key.is_ascii_alphabetic() && key.is_ascii_lowercase() {
            1 << (key as u8 - b'a')
        } else {
            0
        }
    }

    fn door_to_bitset(door: char) -> u32 {
        if door.is_ascii_alphabetic() && door.is_ascii_uppercase() {
            1 << (door as u8 - b'A')
        } else {
            0
        }
    }
}

#[test]
fn doors_are_locked_without_key() {
    assert_eq!(true, KeySet::default().door_is_locked('A'))
}

#[test]
fn doors_are_unlocked_with_key() {
    assert_eq!(false, KeySet::default().add_key('a').door_is_locked('A'))
}

#[test]
fn example_1() {
    let maze: Maze = r"
#########
#b.A.@.a#
#########
"
    .split('\n')
    .map(|s| s.to_string())
    .collect();

    assert_eq!(maze.shortest_route(), 8);
}

#[test]
fn example_2() {
    let maze: Maze = r"
#################
#i.G..c...e..H.p#
########.########
#j.A..b...f..D.o#
########@########
#k.E..a...g..B.n#
########.########
#l.F..d...h..C.m#
#################
"
    .split('\n')
    .map(|s| s.to_string())
    .collect();

    assert_eq!(maze.shortest_route(), 136);
}

#[test]
fn example_3() {
    let maze: Maze = r"
#############
#DcBa.#.GhKl#
#.###@#@#I###
#e#d#####j#k#
###C#@#@###J#
#fEbA.#.FgHi#
#############
"
    .split('\n')
    .map(|s| s.to_string())
    .collect();

    assert_eq!(maze.shortest_route(), 32);
}

#[test]
fn example_4() {
    let maze: Maze = r"
#############
#g#f.D#..h#l#
#F###e#E###.#
#dCba@#@BcIJ#
#############
#nK.L@#@G...#
#M###N#H###.#
#o#m..#i#jk.#
#############
"
    .split('\n')
    .map(|s| s.to_string())
    .collect();

    assert_eq!(maze.shortest_route(), 74);
}