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

#include <doctest.h>

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

using namespace aocpp;

namespace {

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

using Name = std::string;
using Portals = std::map<std::string, Coord>;
using Distances = std::map<Name,std::vector<std::pair<Name, std::int64_t>>>;

auto FindPortals(Grid const& grid) -> Portals {
  Portals portals;

  std::int64_t w = grid.Cols();
  std::int64_t h = grid.Rows();

  for (std::int64_t x = 1; x < w-1; x++) {
    for (std::int64_t y = 1; y < h-1; y++) {
      Coord const c = {x,y};
      auto const v = grid[c];
      if (std::isupper(v)) {
        char polarity = x==1 || x==w-2 || y==1 || y==h-2
                      ? '-' : '+';
        if (grid[Up(c)] == '.') {
          portals[{polarity, v, grid[Down(c)]}] = Up(c);
        } else if (grid[Down(c)] == '.') {
          portals[{polarity, grid[Up(c)], v}] = Down(c);
        } else if (grid[Left(c)] == '.') {
          portals[{polarity, v, grid[Right(c)]}] = Left(c);
        } else if (grid[Right(c)] == '.') {
          portals[{polarity, grid[Left(c)], v}] = Right(c);
        }
      }
    }
  }
  return portals;
}

auto FindDistancesFrom(
  Grid const& grid,
  std::map<Coord, Name> const& names,
  std::string const start_name,
  Coord const start
) {
  std::vector<std::pair<Name, std::int64_t>> result;
  std::set<Coord> seen;
  for (
    std::queue<std::pair<std::int64_t, Coord>> todo
      {decltype(todo)::container_type{{{0, start}}}};
    !todo.empty();
    todo.pop()) {
    auto const [steps, here] = 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 (auto it = names.find(here); it != names.end() && it->second != start_name) {
        result.emplace_back(it->second, steps);
        continue; // don't walk beyond the portal
    }

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

  return result;
}

auto OtherName(Name name) {
  name[0] = name[0] == '+' ? '-' : '+';
  return name;
}

auto FindDistances(Grid const& grid, Portals const& portals) {
  Distances distances;

  std::map<Coord, Name> names;
  for (auto const& [k,v] : portals) {
    names[v] = k;
  }

  for (auto const& [start_name, start_coord] : portals) {
    distances[start_name] = FindDistancesFrom(grid, names, start_name, start_coord);
  }

  return distances;
}

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

  // Priority queue returning lowest path cost states first.
  using PqElt = std::tuple<std::int64_t, std::int64_t, Name>;
  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, 0, "-AA");

  while(!todo.empty()) {
    auto const [steps, depth, name] = std::move(todo.top());
    todo.pop();
    if (name == "-ZZ") { return steps; }
    if (seen.emplace(depth, name).second) {
      if (auto const it = distances.find(name); it != distances.end()) {
        for (auto const& [next, cost] : it->second) {
          if (next == "-ZZ") {
            if (depth == 0) todo.emplace(steps + cost, depth, "-ZZ");
          } else {
            auto const depth_ = depth + (recursive ? (next[0]=='+' ? 1 : -1) : 0);
            if (depth_ >= 0) todo.emplace(steps + cost + 1, depth_, OtherName(next));
          }
        }
      }
    }
  }
  throw std::runtime_error{"no solution"};
}

} // namespace

TEST_SUITE("2019-20 examples") {
  TEST_CASE("part 1 a") {
    std::istringstream in {
      "         A           \n"
      "         A           \n"
      "  #######.#########  \n"
      "  #######.........#  \n"
      "  #######.#######.#  \n"
      "  #######.#######.#  \n"
      "  #######.#######.#  \n"
      "  #####  B    ###.#  \n"
      "BC...##  C    ###.#  \n"
      "  ##.##       ###.#  \n"
      "  ##...DE  F  ###.#  \n"
      "  #####    G  ###.#  \n"
      "  #########.#####.#  \n"
      "DE..#######...###.#  \n"
      "  #.#########.###.#  \n"
      "FG..#########.....#  \n"
      "  ###########.#####  \n"
      "             Z       \n"
      "             Z       \n"
    };
    auto map = Grid::Parse(in);
    auto portals = FindPortals(map);
    auto distances = FindDistances(map, portals);
    CHECK(SolveMaze(distances, false) == 23);
  }

  TEST_CASE("part 1") {
    std::istringstream in {
      "                   A               \n"
      "                   A               \n"
      "  #################.#############  \n"
      "  #.#...#...................#.#.#  \n"
      "  #.#.#.###.###.###.#########.#.#  \n"
      "  #.#.#.......#...#.....#.#.#...#  \n"
      "  #.#########.###.#####.#.#.###.#  \n"
      "  #.............#.#.....#.......#  \n"
      "  ###.###########.###.#####.#.#.#  \n"
      "  #.....#        A   C    #.#.#.#  \n"
      "  #######        S   P    #####.#  \n"
      "  #.#...#                 #......VT\n"
      "  #.#.#.#                 #.#####  \n"
      "  #...#.#               YN....#.#  \n"
      "  #.###.#                 #####.#  \n"
      "DI....#.#                 #.....#  \n"
      "  #####.#                 #.###.#  \n"
      "ZZ......#               QG....#..AS\n"
      "  ###.###                 #######  \n"
      "JO..#.#.#                 #.....#  \n"
      "  #.#.#.#                 ###.#.#  \n"
      "  #...#..DI             BU....#..LF\n"
      "  #####.#                 #.#####  \n"
      "YN......#               VT..#....QG\n"
      "  #.###.#                 #.###.#  \n"
      "  #.#...#                 #.....#  \n"
      "  ###.###    J L     J    #.#.###  \n"
      "  #.....#    O F     P    #.#...#  \n"
      "  #.###.#####.#.#####.#####.###.#  \n"
      "  #...#.#.#...#.....#.....#.#...#  \n"
      "  #.#####.###.###.#.#.#########.#  \n"
      "  #...#.#.....#...#.#.#.#.....#.#  \n"
      "  #.###.#####.###.###.#.#.#######  \n"
      "  #.#.........#...#.............#  \n"
      "  #########.###.###.#############  \n"
      "           B   J   C               \n"
      "           U   P   P               \n"
    };
    auto map = Grid::Parse(in);
    auto portals = FindPortals(map);
    auto distances = FindDistances(map, portals);
    CHECK(SolveMaze(distances, false) == 58);
  }

  TEST_CASE("part 2") {
    std::istringstream in {
      "             Z L X W       C                 \n"
      "             Z P Q B       K                 \n"
      "  ###########.#.#.#.#######.###############  \n"
      "  #...#.......#.#.......#.#.......#.#.#...#  \n"
      "  ###.#.#.#.#.#.#.#.###.#.#.#######.#.#.###  \n"
      "  #.#...#.#.#...#.#.#...#...#...#.#.......#  \n"
      "  #.###.#######.###.###.#.###.###.#.#######  \n"
      "  #...#.......#.#...#...#.............#...#  \n"
      "  #.#########.#######.#.#######.#######.###  \n"
      "  #...#.#    F       R I       Z    #.#.#.#  \n"
      "  #.###.#    D       E C       H    #.#.#.#  \n"
      "  #.#...#                           #...#.#  \n"
      "  #.###.#                           #.###.#  \n"
      "  #.#....OA                       WB..#.#..ZH\n"
      "  #.###.#                           #.#.#.#  \n"
      "CJ......#                           #.....#  \n"
      "  #######                           #######  \n"
      "  #.#....CK                         #......IC\n"
      "  #.###.#                           #.###.#  \n"
      "  #.....#                           #...#.#  \n"
      "  ###.###                           #.#.#.#  \n"
      "XF....#.#                         RF..#.#.#  \n"
      "  #####.#                           #######  \n"
      "  #......CJ                       NM..#...#  \n"
      "  ###.#.#                           #.###.#  \n"
      "RE....#.#                           #......RF\n"
      "  ###.###        X   X       L      #.#.#.#  \n"
      "  #.....#        F   Q       P      #.#.#.#  \n"
      "  ###.###########.###.#######.#########.###  \n"
      "  #.....#...#.....#.......#...#.....#.#...#  \n"
      "  #####.#.###.#######.#######.###.###.#.#.#  \n"
      "  #.......#.......#.#.#.#.#...#...#...#.#.#  \n"
      "  #####.###.#####.#.#.#.#.###.###.#.###.###  \n"
      "  #.......#.....#.#...#...............#...#  \n"
      "  #############.#.#.###.###################  \n"
      "               A O F   N                     \n"
      "               A A D   M                     \n"
    };
    auto map = Grid::Parse(in);
    auto portals = FindPortals(map);
    auto distances = FindDistances(map, portals);
    CHECK(SolveMaze(distances, true) == 396);
  }
}


auto main(int argc, char** argv) -> int {
  auto map = Grid::Parse(*Startup(argc, argv));
  auto portals = FindPortals(map);
  auto distances = FindDistances(map, portals);

  std::cout << "Part 1: " << SolveMaze(distances, false) << std::endl;
  std::cout << "Part 2: " << SolveMaze(distances, true) << std::endl;
}