#include <array>
#include <bit>
#include <cstdint>
#include <iostream>
#include <optional>
#include <string>
#include <vector>

#include <doctest.h>

#include <aocpp/Startup.hpp>
#include <aocpp/Parsing.hpp>
#include <knothash.hpp>

namespace {

using Grid = std::vector<std::array<std::uint8_t, 16>>;

auto MakeRows(std::string input) -> Grid {
  Grid rows;
  rows.reserve(128);

  input += "-";
  for (int i = 0; i < 128; i++) {
    rows.emplace_back(knothash::hash(input + std::to_string(i)));
  }

  return rows;
}

auto CountBits(Grid const& grid) -> std::size_t {
  std::size_t bits = 0;
  for (auto const& row : grid) {
    for (auto const x : row) {
      bits += std::popcount(x);
    }
  }
  return bits;
}

auto Connect(std::vector<std::uint16_t> & leaders, std::uint16_t who, std::uint16_t leader) {
  auto who_ = leaders[who];
  if (who_ == std::uint16_t(-1)) { return; } // there wasn't a node here, abort
  for (;;) {
    leaders[who] = leader;
    if (who_ == who) { return; }
    who = who_;
    who_ = leaders[who];
  }
}

auto CountComponents(Grid const& grid) -> std::uint16_t {

  std::vector<std::uint16_t> leaders(128 * 128);

  // Visit all pairs of neighbors linking them if both are set
  for (std::uint16_t r = 0; r < 128; r++) {
    bool prev = false;
    for (std::uint16_t c = 0; c < 128; c++) {
      std::uint16_t me = r*128+c;

      // first time we've visited this node, if it's in the graph it's a leader!
      if (grid[r][c/8] & (1<<(7-(c%8)))) {
        leaders[me] = me;

        // If previous cell in row is a node, link it to me immediately
        // It will have just been set as its own group leader.
        if (prev) { leaders[me-1] = me; }

        // Connect up's leader _after_ updating left's leader in case they were connected
        if (0 < r) { Connect(leaders, me - 128, me); }
        prev = true;
      } else {
        prev = false;
        leaders[me] = -1;
      }
    }
  }

  // count component leaders
  std::uint16_t n = 0;
  for (std::size_t i = 0; i < leaders.size(); i++) {
    if (leaders[i] == i) { n++; }
  }

  return n;
}

} // namespace

TEST_SUITE("2017-14 examples") {
TEST_CASE("flqrgnkx") {
  auto rows = MakeRows("flqrgnkx");
  CHECK(CountBits(rows) == 8108);
  CHECK(CountComponents(rows) == 1242);
}
}

auto main(int argc, char** argv) -> int {
  std::string input;
  std::getline(*aocpp::Startup(argc, argv), input);
  auto rows = MakeRows(std::move(input));
  std::cout << "Part 1: " << CountBits(rows) << std::endl;
  std::cout << "Part 2: " << CountComponents(rows) << std::endl;
}