#include <algorithm>
#include <cstdint>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <iterator>
#include <stdexcept>
#include <vector>
#include <set>
#include <string>
#include <tuple>
#include <bitset>
#include <queue>
#include <cctype>
#include <queue>
#include <map>

#include <aocpp/Startup.hpp>
#include <aocpp/Coord.hpp>
#include <aocpp/Grid.hpp>

using namespace aocpp;

// lowercase key
// uppercase door
// start at @
// wall #
// boring .

namespace {

Coord(* const directions[4])(Coord) = {Up, Down, Left, Right};

using Features = std::map<char, Coord>;
using Doors = std::bitset<26>;
using Distances = std::map<char,std::map<char, std::pair<std::int64_t, Doors>>>;

auto SetKey(Doors& doors, char key) {
  return doors.set(std::toupper(key) - 'A');
}

auto FindFeatures(Grid const& grid) -> Features {
  Features features;
  grid.each([&](Coord c, char v) {
    if ('#' != v && '.' != v) {
      features[v] = c;
    }
  });
  return features;
}

auto FindDistancesFrom(
  Grid const& grid,
  char const start_letter,
  Coord const start
) {
  std::map<char, std::pair<std::int64_t, Doors>> result;

  std::queue<std::tuple<std::int64_t, Coord, Doors>> todo;
  todo.push({0, start, {}});

  std::set<Coord> seen;

  for (; !todo.empty(); todo.pop()) {
    auto [steps, here, doors] = todo.front();

    // Don't visit the same coordinate twice
    if (!seen.insert(here).second) continue;

    auto const c = grid[here];
    if (c == '#') continue; // avoid walls

    // success, we've found a key, record the path
    if (c != start_letter && std::islower(c)) {
        result[c] = {steps, doors};
        continue; // don't walk beyond the key
    }

    // Note any keys we encounter on our journey
    if (std::isupper(c)) {
      SetKey(doors, c);
    }

    // Visit all neighbors
    for (auto const fn : directions) {
      todo.emplace(steps+1, fn(here), doors);
    }
  }

  return result;
}

auto FindDistances(Grid const& grid, Features const& features) {
  Distances distances;
  for (auto const [start_letter, start_coord] : features) {
    if (!std::isupper(start_letter)) {
      distances[start_letter] = FindDistancesFrom(grid, start_letter, start_coord);
    }
  }
  return distances;
}

auto SolveMaze(
  Distances const& distances,
  std::string const initial_locations
) -> std::int64_t
{
  // Track current positions and current set of keys in easy to compare form
  using Visited = std::pair<std::string, unsigned long long>;
  std::set<Visited> seen;

  // Priority queue returning lowest path cost states first.
  using PqElt = std::tuple<std::int64_t, std::string, Doors>;
  using PqCmp = decltype([](PqElt const& x, PqElt const& y) {
                return std::get<0>(x) > std::get<0>(y); });
  std::priority_queue<PqElt, std::vector<PqElt>, PqCmp> todo;

  todo.emplace(0, initial_locations, Doors());

  while(!todo.empty()) {
    auto [steps, locations, keys] = todo.top();
    todo.pop();

    if (keys.all()) { return steps; }

    std::sort(locations.begin(), locations.end());
    if (seen.emplace(locations, keys.to_ullong()).second) {

      for (auto& location : locations) {
        auto const save = location;

        for (auto const& [next, costneed] : distances.at(location)) {
          auto const [cost, need] = costneed;
          if ((need & ~keys).none()) { // no missing keys
            location = next;
            auto keys_ = keys; SetKey(keys_, next);
            todo.emplace(steps + cost, locations, keys_);
            location = save;
          }
        }
      }
    }
  }

  throw std::runtime_error{"no solution to part 1"};
}

// Part 2 instructs us to update the map splitting it into 4 quadrants
auto Part2(Grid & grid, Features & features) {
  auto const start = features['@'];
  for (auto const fn : directions) {
    grid[fn(start)] = '#';
  }
  features.erase('@');
  features['^'] =   Up(Left (start));
  features['&'] = Down(Left (start));
  features['*'] =   Up(Right(start));
  features['$'] = Down(Right(start));
}

} // namespace

auto Main(std::istream & in) -> void
{
  auto grid = Grid::Parse(in);
  auto features = FindFeatures(grid);
  auto distances = FindDistances(grid, features);
  std::cout << "Part 1: " << SolveMaze(distances, "@") << std::endl;

  Part2(grid, features);
  distances = FindDistances(grid, features);
  std::cout << "Part 2: " << SolveMaze(distances, "^&*$") << std::endl;
}