#include <bitset>
#include <cstdint>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>

#include <boost/multi_array.hpp>
#include <boost/phoenix.hpp>
#include <boost/range/adaptors.hpp>
#include <boost/range/algorithm.hpp>
#include <boost/range/irange.hpp>
#include <boost/spirit/include/qi.hpp>

#include <doctest.h>

#include <aocpp/Startup.hpp>

namespace {

using distance_array = boost::multi_array<std::uint64_t, 2>;

/// @brief Update distance matrix with transitive shortest paths
/// @param dist Single-step distances between nodes (must be square)
auto ShortestDistances(distance_array & dist) -> void
{
  // Floyd–Warshall_algorithm
  auto const range = boost::irange(dist.size());
  for (auto const k : range) {
    for (auto const i : range) {
      for (auto const j : range) {
        auto const new_dist = dist[i][k] + dist[k][j];
        auto & old_dist = dist[i][j];
        if (old_dist > new_dist) old_dist = new_dist;
      }
    }
  }
}

struct Room {
  /// @brief Name of the room
  std::string name;
  /// @brief Flow rate of the valve in the room
  std::uint64_t flow;
  /// @brief Directly adjacent rooms
  std::vector<std::string> connections;
};

auto Parse(std::istream & in) -> std::vector<Room>
{
  std::vector<Room> result;
  std::string line;
  while (std::getline(in, line)) {
    namespace phx = boost::phoenix;
    namespace qi = boost::spirit::qi;
    using namespace qi::labels;
    using It = std::string::const_iterator;

    qi::rule<It, std::string()> const name = qi::as_string[+qi::alpha];
    qi::rule<It, Room()> const room_description =
      "Valve " >>
      name [ phx::bind(&Room::name, _val) = _1 ] >>
      " has flow rate=" >>
      qi::ulong_long [ phx::bind(&Room::flow, _val) = _1 ] >>
      "; tunnel"  >> -qi::string("s") >>
      " lead"     >> -qi::string("s") >>
      " to valve" >> -qi::string("s") >>
      " " >>
      (name % ", ") [ phx::bind(&Room::connections, _val) = _1 ];

    It b = line.begin();
    It e = line.end();
    result.emplace_back();
    if (!qi::parse(b, e, room_description, result.back())) {
      throw std::runtime_error{"bad input line"};
    }
  }
  return result;
}

/// @brief Computes the distances between rooms and finds the address of the starting room
/// @returns starting index and distances
auto GenerateDistances(
  std::vector<Room> const& rooms
) -> std::pair<std::size_t, distance_array>
{
  auto const N = rooms.size();

  // Associate the names and indexes of each room
  std::unordered_map<std::string, std::size_t> names;
  for (auto [i,room] : rooms | boost::adaptors::indexed()) {
    names[room.name] = i;
  }

  distance_array distances{boost::extents[N][N]};

  // N is longer than any optimal distance by at least 1
  std::fill_n(distances.data(), distances.num_elements(), N);

  for (auto [i, room] : rooms | boost::adaptors::indexed()) {
    auto di = distances[i];

    // each room is one away from adjacent rooms
    for (auto const& name : room.connections) {
      di[names[name]] = 1;
    }

    // zero distance to self
    di[i] = 0;
  }

  ShortestDistances(distances);

  return {names.at("AA"), std::move(distances)};
}

using Valves = std::bitset<64>;

struct State {
  std::uint64_t time;
  std::uint64_t flow;
  std::size_t location;
  Valves valves;
};

/// @brief Compute all the flows achievable with a set of values
/// @param start Index of starting room
/// @param initial_time Initial amount of time
/// @param rooms Array of rooms from input file
/// @param distances Shortest paths between all pairs of rooms
/// @return Mapping of maximum flow achievable using a particular set of valves
auto Routes(
  std::size_t const start,
  std::uint64_t const initial_time,
  std::vector<Room> const& rooms,
  distance_array const& distances
) -> std::unordered_map<Valves, std::uint64_t>
{
  std::vector<State> states { State{initial_time, 0, start, {}} };
  std::unordered_map<Valves, std::uint64_t> result;

  while (!states.empty()) {
    auto const state = states.back();
    states.pop_back();
    auto const i = state.location;

    for (auto [j, room] : rooms | boost::adaptors::indexed()) {
      // don't revisit a valve
      if (state.valves.test(j)) { continue; }

      // don't visit rooms with useless valves
      if (room.flow == 0) { continue; }

      // don't visit rooms we can't get to in time
      auto const cost = distances[i][j];
      if (cost+1 >= state.time) { continue; }

      auto const time = state.time - cost - 1;
      auto const flow = state.flow + room.flow * time;
      auto valves = state.valves;
      valves.set(j);

      // remember the best we've seen for this valve-set
      if (result[valves] < flow) {
        result[valves] = flow;
      }

      states.push_back({time, flow, static_cast<std::size_t>(j), valves});
    }
  }

  return result;
}

/// @brief Maximize the water flow using a single actor and 30 minutes
/// @param start Index of the starting room
/// @param rooms Rooms from input file
/// @param distances Shortest distances between pairs of rooms
/// @return Maximum flow achievable
auto Part1(
  std::size_t const start,
  std::vector<Room> const& rooms,
  distance_array const& distances
) -> std::uint64_t
{
  auto const routes = Routes(start, 30, rooms, distances);
  return *boost::range::max_element(routes | boost::adaptors::map_values);
}

/// @brief Maximize the water flow using two actos and 26 minutes
/// @param start Index of the starting room
/// @param rooms Rooms from input file
/// @param distances Shortest distances between pairs of rooms
/// @return Maximum flow achievable
auto Part2(
  std::size_t const start,
  std::vector<Room> const& rooms,
  distance_array const& distances
) -> std::uint64_t
{
  auto const routes = Routes(start, 26, rooms, distances);
  auto const end = routes.end();
  std::uint64_t result {0};
  for (auto it1 = routes.begin(); it1 != end; std::advance(it1, 1)) {
    for (auto it2 = std::next(it1); it2 != end; std::advance(it2, 1)) {
      // only consider pairs that have disjoint sets of valves
      if ((it1->first & it2->first).none()) {
        result = std::max(result, it1->second + it2->second);
      }
    }
  }
  return result;
}

} // namespace

TEST_SUITE("2022-16 examples") {
  TEST_CASE("example") {
    std::istringstream in {
R"(Valve AA has flow rate=0; tunnels lead to valves DD, II, BB
Valve BB has flow rate=13; tunnels lead to valves CC, AA
Valve CC has flow rate=2; tunnels lead to valves DD, BB
Valve DD has flow rate=20; tunnels lead to valves CC, AA, EE
Valve EE has flow rate=3; tunnels lead to valves FF, DD
Valve FF has flow rate=0; tunnels lead to valves EE, GG
Valve GG has flow rate=0; tunnels lead to valves FF, HH
Valve HH has flow rate=22; tunnel leads to valve GG
Valve II has flow rate=0; tunnels lead to valves AA, JJ
Valve JJ has flow rate=21; tunnel leads to valve II
)"};
    auto const rooms = Parse(in);
    auto const [start, distances] = GenerateDistances(rooms);
    CHECK(1651 == Part1(start, rooms, distances));
    CHECK(1707 == Part2(start, rooms, distances));
  }
}

auto main(int argc, char** argv) -> int
{
  auto const rooms = Parse(*aocpp::Startup(argc, argv));
  auto const [start, distances] = GenerateDistances(rooms);
  std::cout << "Part 1: " << Part1(start, rooms, distances) << std::endl;
  std::cout << "Part 2: " << Part2(start, rooms, distances) << std::endl;
}