Day 23: Amphipod

impl Burrow {
	fn solve<const N: usize>(&self, initial_locs: &AmphipodIndexed<Point, N>) -> Cost {
		const N_KINDS: usize = AmphipodKind::n_kinds();

		#[derive(Debug)]
		struct Update {
			amphipod: Amphipod,
			n_steps: usize,
			new_state: SingleAmphipodState,
		}

		#[derive(Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord, Debug)]
		struct SingleAmphipodState {
			loc: Point,
			is_done: bool,
		}
		type AllAmphipodStates<const N: usize> = AmphipodIndexed<SingleAmphipodState, N>;

		fn enqueue<H: std::hash::BuildHasher + Default, const N: usize>(
			seen: &mut BTreeSet<AllAmphipodStates<N>>,
			pq: &mut PriorityQueue<AllAmphipodStates<N>, Cost, H>,
			state: AllAmphipodStates<N>,
			cost: Cost,
		) {
			if seen.contains(&state) {
				return;
			}
			seen.insert(state);
			pq.push_increase(state, cost);
		}

		let hallway_row = 1;
		let sideroom_max_row = self.tiles.nrows() - 2;

		let mut pq = PriorityQueue::<_, _, DefaultHashBuilder>::with_default_hasher();
		let mut seen = BTreeSet::new();

		let tiles = &self.tiles;

		let initial_state = AmphipodIndexed(initial_locs.0.map(|loc| SingleAmphipodState {
			loc,
			is_done: false,
		}));

		pq.push(initial_state, Cost(0));

		let mut updates = Vec::new();

		while let Some((total_state, cost)) = pq.pop() {
			updates.clear();

			if total_state
				.iter_items()
				.all(|(am, state)| tiles[state.loc] == Tile::SideRoom(am.kind))
			{
				return cost;
			}

			let amphipod_locs = total_state
				.iter_items()
				.map(|(am, state)| (state.loc, am))
				.collect::<Map<_, _>>();

			let n_in_correct_sideroom_by_kind = {
				let mut arr = [0_usize; N_KINDS];
				for (am, state) in total_state.iter_items() {
					if tiles[state.loc] == Tile::SideRoom(am.kind) {
						arr[am.kind as usize] += 1;
					}
				}
				arr
			};

			let siderooms_available_for_amphipods =
				[0, 1, 2, 3].zip([3, 5, 7, 9_usize]).map(|(i, col)| {
					(hallway_row + 1..=sideroom_max_row).all(|row| {
						amphipod_locs
							.get(&[row, col])
							.map_or(true, |am| am.kind as usize == i)
					})
				});

			let update_that_moves_am_to_final_loc =
				total_state.iter_items().find_map(|(am, &state)| {
					if state.is_done || !siderooms_available_for_amphipods[am.kind as usize] {
						return None;
					}

					let curr_loc = state.loc;
					let dest_col = am.kind.dest_col();
					let sideroom_last_empty_row =
						sideroom_max_row - n_in_correct_sideroom_by_kind[am.kind as usize];
					let new_loc = [sideroom_last_empty_row, dest_col];

					if curr_loc[COL] == dest_col && curr_loc[ROW] >= sideroom_last_empty_row {
						return Some(Update {
							amphipod: am,
							n_steps: 0,
							new_state: SingleAmphipodState {
								loc: curr_loc,
								is_done: true,
							},
						});
					}

					let [curr_row, curr_col] = curr_loc;

					for row in hallway_row..curr_row {
						if amphipod_locs.contains_key(&[row, curr_col]) {
							return None;
						}
					}

					#[allow(clippy::range_minus_one)]
					let col_range = if curr_col < dest_col {
						curr_col + 1..=dest_col
					} else {
						dest_col..=curr_col - 1
					};
					for col in col_range {
						if amphipod_locs.contains_key(&[hallway_row, col]) {
							return None;
						}
					}

					let n_steps_up = curr_row - hallway_row;
					let n_steps_sideways = abs_diff(curr_col, dest_col);
					let n_steps_down = sideroom_last_empty_row - hallway_row;
					let n_steps = n_steps_up + n_steps_sideways + n_steps_down;

					Some(Update {
						amphipod: am,
						n_steps,
						new_state: SingleAmphipodState {
							loc: new_loc,
							is_done: true,
						},
					})
				});

			if let Some(update) = update_that_moves_am_to_final_loc {
				updates.push(update);
			} else {
				for (am, &state) in total_state.iter_items() {
					if state.is_done {
						continue;
					}

					let curr_loc = state.loc;
					let curr_tile = tiles[curr_loc];

					if !(matches!(curr_tile, Tile::SideRoom(_))
						&& (hallway_row..curr_loc[ROW])
							.all(|row| !amphipod_locs.contains_key(&[row, curr_loc[COL]])))
					{
						continue;
					}

					let n_steps_to_hallway = curr_loc[ROW] - hallway_row;

					// Rust...
					for range in [
						&mut (0..curr_loc[COL]).rev(),
						&mut (curr_loc[COL] + 1..tiles.ncols()),
					] as [&mut dyn Iterator<Item = usize>; 2]
					{
						for col in range {
							let new_loc = [hallway_row, col];
							if amphipod_locs.contains_key(&new_loc) || tiles[new_loc] == Tile::Wall
							{
								break;
							}
							if tiles[new_loc] == Tile::Doorway {
								continue;
							}
							updates.push(Update {
								amphipod: am,
								n_steps: n_steps_to_hallway + abs_diff(col, curr_loc[COL]),
								new_state: SingleAmphipodState {
									loc: new_loc,
									is_done: false,
								},
							});
						}
					}
				}
			}

			for &Update {
				amphipod: am,
				n_steps,
				new_state,
			} in &updates
			{
				let mut new_total_state = total_state;
				new_total_state[am] = new_state;

				enqueue(
					&mut seen,
					&mut pq,
					new_total_state,
					Cost(cost.0 + n_steps * am.kind.energy()),
				);
			}
		}

		panic!("Could not find a path to the finish line!");
	}
}

pub fn ans() -> Answer<usize, usize> {
	(23, (pt1(), pt2())).into()
}

fn pt1() -> usize {
	let b = InstantiatedBurrow::<8>::from_str(include_str!("input_1.txt")).unwrap();
	b.burrow.solve(&b.amphipod_locs).0
}

fn pt2() -> usize {
	let b = InstantiatedBurrow::<16>::from_str(include_str!("input_2.txt")).unwrap();
	b.burrow.solve(&b.amphipod_locs).0
}