#include "allocator.h"
#include "aoc.h"
#include "dvec_s.h"
#include "util.h"

#include <assert.h>
#include <stddef.h>
#include <stdlib.h>

size_t
imul(size_t a, size_t b, size_t n)
{
	size_t i, out;

	out = 0;
	for (i = 0; i < sizeof(size_t) * 8; i++, b <<= 1) {
		out = (2 * out) % n;
		if (b & (1ULL << 63))
			out = (out + a) % n;
	}

	return out;
}

size_t
ipow(size_t a, size_t b, size_t n)
{
	size_t i, out;

	out = 1;
	for (i = 0; i < sizeof(size_t) * 8; i++, b <<= 1) {
		out = imul(out, out, n);
		if (b & (1ULL << 63))
			out = imul(out, a, n);
	}

	return out;
}

size_t
modinv(size_t a, size_t n)
{
	return ipow(a, n - 2, n);
}

void
part2(void)
{
	char line[256];
	const char *lpos, *lend;
	size_t answer;
	size_t count;
	size_t times;
	size_t target;
	ssize_t sval;
	size_t mul, off;
	size_t _mul, _off;
	size_t val;

	times = 101741582076661;
	count = 119315717514047;
	target = 2020;

	/* reduce operations to a single multiplication and
	 * addition under modulus 'count' */
	mul = 1;
	off = 0;
	lpos = aoc.input;
	lend = lpos + aoc.input_size;
	while (readtok(line, sizeof(line), '\n', &lpos, lend)) {
		if (sscanf(line, "deal with increment %lu", &val) == 1) {
			/* multiplication : pos = val * pos */
			_mul = val;
			_off = 0;
		} else if (sscanf(line, "cut %li", &sval) == 1) {
			/* addition : pos = pos + val */
			_mul = 1;
			_off = ((size_t) ((ssize_t) count - sval)) % count;
		} else if (!strcmp(line, "deal into new stack")) {
			/* invert : pos = - pos + n - 1 */
			_mul = count - 1;
			_off = count - 1;
		} else if (strspn(line, " \n\r\t\v") == strlen(line)) {
			continue;
		} else {
			assert(false);
		}
		mul = imul(mul, _mul, count);
		off = (imul(off, _mul, count) + _off) % count;
	}

	/* apply multiplication and addition 'times' times.. */
	val = (ipow(mul, times, count) + count - 1) % count;
	val = imul(val, modinv(mul + count - 1, count), count);
	off = imul(off, val, count); /* b' = b * ((a^n - 1) / (a - 1)) */
	mul = ipow(mul, times, count); /* a' = a^n */

	answer = imul((target + count - off) % count, modinv(mul, count), count);
	aoc.answer = aprintf("%lu", answer);
	aoc.solution = "73394009116480";
}