from aoc import AOCDay
from math import lcm
from re import compile, findall
from tools import compare
from typing import Any, List

scanline_regex = compile(r"^<x=(.*), y=(.*), z=(.*)>$")


class Moon:
    x: int
    y: int
    z: int
    vel_x: int = 0
    vel_y: int = 0
    vel_z: int = 0

    def __init__(self, x: int, y: int, z: int):
        self.x = x
        self.y = y
        self.z = z

    def getEnergy(self):
        potential = abs(self.x) + abs(self.y) + abs(self.z)
        kinetic = abs(self.vel_x) + abs(self.vel_y) + abs(self.vel_z)
        return potential * kinetic


def moveMoons(moons: List[Moon]):
    for i, moon in enumerate(moons):
        for t in range(len(moons)):
            if i == t:
                continue

            moon.vel_x += compare(moons[t].x, moon.x)
            moon.vel_y += compare(moons[t].y, moon.y)
            moon.vel_z += compare(moons[t].z, moon.z)

    for moon in moons:
        moon.x += moon.vel_x
        moon.y += moon.vel_y
        moon.z += moon.vel_z


class Day(AOCDay):
    test_solutions_p1 = [183, 14645]
    test_solutions_p2 = [2772, 4686774924]

    def part1(self) -> Any:
        moons = [Moon(*map(int, findall(scanline_regex, line)[0])) for line in self.getInput()]

        for step in range(1000):
            moveMoons(moons)

        return sum([x.getEnergy() for x in moons])

    def part2(self) -> Any:
        moons = [Moon(*map(int, findall(scanline_regex, line)[0])) for line in self.getInput()]

        init_moons = moons.copy()
        vel_0_pos = [0, 0, 0, 0]
        moveCounter = 0
        #while 0 in vel_0_pos:
        for bla in range(3000):
            moveCounter += 1
            moveMoons(moons)
            for i, m in enumerate(moons):
                if m.x == init_moons[i].x and m.y == init_moons[i].y and m.z == init_moons[i].z and m.vel_x == 0 and m.vel_y == 0 and m.vel_z == 0:
                    print(i, moveCounter)
                if (vel_0_pos[i] == 0 or vel_0_pos[i] > moveCounter - vel_0_pos[i]) and m.vel_x == 0 and m.vel_y == 0 and m.vel_z == 0:
                    vel_0_pos[i] = moveCounter - vel_0_pos[i]

        print(vel_0_pos)
        return lcm(*vel_0_pos)