/// @file 16.cpp
/// @brief Solution to Advent of Code 2022 Day 16
///
/// This solution follows the following process:
/// 1. Parse the input source into a list of rooms
/// 2. Optimization: put rooms with valves first in the array
/// 3. Compute the shortest paths between each room
/// 4. Enumerate all the paths through the graph to find maximum water flow per valve-set.
/// 5. Use the flow/valve summaries to compute the two answers.

#include <bitset>
#include <cstdint>
#include <iostream>
#include <sstream>
#include <stack>
#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/Parsing.hpp>
#include <aocpp/Startup.hpp>

namespace {

namespace phx = boost::phoenix;
namespace qi = boost::spirit::qi;

/// @brief The starting room as defined in the problem statement
constexpr char const* STARTING_ROOM = "AA";

/// @brief Arbitrary number of valves this solution supports
constexpr std::size_t MAX_VALVES = 64;

/// @brief Array of distances from one node to another.
/// @tparam T distance type
///
/// `distance[i][j]` is the cost to move from node `i` to node `j`
template <typename T>
using distance_array = boost::multi_array<T, 2>;

/// @brief Update single-step distance matrix with transitive shortest paths
/// @tparam T distance type
/// @param[in,out] dist distance matrix
///
/// This implementation uses the Floyd–Warshall algorithm and assumes that
/// there are no negative-cost cycles. It also assumes that a path exists
/// between all pairs of nodes.
template <typename T>
auto ShortestDistances(distance_array<T> & dist) -> void
{
  auto const range = boost::irange(dist.size());
  for (auto const k : range) {
    for (auto const i : range) {
      for (auto const j : range) {
        auto const d_ikj = dist[i][k] + dist[k][j];
        auto & d_ij = dist[i][j];
        if (d_ij > d_ikj) d_ij = d_ikj;
      }
    }
  }
}

/// @brief A single record from the problem input.
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;
};

class Grammar : public qi::grammar<std::string::const_iterator, std::vector<Room>()> {
  qi::rule<iterator_type, std::string()> name;
  qi::rule<iterator_type, Room()> room_description;
  qi::rule<iterator_type, std::vector<Room>()> room_descriptions;

public:
  Grammar() : grammar::base_type{room_descriptions} {
    using namespace qi::labels;
    name = qi::as_string[+qi::alpha];
    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 ] >>
      "\n";
    room_descriptions = *room_description;
  }
};

/// @brief Rearrange the rooms so that those with flows come first
/// @param[in,out] rooms vector of rooms that gets reordered
/// @return number of rooms with with non-zero flows
auto FlowsFirst(
  std::vector<Room> & rooms
) -> std::size_t
{
  using namespace phx::placeholders;
  auto const zeros = boost::range::partition(rooms, phx::bind(&Room::flow, arg1) > 0);
  return std::distance(boost::begin(rooms), zeros);
}

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

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

  auto distances = distance_array<std::uint64_t>{boost::extents[N][N]};

  for (auto const i : boost::irange(rooms.size())) {
    auto di = distances[i];
  
    // Default value: N is longer than any optimal distance by at least 1
    boost::range::fill(di, N);

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

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

  ShortestDistances(distances);

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

/// @brief Bitset used to track which valves have been turned on
using Valves = std::bitset<MAX_VALVES>;

/// @brief Intermediate states for depth-first search in Routes
struct State {
  /// @brief Time remaining
  std::uint64_t time;
  /// @brief Water flow achieved so far
  std::uint64_t flow;
  /// @brief Current actor location
  std::size_t location;
  /// @brief Set of valves already opened
  Valves valves;
};

/// @brief Compute all the flows achievable with a set of values
/// @param[in] start Index of starting room
/// @param[in] initial_time Initial amount of time
/// @param[in] rooms Array of rooms from input file
/// @param[in] 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<std::uint64_t> const& distances
) -> std::unordered_map<Valves, std::uint64_t>
{
  if (rooms.size() > MAX_VALVES) {
    throw std::runtime_error{"Too many valves"};
  }

  // Maximal flow seen at each set of open valves
  auto result = std::unordered_map<Valves, std::uint64_t>{};

  // Remaining states for depth first search
  auto states = std::stack<State>{};
  states.push({initial_time, 0, start, {}});

  while (!states.empty()) {
    auto const state = states.top();
    states.pop();

    if (auto & best = result[state.valves]; best < state.flow) {
      best = state.flow;
    }

    auto const distances_i = distances[state.location];

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

      // don't visit rooms we can't get to in time
      // +1 accounts for the cost of actually turning the valve
      auto const cost = distances_i[j] + 1;
      if (cost >= state.time) { continue; }

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

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

  return result;
}

/// @brief Maximize the water flow using a single actor and 30 minutes
/// @param[in] start Index of the starting room
/// @param[in] rooms Rooms from input file
/// @param[in] 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<std::uint64_t> 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 actors and 26 minutes
/// @param[in] start Index of the starting room
/// @param[in] rooms Rooms from input file
/// @param[in] 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<std::uint64_t> const& distances
) -> std::uint64_t
{
  auto const routes = Routes(start, 26, rooms, distances);
  auto const end = routes.end();

  auto best = std::uint64_t{0};
  for (auto it1 = routes.begin(); it1 != end; ++it1) {
    for (auto it2 = std::next(it1); it2 != end; ++it2) {
      // only consider pairs that have disjoint sets of valves
      if ((it1->first & it2->first).none()) {
        best = std::max(best, it1->second + it2->second);
      }
    }
  }
  return best;
}

} // namespace

/// @brief Print solutions to parts 1 and 2
/// @param[in,out] in input text
/// @param[in,out] out output text
auto Main(std::istream & in, std::ostream & out) -> void
{
  auto rooms = aocpp::ParseGrammar(Grammar{}, in);
  auto const n = FlowsFirst(rooms); // reorders rooms
  auto const [start, distances] = GenerateDistances(rooms);
  rooms.resize(n); // forget about the rooms with no flow
  out << "Part 1: " << Part1(start, rooms, distances) << std::endl;
  out << "Part 2: " << Part2(start, rooms, distances) << std::endl;
}

TEST_SUITE("2022-16") {
  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 out = std::ostringstream{};
    Main(in, out);
    CHECK(out.str() == "Part 1: 1651\nPart 2: 1707\n");
  }

  TEST_CASE("shortest path") {
    auto distances = distance_array<int>{boost::extents[4][4]};
    std::fill_n(distances.data(), distances.num_elements(), 100);

    distances[0][2] = -2;
    distances[0][0] = 0;
    distances[1][0] = 4;
    distances[1][1] = 0;
    distances[1][2] = 3;
    distances[2][2] = 0;
    distances[2][3] = 2;
    distances[3][1] = -1;
    distances[3][3] = 0;
    ShortestDistances(distances);

    CHECK(distances[0][0] == 0);
    CHECK(distances[0][1] == -1);
    CHECK(distances[0][2] == -2);
    CHECK(distances[0][3] == 0);

    CHECK(distances[1][0] == 4);
    CHECK(distances[1][1] == 0);
    CHECK(distances[1][2] == 2);
    CHECK(distances[1][3] == 4);

    CHECK(distances[2][0] == 5);
    CHECK(distances[2][1] == 1);
    CHECK(distances[2][2] == 0);
    CHECK(distances[2][3] == 2);

    CHECK(distances[3][0] == 3);
    CHECK(distances[3][1] == -1);
    CHECK(distances[3][2] == 1);
    CHECK(distances[3][3] == 0);
  }
}