#include <algorithm>
#include <cstdint>
#include <string>
#include <iostream>
#include <tuple>
#include <array>
#include <map>

#include <doctest.h>

#include <aocpp/Startup.hpp>

namespace {

template <std::size_t N>
using Banks = std::array<std::int64_t, N>;

auto Parse(std::istream & in) -> Banks<16> {
  Banks<16> result;
  for (auto & x : result) {
    if (!(in >> x)) {
      throw std::runtime_error{"bad input"};
    }
  }

  return result;
}

template <std::size_t N>
auto Redistribute(Banks<N> & banks) {
  auto it = std::max_element(banks.begin(), banks.end());
  auto val = *it;
  *it++ = 0;
  for (; val > 0; val--) {
    if (it == banks.end()) it = banks.begin();
    (*it++)++;
  }
}

template <std::size_t N>
auto CycleSearch(Banks<N> & banks)
-> std::pair<std::size_t, std::size_t>
{
  std::map<Banks<N>, std::size_t> seen {{banks, 0}};
  std::size_t n = 0;
  for(;;) {
    Redistribute(banks);
    n++;
    auto [it, success] = seen.try_emplace(banks, n);
    if (!success) {
      return {n, n - it->second};
    }
  }
}

} // namespace

TEST_SUITE("2017-06 examples") {
  TEST_CASE("example") {
    Banks<4> banks {0, 2, 7, 0};

    SUBCASE("single step") {
      Redistribute(banks);
      REQUIRE(banks == Banks<4>{2, 4, 1, 2});
      Redistribute(banks);
      REQUIRE(banks == Banks<4>{3, 1, 2, 3});
      Redistribute(banks);
      REQUIRE(banks == Banks<4>{0, 2, 3, 4});
      Redistribute(banks);
      REQUIRE(banks == Banks<4>{1, 3, 4, 1});
      Redistribute(banks);
      REQUIRE(banks == Banks<4>{2, 4, 1, 2});
    }

    SUBCASE("full example") {
      auto [p1,p2] = CycleSearch(banks);
      CHECK(p1 == 5);
      CHECK(p2 == 4);
    }
  }
}

auto main(int argc, char** argv) -> int {
  auto input = Parse(*aocpp::Startup(argc, argv));
  auto answer = CycleSearch(input);
  std::cout << "Part 1: " << answer.first << std::endl;
  std::cout << "Part 2: " << answer.second << std::endl;
}