from math import lcm
from re import compile, findall
from tools.aoc import AOCDay
from tools.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 getKineticEnergy(self):
        return abs(self.vel_x) + abs(self.vel_y) + abs(self.vel_z)

    def getPotentialEnergy(self):
        return abs(self.x) + abs(self.y) + abs(self.z)

    def getEnergy(self):
        return self.getPotentialEnergy() * self.getKineticEnergy()

    def __str__(self):
        return "x=%3d y=%3d z=%3d" % (self.x, self.y, self.z)

    def __repr__(self):
        return str(self)


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):
    inputs = [
        [
            (183, "test_input12"),
            (14645, "test_input12_2"),
            (12644, "input12")
        ],
        [
            (2772, "test_input12"),
            (4686774924, "test_input12_2"),
            (None, "input12")
        ]
    ]

    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 = [Moon(*map(int, findall(scanline_regex, line)[0])) for line in self.getInput()]
        vel_0_pos = [0] * len(moons)
        moveCounter = 0
        while 0 in vel_0_pos:
            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:
                    # if i == 1:
                    #     print("init pos:", i, "after", moveCounter, "moves =>", m.getPotentialEnergy(), m.getKineticEnergy())
                    if m.getKineticEnergy() == 0:
                        print("init state:", i, "after", moveCounter, "moves =>", m.getPotentialEnergy())
                        vel_0_pos[i] = moveCounter

        return lcm(*vel_0_pos)


if __name__ == '__main__':
    day = Day(2019, 12)
    day.run(verbose=True)