#include <algorithm>
#include <cstdint>
#include <iostream>
#include <iterator>
#include <numeric>
#include <sstream>
#include <stdexcept>
#include <string>
#include <tuple>
#include <vector>

#include <boost/range/irange.hpp>

#include <doctest.h>

#include <aocpp/DivMod.hpp>
#include <aocpp/Startup.hpp>

namespace {

using Input = std::vector<std::int64_t>;

/// @brief Parse the input stream as a newline-delimited list of lists of integers
/// @param in input file parsed until EOF
/// @return outer list of elves, inner lists of calories
auto Parse(std::istream & in) -> Input
{
  auto result = Input{};
  using It = std::istream_iterator<std::int64_t>;
  std::copy<It>(in, {}, std::back_inserter(result));
  return result;
}

auto Solve(Input const& input, std::size_t const rounds) -> std::int64_t
{
  auto const len = input.size();
  auto fwds = std::vector<std::size_t>(len);
  auto bwds = std::vector<std::size_t>(len);

  if (len == 0) {
    throw std::runtime_error{"Input too small"};
  }

  auto const zero_it = std::find(input.begin(), input.end(), 0);
  if (zero_it == input.end()) {
    throw std::runtime_error{"no zero element in input"};
  }

  std::iota(fwds.begin(), fwds.end()-1, 1); fwds.back() = 0;
  std::iota(bwds.begin()+1, bwds.end(), 0); bwds.front() = len-1; 

  for (auto const _ : boost::irange(rounds)) {
    for (auto const i : boost::irange(len)) {
      
      // Remove i from the ring
      auto next = fwds[i];
      auto prev = bwds[i];
      fwds[prev] = next;
      bwds[next] = prev;

      auto const steps = aocpp::Mod<std::int64_t>(input[i], len - 1);
      if (steps < len/2) {
        for (auto const _ : boost::irange(steps)) {
          next = fwds[next];
        }
      } else {
        for (auto const _ : boost::irange(len - 1 - steps)) {
          next = bwds[next];
        }
      }
      prev = bwds[next];

      // insert i into the ring
      fwds[i] = next;
      bwds[i] = prev;
      fwds[prev] = i;
      bwds[next] = i;
    }
  }

  auto sum = std::int64_t{0};
  auto cursor = std::distance(input.begin(), zero_it);
  for (auto const _ : boost::irange(3)) {
    for (auto const _ : boost::irange(1000)) {
      cursor = fwds[cursor];
    }
    sum += input[cursor];
  }

  return sum;
}

auto Part1(Input const& input) -> std::int64_t
{
  return Solve(input, 1);
}

auto Part2(Input input) -> std::int64_t
{
  for (auto& x : input) { x *= 811'589'153; }
  return Solve(input, 10);
}

} // namespace

TEST_SUITE("2022-20 examples") {
  TEST_CASE("example") {
    std::istringstream in {
      "1\n"
      "2\n"
      "-3\n"
      "3\n"
      "-2\n"
      "0\n"
      "4\n"};
      auto input = Parse(in);
      CHECK(3 == Part1(input));
      CHECK(1623178306 == Part2(std::move(input)));
  }
}

auto Main(std::istream & in, std::ostream & out) -> void
{
  auto input = Parse(in);
  out << "Part 1: " << Part1(input) << std::endl;
  out << "Part 2: " << Part2(std::move(input)) << std::endl;
}