use std::fs;
use std::path::Path;
use std::collections::{BinaryHeap,HashMap};
use std::fmt::{Display,Error,Formatter};

type Amphipod = u8;
type Cost = u32;

#[derive(Debug,Clone,PartialEq,Eq,PartialOrd,Ord,Hash)]
struct Room<const SIZE: usize> {
    id: Amphipod,
    layout: [Option<Amphipod>;SIZE]
}

impl<const SIZE: usize> Room<SIZE> {
    fn entrance(&self) -> usize {
        self.id as usize * 2 + 2
    }

    fn solved(&self) -> bool {
        self.layout.iter().all(|a| a == &Some(self.id))
    }

    fn accept(&mut self) -> Option<Cost> {
        if let Some(tile) = (0..SIZE).find(|&i| self.layout[i].is_none() && self.layout[i+1..].iter().all(|a| a == &Some(self.id))) {
            self.layout[tile] = Some(self.id);
            Some(1+tile as Cost)
        } else { None }
    }

    fn take(&mut self) -> Option<(Amphipod,Cost)> {
        for i in 0..SIZE {
            if let Some(amphipod) = self.layout[i] {
                if amphipod != self.id || self.layout[i+1..].iter().any(|a| a != &Some(self.id)) {
                    return self.layout[i].take().map(|a| (a,1+i as Cost))
                }
            }
        }
        None
    }
}

#[derive(Debug,Clone,PartialEq,Eq,PartialOrd,Ord,Hash)]
struct Corridor {
    layout: [Option<Amphipod>;11]
}

impl Corridor {
    fn clear(&self, from: usize, to: usize) -> Option<Cost> {
        if from < to && (from+1..=to).all(|i| self.layout[i] == None) {
            Some((to-from) as Cost)
        } else if to < from && (to..from).all(|i| self.layout[i] == None) { 
            Some((from-to) as Cost)
        } else { None }
    }

    fn accept(&mut self, tile: usize, amphipod: Amphipod) {
        self.layout[tile] = Some(amphipod);
    }

    fn take(&mut self, tile: usize) -> Option<Amphipod> {
        self.layout[tile].take()
    }
}

#[derive(Debug,Clone,PartialEq,Eq,PartialOrd,Ord,Hash)]
struct Cave<const SIZE: usize> {
    rooms: [Room<SIZE>;4],
    corridor: Corridor
}

impl<const SIZE: usize> Cave<SIZE> {
    fn moves(&self) -> Vec<(Cave<SIZE>, Cost)> {
        let mut moves = vec![];
        for tile in 0..self.corridor.layout.len() {
            let mut cave = self.clone();
            if let Some(amphipod) = cave.corridor.take(tile) {
                if let Some(corridor_cost) = cave.corridor.clear(tile, cave.rooms[amphipod as usize].entrance()) {
                    if let Some(in_cost) = cave.rooms[amphipod as usize].accept() {
                        moves.push((cave, 10_u32.pow(amphipod as Cost) * (corridor_cost + in_cost)));
                    }
                }
            }
        }
        for room in 0..self.rooms.len() {
            let mut cave = self.clone();
            if let Some((amphipod, out_cost)) = cave.rooms[room].take() {
                if let Some(corridor_cost) = cave.corridor.clear(cave.rooms[room].entrance(), cave.rooms[amphipod as usize].entrance()) {
                    if let Some(in_cost) = cave.rooms[amphipod as usize].accept() {
                        moves.push((cave, 10_u32.pow(amphipod as Cost) * (out_cost + corridor_cost + in_cost)));
                        continue;
                    }
                }
                for tile in (0..self.corridor.layout.len()).filter(|&i| i < 2 || i > 8 || (i-2)%2 == 1) {
                    let mut cave = cave.clone();
                    if let Some(corridor_cost) = cave.corridor.clear(cave.rooms[room].entrance(), tile) {
                        cave.corridor.accept(tile, amphipod);
                        moves.push((cave, 10_u32.pow(amphipod as Cost) * (out_cost + corridor_cost)));
                    }
                }
            }
        }
        moves
    }

    fn solved(&self) -> bool {
        self.rooms.iter().all(|r| r.solved())
    }

    fn solve(&mut self) -> u32 {
        let mut heap = BinaryHeap::new();
        let mut done = HashMap::new();
        heap.push(Item::new(self.clone(),0));

        while let Some(item) = heap.pop() {
            if item.cave.solved() {
                return item.cost;
            }
            let moves = item.cave.moves();
            for (mv,cost) in moves {
                let new_cost = item.cost + cost;
                let entry = done.entry(mv.clone()).or_insert(Cost::MAX);
                if *entry > new_cost {
                    *entry = new_cost;
                    heap.push(Item::new(mv, new_cost));
                }
            }
        }
        0
    }
}

impl Cave<2> {
    fn unfold(&self) -> Cave<4> {
        let r0 = self.rooms[0].layout;
        let r1 = self.rooms[1].layout;
        let r2 = self.rooms[2].layout;
        let r3 = self.rooms[3].layout;
        Cave {
            rooms: [
                Room { id: 0, layout: [r0[0], Some(3), Some(3), r0[1]] },
                Room { id: 1, layout: [r1[0], Some(2), Some(1), r1[1]] },
                Room { id: 2, layout: [r2[0], Some(1), Some(0), r2[1]] },
                Room { id: 3, layout: [r3[0], Some(0), Some(2), r3[1]] },
            ],
            corridor: Corridor {
                layout: [None;11]
            }
        }

    }
}

impl<const SIZE: usize> Display for Cave<SIZE> {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
        write!(f,"#############\n")?;
        write!(f,"#{}#\n", self.corridor.layout.iter().map(|a| if let Some(a) = a { (a + b'A') as char } else { '.' }).collect::<String>())?;
        for i in 0..SIZE {
            write!(f,"###{}#{}#{}#{}###\n",
                   if let Some(a) = self.rooms[0].layout[i] { (a + b'A') as char } else { '.' },
                   if let Some(a) = self.rooms[1].layout[i] { (a + b'A') as char } else { '.' },
                   if let Some(a) = self.rooms[2].layout[i] { (a + b'A') as char } else { '.' },
                   if let Some(a) = self.rooms[3].layout[i] { (a + b'A') as char } else { '.' })?;
        }
        write!(f,"#############\n")?;
        Ok(())
    }
}

impl From<&str> for Cave<2> {
    fn from(s: &str) -> Self {
        let mut lines = s.lines();
        lines.next(); lines.next();
        let r0: Vec<_> = lines.next().unwrap().bytes().filter_map(|b| if b != b'#' { Some(b - b'A') } else { None }).collect();
        let r1: Vec<_> = lines.next().unwrap().bytes().filter_map(|b| if b != b'#' && b != b' ' { Some(b - b'A') } else { None }).collect();
        Cave {
            rooms: [
                Room { id: 0, layout: [Some(r0[0]),Some(r1[0])] },
                Room { id: 1, layout: [Some(r0[1]),Some(r1[1])] },
                Room { id: 2, layout: [Some(r0[2]),Some(r1[2])] },
                Room { id: 3, layout: [Some(r0[3]),Some(r1[3])] },
            ],
            corridor: Corridor {
                layout: [None;11]
            }
        }
    }
}

#[derive(PartialEq,Eq)]
struct Item<T> {
    cave: T,
    cost: Cost
}

impl<T> Item<T> {
    fn new(cave: T, cost: Cost) -> Self {
        Item { cave, cost }
    }
}

impl<T: Eq> PartialOrd for Item<T> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    } 
}

impl<T: Eq> Ord for Item<T> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        other.cost.cmp(&self.cost)
    } 
}

fn main() {
    let input = Path::new("resources").join("input.txt");
    let content = fs::read_to_string(input).expect("Unable to read input file");
    println!("Ex1: the result is {}", Cave::from(&content[..]).solve());
    println!("Ex2: the result is {}", Cave::from(&content[..]).unfold().solve());
}

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

    const INPUT: &str = "#############
#...........#
###B#C#B#D###
  #A#D#C#A#
  #########";

    #[test]
    fn example() {
        assert_eq!(12521, Cave::from(INPUT).solve());
    }

    #[test]
    fn example_unfolded() {
        assert_eq!(44169, Cave::from(INPUT).unfold().solve());
    }
}