#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 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 c = map[y][x];
      switch (c) {
        default: features[c] = Coord{x,y};
        case '#':
        case '.': {}
      }
    }
  }
  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::set<Coord> seen;
  std::deque<std::tuple<std::int64_t, Coord, Doors>> todo
    {{0, start, {}}};

  while (!todo.empty()) {
    auto [steps, here, doors] = todo.front();
    todo.pop_front();
    if (seen.insert(here).second) {
      auto c = map[here.y][here.x];
      if (c == '#') {
        continue;
      }
      if (c != start_letter && std::islower(c)) {
        distances[start_letter][c] = {steps, doors};
        continue;
      }
      if (std::isupper(c)) {
        doors.set(c - 'A');
      }
      for (auto && fn : {Up,Down,Left,Right}) {
        auto here_ = fn(here);
        todo.push_back({steps+1, 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 HasKey(Doors doors, char key) {
  return doors.test(std::toupper(key) - 'A');
}

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

template <class Elt, class Compare>
auto MakePQueue(Compare const& compare)
  -> std::priority_queue<Elt, std::vector<Elt>, Compare>
{
  return {compare};
}

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

  auto todo = MakePQueue<std::tuple<std::int64_t, std::set<char>, Doors>>(
    [](auto const& x, auto const& y) {
      return std::get<0>(x) > std::get<0>(y);
    }
  );

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

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

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

    if (seen.insert({locations,keys.to_ullong()}).second) {
      for (auto location : locations) {
        auto locations1 = locations; locations1.erase(location);

        for (auto && [next, costneed] : distances.at(location)) {
          auto [cost, need] = costneed;
          if ((need & ~keys).none()) {
            auto locations2 = locations1; locations2.insert(next);
            todo.emplace(steps + cost, locations2, SetKey(keys, next));
          }
        }
      }
    }
  }

  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;
}