Day 23: Amphipod

use std::env;
use std::io::{prelude::*, BufReader};
use std::fs::File;
use std::collections::VecDeque;
use std::collections::HashMap;
use std::fmt;

use std::collections::BinaryHeap;
use std::cmp::Ordering;

#[macro_use] extern crate cached;

const HALLWAY_ROW: i64 = 1;
// Column number of each amphipod's room
const ROOM_COLUMN_AMBER:  i64 = 3;
const ROOM_COLUMN_BRONZE: i64 = 5;
const ROOM_COLUMN_COPPER: i64 = 7;
const ROOM_COLUMN_DESERT: i64 = 9;

fn amphipod_room_column(species: &AmphipodSpecies) -> i64 {
    match species {
        AmphipodSpecies::Amber  => ROOM_COLUMN_AMBER,
        AmphipodSpecies::Bronze => ROOM_COLUMN_BRONZE,
        AmphipodSpecies::Copper => ROOM_COLUMN_COPPER,
        AmphipodSpecies::Desert => ROOM_COLUMN_DESERT,
    }
}

type AmphipodMap = HashMap<(i64,i64),MapType>;

#[macro_use] extern crate lazy_static;
lazy_static! {
    static ref AMPHI_MAP: AmphipodMap = {
        let args: Vec<String> = env::args().collect();
        let filename = &args[1];
        let file = File::open(filename).expect("Input file not found.");
        let reader = BufReader::new(file);
        // Process input file char by char
        let mut amphimap: AmphipodMap = AmphipodMap::new();
        for (y,line) in reader.lines().enumerate() {
            for (x,c) in line.unwrap().chars().enumerate() {
                let (x,y) = (x as i64, y as i64);
                amphimap.insert((x,y),MapType::from_char_col(c,x));
            }
        }
        amphimap
    };
}

enum MapType {
    OutsideMap,
    Wall,
    Hallway,
    Room,
    Entryway,
}
impl MapType {
    pub fn from_char_col(c: char, col: i64) -> Self {
        match (c, col) {
            (' ',_)       => Self::OutsideMap,
            ('#',_)       => Self::Wall,
            ('A'..='D',_) => Self::Room,
            ('.',col)     => 
                match col {
                    ROOM_COLUMN_AMBER  |
                    ROOM_COLUMN_BRONZE |
                    ROOM_COLUMN_COPPER |
                    ROOM_COLUMN_DESERT => Self::Entryway,
                    _ => Self::Hallway,
                },
            (other,_) => panic!("Unknown map character: {}", other),
        }
    }
}
impl fmt::Display for MapType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::OutsideMap => write!(f, " "),
            Self::Wall       => write!(f, "#"),
            Self::Hallway    => write!(f, "."),
            Self::Room       => write!(f, "_"),
            Self::Entryway   => write!(f, "*"),
        }
    }
}

#[derive(Clone, Eq, PartialEq)]
enum AmphipodSpecies {
    Amber,
    Bronze,
    Copper,
    Desert,
}
impl AmphipodSpecies {
    pub fn from_char(c: char) -> Self {
        match c {
            'A' => Self::Amber ,
            'B' => Self::Bronze,
            'C' => Self::Copper,
            'D' => Self::Desert,
            _ => unreachable!(),
        }
    }
}

#[derive(Clone)]
struct Amphipod {
    species: AmphipodSpecies,
    x: i64,
    y: i64,
}
impl Amphipod {
    pub fn new(x: i64, y: i64, c: char) -> Self {
        match c {
            'A' | 'B' | 'C' | 'D' => Self { x: x, y: y, species: AmphipodSpecies::from_char(c)},
            _ => unreachable!(),
        }
    }
    pub fn energy_cost(&self) -> usize {
        match self.species {
            AmphipodSpecies::Amber  =>    1,
            AmphipodSpecies::Bronze =>   10,
            AmphipodSpecies::Copper =>  100,
            AmphipodSpecies::Desert => 1000,
        }
    }
}
impl fmt::Display for Amphipod {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.species {
            AmphipodSpecies::Amber  => write!(f, "A"),
            AmphipodSpecies::Bronze => write!(f, "B"),
            AmphipodSpecies::Copper => write!(f, "C"),
            AmphipodSpecies::Desert => write!(f, "D"),
        }
    }
}

struct State {
    amphipods: Vec<Amphipod>,
    energy: usize,
}
impl Ord for State {
    fn cmp(&self, other: &Self) -> Ordering {
        other.energy.cmp(&self.energy)
    }
}
impl PartialOrd for State {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(other.energy.cmp(&self.energy))
    }
}
impl Eq for State {}
impl PartialEq for State {
    fn eq(&self, other: &Self) -> bool {
        other.energy == self.energy
    }
}

fn move_from_to(x0: i64, y0: i64, x1: i64, y1: i64, amphipods: &mut Vec<Amphipod>) {
    amphipods.iter_mut().filter(|a| (a.x,a.y) == (x0,y0)).for_each(|mut a| {a.x = x1; a.y = y1;})
}

// Get valid destinations for amphipod moves
fn valid_destinations(x: i64, y: i64, amphipods: &Vec<Amphipod>) -> Vec<(i64,i64)> {
    let mut destinations: Vec<(i64,i64)> = Vec::new();

    let unit = amphipods.iter().find(|a| (a.x,a.y) == (x,y)).unwrap();

    // Handle unit in final room or final room ready for unit
    if let Some(target) = room_target(&unit.species,&amphipods) {
        if amphipod_room_column(&unit.species) == x { return destinations; } // Unit is already in its final room
        if accessible(x,y,target.0,target.1,&amphipods) {
            destinations.push(target);
            return destinations;
        }
    }

    // If unit is currently in the wrong room, valid destinations are:
    // - any accessible hallway spot (handled below)
    // - the deepest cell in target room if target room is ready (handled above)
    // If unit is already in hallway, the only valid spot is the target room (handled above)
    if y != HALLWAY_ROW {
        let (xmax,_) = map_limits();
        for xnew in 1..xmax {
            // Skip entryways
            match xnew {
                ROOM_COLUMN_AMBER  |
                ROOM_COLUMN_BRONZE |
                ROOM_COLUMN_COPPER |
                ROOM_COLUMN_DESERT => { continue; },
                _ => {},
            }

            // Check if no unit in spot and spot is reachable
            if amphipods.iter().find(|a| (a.x,a.y) == (xnew,HALLWAY_ROW)).is_none() {
                if accessible(x,y, xnew,HALLWAY_ROW, &amphipods) {
                    destinations.push((xnew,HALLWAY_ROW));
                }
            }
        }
    }
    destinations
}

// Assess whether a room is ready for amphipods to move into, if so, return deepest cell in room
fn room_target(room: &AmphipodSpecies, amphipods: &Vec<Amphipod>) -> Option<(i64,i64)> {
    let (_,ymax) = map_limits();
    let column = amphipod_room_column(&room);
    let mut target = (column,ymax-1);
    for space in ((HALLWAY_ROW+1)..ymax).rev() {
        match amphipods.iter().find(|a| (a.x,a.y) == (column,space)) {
            Some(amphi) if amphi.species != *room => { return None; },
            Some(_) => { target = (column,space); }
            None => { return Some((column,space)); },
        }
    }
    Some(target)
}

// Assess whether amphipods have moved into their final destinations
fn map_complete(amphipods: &Vec<Amphipod>) -> bool {
    for amphipod in amphipods {
        if amphipod_room_column(&amphipod.species) != amphipod.x { return false; }
    }
    true
}

fn solve(input: &str, debug: bool) {
    let file = File::open(input).expect("Input file not found.");
    let reader = BufReader::new(file);
    
    // Generate initial map
    let mut amphipods: Vec<Amphipod> = Vec::new();
    for (y,line) in reader.lines().enumerate() {
        for (x,c) in line.unwrap().chars().enumerate() {
            let (x,y) = (x as i64, y as i64);
            match c {
                'A' | 'B' | 'C' | 'D' => amphipods.push(Amphipod::new(x,y,c)),
                _ => {},
            }
        }
    }

    // Debug info
    if debug {
        println!("Initial map:");
        print_map(&amphipods);
    }
    
    // Solve
    let mut heap: BinaryHeap<State> = BinaryHeap::new();
    let start = State { amphipods: amphipods, energy: 0 };
    heap.push(start);
    while heap.len() > 0 {

        let state = heap.pop().unwrap();
        // Amphipods are organized
        if map_complete(&state.amphipods) {
            print_map(&state.amphipods);
            println!("Solution: {}", state.energy);
            // Part 1: 14510
            // Part 2: 49180
            break;
        }

        if debug {
            print_map(&state.amphipods);
            println!("Energy: {}", state.energy);
        }

        // Move amphipods
        for amphipod in &state.amphipods {
            let (x,y) = (amphipod.x,amphipod.y);
            for (newx,newy) in valid_destinations(x,y,&state.amphipods) {
                let mut new_map = state.amphipods.clone();
                let energy_delta = amphipod.energy_cost() * distance(x,y,newx,newy);
                let new_energy = state.energy + energy_delta;
                move_from_to(x,y,newx,newy, &mut new_map);
                let new_state = State { amphipods: new_map, energy: new_energy };
                heap.push(new_state);
            }
        }
    }
}

// Determine whether point (x1,y1) is accessible from point (x0,y0) by counting amphipods in the way
fn accessible(x0: i64, y0: i64, x1: i64, y1: i64, amphipods: &Vec<Amphipod>) -> bool {
    amphipods
        .iter()
        .filter(|a| (a.x,a.y) != (x0,y0))                               // exclude self
        .filter(|a| (a.x == x0 && a.y <  y0) ||                         // above self in same room
                    (a.x == x1 && a.y <= y1) ||                         // above target in target room
                    ((x0..=x1).contains(&a.x) && a.y == HALLWAY_ROW) || // in the way in hallway
                    ((x1..=x0).contains(&a.x) && a.y == HALLWAY_ROW))   // in the way in hallway
        .count() == 0
}

// Determine what neighboring cells can be moved to FOR DISTANCE CALCULATION
// PURPOSES ONLY. Not valid for determining whether a point can be a unit's
// final destination because room entryways are valid neighbors but not valid
// final destinations.
cached! { NEIGHBORS; 
    fn neighbors(x: i64, y: i64) -> Vec<(i64,i64)> = {
        let mut destinations: Vec<(i64,i64)> = Vec::new();
        let newpts = vec![(x-1,y  ),(x+1,y  ),(x  ,y-1),(x  ,y+1)];

        for newpt in newpts {
            let map_point = AMPHI_MAP.get(&(newpt.0,newpt.1)).unwrap();
            use MapType::*;
            match map_point {
                OutsideMap | Wall => {}, // not valid
                Hallway | Entryway | Room => destinations.push(newpt),
            }
        }
        destinations
    }
}

struct PointDistance {
    x: i64,
    y: i64,
    d: usize,
}

cached! { DISTANCE; 
    fn distance(x0: i64, y0: i64, x1: i64, y1: i64) -> usize = {
        let (xmax,ymax) = map_limits();
        let mut visited = vec![vec![false; ymax as usize]; xmax as usize];
        visited[x0 as usize][y0 as usize] = true;

        // Insert first point
        let mut q: VecDeque<PointDistance> = VecDeque::new();
        let initial = PointDistance { x: x0, y: y0, d: 0 };
        q.push_back(initial);

        // Queue
        while q.len() > 0 {
            let node = q.pop_front().unwrap();

            // Are we there yet
            if (node.x,node.y) == (x1,y1) { return node.d; }

            // Explore neighbors
            for neighbor in neighbors(node.x, node.y) {
                let (nx,ny) = (neighbor.0 as usize, neighbor.1 as usize);
                if visited[nx][ny] { continue; }
                visited[nx][ny] = true;
                let new_node = PointDistance { x: neighbor.0, y: neighbor.1, d: &node.d+1 };
                q.push_back(new_node);
            }
        }
        std::usize::MAX
    }
}

// Debug: print picture of current amphipod map
fn print_map(amphipods: &Vec<Amphipod>) {
    let (xmax,ymax) = map_limits();
    println!("Map:");
    for y in 0..=ymax {
        for x in 0..=xmax {
            if let Some(unit) = amphipods.iter().find(|a| (a.x,a.y) == (x,y)) {
                print!("{}",unit);
            } else {
                if let Some(cell) = AMPHI_MAP.get(&(x,y)) {
                    print!("{}", cell);
                }
            }
        }
        println!();
    }
    println!();
}

// Return limits of map
cached! { MAP_LIM;
    fn map_limits() -> (i64,i64) = {
        (AMPHI_MAP.keys().map(|k| k.0).max().expect("Failed to determine maximum x-dimension."),
         AMPHI_MAP.keys().map(|k| k.1).max().expect("Failed to determine maximum y-dimension."))
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let filename = &args[1];
    let debug = if args.len() >= 3 { &args[2] == "--debug" } else { false };
    solve(filename, debug);
}