/// @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 <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 {

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

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

/// @brief Parse the input file
/// @param[in,out] in input stream
/// @return Vector of parsed rooms, one per input line
///
/// The parser will consume input until the end of the stream.
///
/// Input lines should follow the following pattern:
/// * Valve **name** has flow rate= **number** ; tunnels lead to valves **name**, **name** ...
/// * Valve **name** has flow rate= **number** ; tunnel leads to valve **name**
auto Parse(std::istream & in) -> std::vector<Room>
{
  std::vector<Room> result;
  std::string line;
  while (std::getline(in, line)) {
    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()) || b != e) {
      throw std::runtime_error{"bad input line"};
    }
  }

  return result;
}

/// @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 [i,room] : rooms | boost::adaptors::indexed()) {
    names[room.name] = i;
  }

  distance_array<std::uint64_t> 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 const i : boost::irange(rooms.size())) {
    auto & room = rooms[i];
    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)};
}

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

/// @struct State
/// @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>
{
  // Maximal flow seen at each set of open valves
  std::unordered_map<Valves, std::uint64_t> result;

  // Remaining states for depth first search
  std::vector<State> states { State{initial_time, 0, start, {}} };
  while (!states.empty()) {
    auto const state = states.back();
    states.pop_back();

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

    auto const distances_i = distances[state.location];

    for (auto const [j, room] : rooms | boost::adaptors::indexed()) {
      // 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_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[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();
  
  std::uint64_t best {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

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 rooms = Parse(in);
    auto const n = FlowsFirst(rooms);
    auto const [start, distances] = GenerateDistances(rooms);
    rooms.resize(n);
    CHECK(1651 == Part1(start, rooms, distances));
    CHECK(1707 == Part2(start, rooms, distances));
  }

  TEST_CASE("shortest path") {
    distance_array<int> distances{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);
  }
}

/// @brief Print solutions to parts 1 and 2 
/// @param in selected input stream
auto Main(std::istream & in, std::ostream & out) -> void
{
  auto rooms = Parse(in);
  auto const n = FlowsFirst(rooms);
  auto const [start, distances] = GenerateDistances(rooms);
  rooms.resize(n);
  out << "Part 1: " << Part1(start, rooms, distances) << std::endl;
  out << "Part 2: " << Part2(start, rooms, distances) << std::endl;
}