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

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

struct vec3 {
	union {
		struct {
			int x, y, z;
		};
		int axis[3];
	};
};

struct moon {
	struct vec3 pos, vel;
};

void
moons_init(struct dvec *vec)
{
	char *p, *tok, *val;
	struct moon m, *slot;
	int i;

	m.vel = (struct vec3) { 0, 0, 0 };
	dvec_init(vec, 100, sizeof(struct moon), &stdlib_strict_heap_allocator);

	tok = aoc.input;
	while ((p = strchr(tok, '\n'))) {
		val = tok;
		for (i = 0; i < 3 && (val = strchr(val + 1, '=')); i++)
			m.pos.axis[i] = (int) strtol(val + 1, NULL, 10);
		slot = dvec_add_slot(vec);
		memcpy(slot, &m, sizeof(struct moon));
		tok = p + 1;
	}
}

void
moons_deinit(struct dvec *vec)
{
	dvec_deinit(vec);
}

int
get_energy(struct moon *m)
{
	int pot, kin;

	pot = ABS(m->pos.x);
	pot += ABS(m->pos.y);
	pot += ABS(m->pos.z);

	kin = ABS(m->vel.x);
	kin += ABS(m->vel.y);
	kin += ABS(m->vel.z);

	return pot * kin;
}

int
rel(int a, int b)
{
	if (a == b) return 0;
	return (a < b) ? +1 : -1;
}

void
sim_step(struct dvec *moons, int show)
{
	struct moon *m1, *m2;
	static int step = 0;
	size_t i, k, l;
	int relv;

	/* initial state */
	for (i = 0; !step && show && i < dvec_len(moons); i++) {
		m1 = dvec_at(moons, i);
		aoc_debug("%i: %i %i %i , %i %i %i = %i\n", step, m1->pos.x, m1->pos.y, m1->pos.z,
			m1->vel.x, m1->vel.y, m1->vel.z, get_energy(m1));
	}

	/* update velocity */
	for (i = 0; i < dvec_len(moons); i++) {
		for (k = i + 1; k < dvec_len(moons); k++) {
			m1 = dvec_at(moons, i);
			m2 = dvec_at(moons, k);

			for (l = 0; l < 3; l++) {
				relv = rel(m1->pos.axis[l], m2->pos.axis[l]);
				m1->vel.axis[l] += relv;
				m2->vel.axis[l] -= relv;
			}
		}
	}

	/* update position */
	for (i = 0; i < dvec_len(moons); i++) {
		m1 = dvec_at(moons, i);
		for (k = 0; k < 3; k++)
			m1->pos.axis[k] += m1->vel.axis[k];

		if (show) aoc_debug("%i: %i %i %i , %i %i %i = %i\n", step+1,
			m1->pos.x, m1->pos.y, m1->pos.z, m1->vel.x,
			m1->vel.y, m1->vel.z, get_energy(m1));
	}

	step++;
}

void
part1(void)
{
	struct dvec moons;
	size_t i;
	int energy;

	moons_init(&moons);

	for (i = 0; i < 1000; i++)
		sim_step(&moons, aoc.debug);

	energy = 0;
	for (i = 0; i < dvec_len(&moons); i++)
		energy += get_energy(dvec_at(&moons, i));

	aoc.answer = aprintf("%i", energy);
	aoc.solution = "12053";

	moons_deinit(&moons);
}

void
part2(void)
{
	struct dvec init, moons;
	struct maxint cycles[3];
	struct maxint imm = MI_INIT;
	mi_ul tmp;
	struct maxint *steps;
	struct moon *m1, *m2;
	int same, poseq, veleq;
	char buf[64];
	size_t i, k;

	moons_init(&init);

	for (i = 0; i < 3; i++)
		mi_init(&cycles[i]);

	for (i = 0; i < 3; i++) {
		moons_init(&moons);

		same = 0;
		cycles[i] = MI_INIT;
		mi_setv(&cycles[i], 0, MI_POS);

		while (!same) {
			sim_step(&moons, aoc.debug == 2);

			same = 1;
			for (k = 0; k < dvec_len(&moons); k++) {
				m1 = dvec_at(&moons, k);
				m2 = dvec_at(&init, k);
				veleq = m1->pos.axis[i] == m2->pos.axis[i];
				poseq = m1->vel.axis[i] == m2->vel.axis[i];
				if (!veleq || !poseq) {
					same = 0;
					break;
				}
			}

			mi_add(&cycles[i], &cycles[i], mi_imm(&imm, &tmp, 1, MI_POS));
		}

		mi_dec(&cycles[i], buf, sizeof(buf), 0);
		aoc_debug("Cycles Axis %i: %s\n", i+1, buf);

		moons_deinit(&moons);
	}

	steps = &cycles[0];
	for (i = 1; i < 3; i++)
		mi_lcm(steps, steps, &cycles[i]);

	mi_dec(steps, buf, sizeof(buf), 0);
	aoc.answer = strdup(buf);
	aoc.solution = "320380285873116";

	for (i = 0; i < 3; i++)
		mi_deinit(&cycles[i]);

	moons_deinit(&init);
}