#include <cstddef>
#include <string>
#include <iostream>
#include <variant>
#include <vector>

#include <gmpxx.h>

#include <aocpp/Startup.hpp>
using namespace aocpp;

template <class... Ts> struct overloaded : Ts... { using Ts::operator()...; };
template <class... Ts> overloaded(Ts...) -> overloaded<Ts...>;

namespace {

template <unsigned long Mod>
class ZMod {
  mpz_class value;

public:
  ZMod() : value {} {}
  ZMod(long value) : ZMod{mpz_class{value}} {}
  ZMod(mpz_class value) {
    mpz_mod_ui(this->value.get_mpz_t(), value.get_mpz_t(), Mod);
  }

  auto operator+(ZMod const& rhs) const -> ZMod {
    return {value + rhs.value};
  }

  auto operator-() const -> ZMod {
    return {-value};
  }

  auto operator-(ZMod const& rhs) const -> ZMod {
    return {value - rhs.value};
  }

  auto operator*(ZMod const& rhs) const -> ZMod {
    return {value * rhs.value};
  }

  auto inverse() const -> ZMod {
    mpz_class m{Mod};
    ZMod result;
    mpz_invert(result.value.get_mpz_t(), value.get_mpz_t(), m.get_mpz_t());
    return result;
  }

  auto friend operator<<(std::ostream & out, ZMod<Mod> const& x) -> std::ostream & {
  return out << x.value.get_ui();
}
};

template <typename Rep>
struct LinearFunction {
  Rep scale;
  Rep offset;

  auto operator()(Rep x) const -> Rep {
    return scale * x + offset;
  }

  auto inverse() const -> LinearFunction {
    // y = a*x+b
    // y-b = a*x
    // a' * (y-b) = x
    auto scale_ = scale.inverse();
    return {scale_, - scale_ * offset};
  }
};

template <typename Rep>
auto Compose(
  LinearFunction<Rep> const& f,
  LinearFunction<Rep> const& g
) -> LinearFunction<Rep>
{
  return {f.scale * g.scale, f.offset + f.scale * g.offset};
}

template <unsigned long Cards>
using Shuffle = LinearFunction<ZMod<Cards>>;

template<unsigned long Cards>
auto Cut(long cards) -> Shuffle<Cards> {
  return Shuffle<Cards>{1, -cards};
}

template<unsigned long Cards>
auto Rev() -> Shuffle<Cards> {
  return Shuffle<Cards>{-1, -1};
}

template<unsigned long Cards>
auto Deal(long inc) -> Shuffle<Cards> {
  return Shuffle<Cards>{inc, 0};
}

struct DoCut { long n; };
struct DoRev {};
struct DoDeal { long n; };
using Instruction = std::variant<DoCut, DoRev, DoDeal>;

auto Parse(std::istream & in) -> std::vector<Instruction> {
  std::vector<Instruction> result;
  std::string line;
  while (std::getline(in, line)) {
    if (line.starts_with("cut")) {
      auto n = std::stol(line.substr(4));
      result.push_back(DoCut{n});
    } else if (line == "deal into new stack") {
      result.push_back(DoRev{});
    } else if (line.starts_with("deal with increment")) {
      auto n = std::stol(line.substr(20));
      result.push_back(DoDeal{n});
    } else {
      throw std::runtime_error{"bad input command"};
    }
  }
  return result;
}

template <std::int64_t Cards>
auto FollowInstructions(std::vector<Instruction> const& instructions) -> Shuffle<Cards> {
  Shuffle<Cards> shuffle = {1,0};
  for (auto const& instruction : instructions) {
    auto step = std::visit(overloaded {
      [](DoRev) { return Rev<Cards>(); },
      [](DoCut c) { return Cut<Cards>(c.n); },
      [](DoDeal c) { return Deal<Cards>(c.n); }
    }, instruction);
    shuffle = Compose(step, shuffle);
  }
  return shuffle;
}

template<unsigned long Cards>
auto Many(Shuffle<Cards> shuffle, unsigned long n) -> Shuffle<Cards> {
  if (n == 0) {
    return Shuffle<Cards>{1,0};
  } else if (n == 1) {
    return shuffle;
  } else if (n & 1) {
    auto shuffle2 = Many(shuffle, n/2);
    return Compose(shuffle, Compose(shuffle2, shuffle2));
  } else {
    auto shuffle2 = Many(shuffle, n/2);
    return Compose(shuffle2, shuffle2);
  }
}

} // namespace

auto main(int argc, char** argv) -> int {
  auto instructions = Parse(Startup(argc, argv));

  auto shuffle1 = FollowInstructions<10007>(instructions);
  std::cout << "Part 1: " << shuffle1(2019) << std::endl;

  auto shuffle2 = FollowInstructions<119315717514047>(instructions);
  shuffle2 = Many(shuffle2, 101741582076661);
  std::cout << "Part 2: " << shuffle2.inverse()(2020) << std::endl;
}