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

#include <assert.h>
#include <string.h>
#include <stdio.h>

static const int freq[] = { 0, 1, 0, -1 };

void
parse_input(struct dvec *phases)
{
	const char *c, *sep;
	int *val;

	sep = strchr(aoc.input, '\n');
	assert(sep != NULL);

	for (c = aoc.input; c < sep; c++) {
		val = dvec_add_slot(phases);
		*val = *c - '0';
	}
}

void
fft(struct dvec *phases, size_t rounds)
{
	struct dvec copy;
	size_t r, i, k, l;
	ssize_t sum;
	int digit;
	int *val;

	dvec_init(&copy, sizeof(int), 0, &stdlib_strict_heap_allocator);

	for (r = 0; r < rounds; r++) {
		dvec_copy(&copy, phases);

		for (i = 0; i < copy.len; i++) {
			sum = 0;
			for (k = 0; k < copy.len; k++) {
				l = (k + 1) / (i + 1);
				digit = *(int *)dvec_at(&copy, k);
				sum += freq[l % 4] * digit;
			}
			sum = ABS(sum) % 10;
			val = dvec_at(phases, i);
			*val = (int) sum;
		}
	}

	dvec_deinit(&copy);
}

void
fft_dup(struct dvec *phases, size_t offset, size_t rounds, size_t dup)
{
	size_t r, i, len, lutlen;
	int *lut, *lutc, *lutp;
	int digit, *slot;

	/*   |    a           b    c
	 * --+------------------------
	 * 1 |  a+b+c        b+c   c
	 * 2 | a+2b+3c       b+2c  c
	 * n |a+nb+n(n+1)/2c b+nc  c
	 *
	 * a+sum(1)b+sum(sum(1))c .. not nice
	 *
	 * sum(1): n
	 * sum(sum(1)): n(n+1)/2
	 * sum(sum(sum(1))): (K2+K1)/2 = (n^2+n)/2 + (n^3-n)/6
	 *
	 * .. calculation possible but not practically feasible since we
	 * would have to do 10000..th sum() .. will have to use a lookup table
	 */

	len = phases->len * dup;
	assert(offset >= len / 2 && offset < len);

	lutlen = len - offset;
	lut = malloc(lutlen * sizeof(int) * (rounds + 1));
	assert(lut != NULL);

	/* top row */
	for (i = 0; i < lutlen; i++) {
		digit = *(int *)dvec_at(phases, (i + offset) % phases->len);
		lut[i] = digit;
	}

	/* right column */
	digit = *(int *)dvec_at(phases, (len - 1) % phases->len);
	for (r = 0; r < rounds; r++) {
		lut[(r + 1) * lutlen + lutlen - 1] = digit;
	}

	for (r = 0; r < rounds; r++) {
		lutp = &lut[r * lutlen + lutlen - 2];
		lutc = &lut[(r + 1) * lutlen + lutlen - 2];
		for (i = 1; i < lutlen; i++, lutp--, lutc--)
			lutc[0] = (lutp[0] + lutc[1]) % 10;
	}

	dvec_clear(phases);
	for (i = 0; i < 8; i++) {
		slot = dvec_add_slot(phases);
		*slot = lut[rounds * lutlen + i];
	}

	free(lut);
}

void
part1(void)
{
	struct dvec phases;
	char *answer;
	size_t i;
	int digit;

	dvec_init(&phases, sizeof(int), 64, &stdlib_strict_heap_allocator);

	parse_input(&phases);

	fft(&phases, 100);

	answer = malloc(9);
	for (i = 0; i < 8; i++) {
		digit = *(int *)dvec_at(&phases, i);
		answer[i] = '0' + (char) digit;
	}
	answer[i] = 0;

	aoc.answer = answer;
	aoc.solution = "29956495";

	dvec_deinit(&phases);
}

void
part2(void)
{
	struct dvec phases;
	char *answer;
	int digit;
	size_t offset;
	size_t base;
	size_t i;

	dvec_init(&phases, sizeof(int), 64, &stdlib_strict_heap_allocator);

	parse_input(&phases);

	base = 1;
	offset = 0;
	for (i = 0; i < 7; i++) {
		digit = *(int *)dvec_at(&phases, 6 - i);
		offset += base * (size_t) digit;
		base *= 10;
	}

	fft_dup(&phases, offset, 100, 10000);

	answer = malloc(9);
	for (i = 0; i < 8; i++) {
		digit = *(int *)dvec_at(&phases, i);
		answer[i] = '0' + (char) digit;
	}
	answer[i] = 0;

	aoc.answer = answer;
	aoc.solution = "73556504";

	dvec_deinit(&phases);
}