#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 <deque>
#include <cctype>
#include <queue>
#include <map>

#include <aocpp/Startup.hpp>
#include <aocpp/Coord.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');
}

struct Map {
  std::vector<std::string> rows;

  auto operator[](Coord c)       -> char& { return rows[c.y][c.x]; }
  auto operator[](Coord c) const -> char  { return rows[c.y][c.x]; }

  static auto Parse(std::istream & in) -> Map {
    Map result;
    std::string line;
    while (std::getline(in, line)) {
      result.rows.emplace_back(std::move(line));
    }
    return {result};
  }
};

auto FindFeatures(Map const& map) -> Features {
  Features features;
  for (std::int64_t y = 0; y < map.rows.size(); y++) {
    for (std::int64_t x = 0; x < map.rows[y].size(); x++) {
      auto const c = map[{x,y}];
      if ('#' != c && '.' != c) {
        features[c] = {x,y};
      }
    }
  }
  return features;
}

auto FindDistancesFrom(
  Map const& map,
  Features const& features,
  Distances & distances,
  char const start_letter,
  Coord const start
) {
  std::map<char, std::pair<std::int64_t, Doors>> result;

  std::deque<std::tuple<std::int64_t, Coord, Doors>> todo {{0, start, {}}};

  std::set<Coord> seen;

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

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

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

    // success, we've found a key, record the path
    if (c != start_letter && std::islower(c)) {
        distances[start_letter][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_back(steps+1, fn(here), doors);
    }
  }

  return result;
}

auto FindDistances(Map const& map, Features const& features) {
  Distances distances;
  for (auto const [start_letter, start_coord] : features) {
    if (!std::isupper(start_letter)) {
      FindDistancesFrom(map, features, distances, 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(Map& map, Features & features) {
  auto const start = features['@'];
  for (auto const fn : directions) {
    map[fn(start)] = '#';
  }
  features.erase('@');
  features['^'] =   Up(Left (start));
  features['&'] = Down(Left (start));
  features['*'] =   Up(Right(start));
  features['$'] = Down(Right(start));
}

} // namespace

auto main(int argc, char** argv) -> int {
  auto map = Map::Parse(Startup(argc, argv));
  auto features = FindFeatures(map);
  auto distances = FindDistances(map, features);
  std::cout << "Part 1: " << SolveMaze(distances, "@") << std::endl;

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