#include <cstdint>
#include <iostream>
#include <numeric>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>

#include <boost/phoenix.hpp>
#include <boost/range/adaptor/indexed.hpp>
#include <boost/spirit/include/qi.hpp>

#include <doctest.h>

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

namespace {

namespace qi = boost::spirit::qi;
namespace phx = boost::phoenix;

struct Monkey {
  std::vector<std::int64_t> items;
  std::int64_t divisor;
  std::size_t true_ix, false_ix;
  std::optional<std::int64_t> rhs;
  char op;

  auto throw_to(std::int64_t const item) const -> std::size_t {
    return item % divisor == 0 ? true_ix : false_ix;
  }

  auto operation(std::int64_t const item) const -> std::int64_t {
    auto const y = rhs.value_or(item);
    switch (op) {
      case '+': return item + y;
      case '*': return item * y;
      default: throw std::runtime_error{"bad input"};
    }
    return item;
  }
};

struct Grammar : public qi::grammar<std::string::const_iterator, std::vector<Monkey>()> {
  qi::rule<iterator_type, Monkey()> monkey;
  qi::rule<iterator_type, std::vector<Monkey>()> monkeys;

  Grammar() : base_type{monkeys} {
    using namespace qi::labels;
    monkeys = monkey % "\n";
    monkey =
      "Monkey " >>
      qi::omit[qi::ulong_long] >>
      ":\n  Starting items: " >>
      (qi::long_long % ", ") [ phx::bind(&Monkey::items, _val) = _1 ] >>
      "\n  Operation: new = old " >>
      qi::char_ [ phx::bind(&Monkey::op, _val) = _1 ] >>
      " " >>
      ("old" | qi::long_long [ phx::bind(&Monkey::rhs, _val) = _1 ]) >>
      "\n  Test: divisible by " >>
      qi::long_long [ phx::bind(&Monkey::divisor, _val) = _1 ] >>
      "\n    If true: throw to monkey " >>
      qi::ulong_long [ phx::bind(&Monkey::true_ix, _val) = _1 ] >>
      "\n    If false: throw to monkey " >>
      qi::ulong_long [ phx::bind(&Monkey::false_ix, _val) = _1 ] >>
      "\n";
  }
};

auto Solve(std::vector<Monkey> const& input, std::size_t const rounds, std::optional<std::int64_t> const modulus) -> std::int64_t
{
  auto throws = std::vector<std::int64_t>(input.size());
  for (auto const& [start, monkey] : boost::adaptors::index(input)) {
    for (auto item : monkey.items) {
      auto cursor = static_cast<std::size_t>(start);
      for (auto n = rounds; n > 0;) {
        throws[cursor]++;
        auto const& m = input[cursor];

        item = m.operation(item);

        if (modulus) {
          item %= *modulus;
        } else {
          item /= 3;
        }

        auto const next = m.throw_to(item);
        if (next < cursor) n--;
        cursor = next;
      }
    }
  }

  // Compute top 2 elements of the vector
  std::partial_sort(throws.begin(), throws.begin()+2, throws.end(), std::greater());

  return throws[0] * throws[1];
}

auto Part1(std::vector<Monkey> const& input) -> std::int64_t
{
  return Solve(input, 20, {});
}

auto Part2(std::vector<Monkey> const& input) -> std::int64_t
{
  auto const modulus =
    std::transform_reduce(
      input.begin(), input.end(),
      std::int64_t{1}, std::lcm<std::int64_t, std::int64_t>,
      [](auto const& m) { return m.divisor; });
  return Solve(input, 10'000, modulus);
}

} // namespace

TEST_SUITE("2022-11 examples") {
  TEST_CASE("example") {
    std::istringstream in {
R"(Monkey 0:
  Starting items: 79, 98
  Operation: new = old * 19
  Test: divisible by 23
    If true: throw to monkey 2
    If false: throw to monkey 3

Monkey 1:
  Starting items: 54, 65, 75, 74
  Operation: new = old + 6
  Test: divisible by 19
    If true: throw to monkey 2
    If false: throw to monkey 0

Monkey 2:
  Starting items: 79, 60, 97
  Operation: new = old * old
  Test: divisible by 13
    If true: throw to monkey 1
    If false: throw to monkey 3

Monkey 3:
  Starting items: 74
  Operation: new = old + 3
  Test: divisible by 17
    If true: throw to monkey 0
    If false: throw to monkey 1
)"
    };
    auto const input = aocpp::ParseGrammar(Grammar{}, in);
    CHECK(Part1(input) ==        10'605);
    CHECK(Part2(input) == 2'713'310'158);
  }
}

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