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

#include <bitset>

#include <boost/spirit/home/qi.hpp>
#include <boost/range/irange.hpp>
#include <boost/phoenix.hpp>
#include <boost/range/algorithm.hpp>
#include <boost/range/adaptor/map.hpp>

#include <doctest.h>

#include <aocpp/Startup.hpp>

namespace {

auto ShortestDistances(std::vector<std::vector<std::uint64_t>> & dist)
{
  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 RawRoom {
  std::string name;
  std::uint64_t flow;
  std::vector<std::string> connections;
};

struct Room {
  std::uint64_t flow;
  std::vector<std::uint64_t> connections;
};

auto Parse(std::istream & in) -> std::vector<RawRoom>
{
  std::vector<RawRoom> 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;

    auto const name = qi::as_string[+qi::alpha];

    qi::rule<It, RawRoom()> room_description =
      "Valve " >>
      name [ phx::bind(&RawRoom::name, _val) = _1 ] >>
      " has flow rate=" >>
      qi::ulong_long [ phx::bind(&RawRoom::flow, _val) = _1 ] >>
      "; tunnel"  >> -qi::string("s") >>
      " lead"     >> -qi::string("s") >>
      " to valve" >> -qi::string("s") >>
      " " >>
      (name % ", ") [ phx::bind(&RawRoom::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;
}

auto GenerateGraph(std::vector<RawRoom> const& raw_rooms) -> std::pair<std::size_t, std::vector<Room>>
{
  std::pair<std::size_t, std::vector<Room>> result {};
  std::unordered_map<std::string, std::size_t> names;

  std::vector<std::vector<std::uint64_t>> distances(
    raw_rooms.size(), std::vector<std::uint64_t>(raw_rooms.size(), raw_rooms.size()));
  
  for (auto const i : boost::irange(raw_rooms.size())) {
    names[raw_rooms[i].name] = i;
    if (raw_rooms[i].name == "AA") result.first = i;
  }

  for (auto const i : boost::irange(raw_rooms.size())) {
    auto & ds = distances[i];
    ds[i] = 0;
    for (auto const& name : raw_rooms[i].connections) {
      ds[names[name]] = 1;
    }
  }

  ShortestDistances(distances);

  for (auto const i : boost::irange(raw_rooms.size())) {
    result.second.push_back({raw_rooms[i].flow, std::move(distances[i])});
  }

  return result;
}

using Valves = std::bitset<64>;

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

auto Routes(
  std::size_t const start,
  std::vector<Room> const& rooms, std::uint64_t initial_time
) -> 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& room = rooms[state.location];

    for (auto const i : boost::irange(rooms.size())) {
      // don't revisit a valve
      if (state.valves.test(i)) { continue; }

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

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

      auto const time = state.time - cost - 1;      
      auto const flow = state.flow + rooms[i].flow * time;
      auto valves = state.valves;
      valves.set(i);

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

      states.push_back({time, flow, i, valves});
    }
  }

  return result;
}

auto Part1(std::size_t start, std::vector<Room> const& rooms) -> std::uint64_t
{
  auto const routes = Routes(start, rooms, 30);
  return *boost::range::max_element(routes | boost::adaptors::map_values);
}

auto Part2(std::size_t start, std::vector<Room> const& rooms) -> std::uint64_t
{
  auto const routes = Routes(start, rooms, 26);
  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 [start, graph] = GenerateGraph(Parse(in));
    CHECK(1651 == Part1(start, graph));
    CHECK(1707 == Part2(start, graph));
  }
}

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