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

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

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

namespace {

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;
}