#include <cstdint>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>
#include <tuple>

#include <doctest.h>

#include <aocpp/Startup.hpp>

using namespace aocpp;

namespace {

auto Parse(std::istream & in)
-> std::vector<std::vector<std::uint16_t>>
{
  std::string line;
  std::vector<std::vector<std::uint16_t>> result;
  while (std::getline(in, line)) {
    line.erase(std::remove_if(line.begin(), line.end(), [](auto c) {
      return c=='<' || c=='-' || c=='>' || c == ',';
    }), line.end());
    std::istringstream lin { line };
    std::uint16_t x;
    lin >> x;
    std::vector<std::uint16_t> links;
    while (lin >> x) {
      links.push_back(x);
    }
    result.emplace_back(std::move(links));
  }
  return result;
}

auto Find(std::vector<std::uint16_t> & leaders, std::uint16_t const who) -> std::uint16_t
{
  std::uint16_t cursor;
  std::uint16_t leader = who;

  do {
    cursor = leader;
    leader = leaders[cursor];
  } while (leader != cursor);

  cursor = who;
  while (cursor != leader) {
    auto tmp = leaders[cursor];
    leaders[cursor] = leader;
    cursor = tmp;
  }

  return leader;
}

auto LinkNodes(std::vector<std::vector<std::uint16_t>> const& input)
-> std::pair<std::size_t, std::size_t>
{
  // Allocate leaders array and assign each node as its own leader
  std::vector<std::uint16_t> leaders(input.size());
  for (std::size_t i = 0; i < leaders.size(); i++) {
    leaders[i] = i;
  }

  // Link up all the nodes as specified in the input file
  for (std::size_t i = 0; i < input.size(); i++) {
    auto x_ = Find(leaders, i);
    for (auto const y : input[i]) {
      leaders[Find(leaders,y)] = x_;
    }
  }

  std::size_t part1 = 0;
  std::size_t part2 = 0;

  auto target1 = Find(leaders, 0);
  for (std::size_t i = 0; i < leaders.size(); i++) {
    if (Find(leaders, i) == target1) part1++;
    if (leaders[i] == i) part2++;
  }

  return {part1, part2};
}

} // namespace

TEST_SUITE("2017-12 examples") {
  TEST_CASE("example") {
    std::istringstream in {
      "0 <-> 2\n"
      "1 <-> 1\n"
      "2 <-> 0, 3, 4\n"
      "3 <-> 2, 4\n"
      "4 <-> 2, 3, 6\n"
      "5 <-> 6\n"
      "6 <-> 4, 5\n"};
    auto [p1,p2] = LinkNodes(Parse(in));
    CHECK(p1 == 6);
    CHECK(p2 == 2);
  }
}

auto main(int argc, char** argv) -> int {
  auto input = Parse(*Startup(argc, argv));
  auto [part1, part2] = LinkNodes(input);
  std::cout << "Part 1: " << part1 << std::endl;
  std::cout << "Part 2: " << part2 << std::endl;
}