#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>

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

namespace {

using namespace aocpp;
using Map = std::vector<std::string>;
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 ReadMap(std::istream & in) -> Map {
  Map result;
  std::string line;
  while (std::getline(in, line)) {
    result.emplace_back(std::move(line));
  }
  return result;
}

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

auto FindDistancesFrom(
  Map const& map, Features const& features,
  Distances & distances,
  char start_letter,
  Coord 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.y][here.x];
    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 && fn : {Up,Down,Left,Right}) {
      todo.emplace_back(steps+1, fn(here), doors);
    }
  }

  return result;
}

auto FindDistances(Map const& map, Features const& features) -> Distances {
  Distances distances;

  for (auto && [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;

  using QElt = std::tuple<std::int64_t, std::string, Doors>;
  struct Cmp { auto operator()(QElt const& x, QElt const& y) -> bool {
    return std::get<0>(x) > std::get<0>(y);
  }};
  std::priority_queue<QElt, std::vector<QElt>, Cmp> 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 save = location;

        for (auto [next, costneed] : distances.at(location)) {
          auto [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"};
}

auto UpdateMap(Map& map, Features & features) {
  auto start = features['@'];
  map[start.y][start.x] = '#';
  map[start.y][start.x+1] = '#';
  map[start.y][start.x-1] = '#';
  map[start.y+1][start.x] = '#';
  map[start.y-1][start.x] = '#';
  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 = ReadMap(Startup(argc, argv));
  auto features = FindFeatures(map);
  auto distances = FindDistances(map, features);

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