from tools.aoc import AOCDay
from tools.coordinate import Coordinate, DistanceAlgorithm
from typing import Any


ZERO = Coordinate(0, 0, 0)


class Particle:
    def __init__(self, pos: Coordinate, vel: Coordinate, acc: Coordinate):
        self.pos = pos
        self.vel = vel
        self.acc = acc

    def move(self):
        self.vel += self.acc
        self.pos += self.vel

    def dist(self):
        return self.pos.getDistanceTo(ZERO, algorithm=DistanceAlgorithm.MANHATTAN, includeDiagonals=False)


class Day(AOCDay):
    inputs = [
        [
            (0, "input20_test"),
            (457, "input20")
        ],
        [
            (1, "input20_test2"),
            (448, "input20")
        ]
    ]

    def getParticles(self) -> list:
        particles = []
        for line in self.getInput():
            p, v, a = line.split(", ")
            p = Coordinate(*map(int, p[3:-1].split(",")))
            v = Coordinate(*map(int, v[3:-1].split(",")))
            a = Coordinate(*map(int, a[3:-1].split(",")))

            particles.append(Particle(p, v, a))

        return particles

    def part1(self) -> Any:
        particles = self.getParticles()
        closest = 0
        last_closest = list(range(-1, -1000, -1))
        while len(set(last_closest)) != 1:
            last_closest.append(closest)
            last_closest.pop(0)
            min_dist = 10e9
            for i, p in enumerate(particles):
                p.move()
                if p.dist() < min_dist:
                    closest = i
                    min_dist = p.dist()

        return closest

    def part2(self) -> Any:
        particles = self.getParticles()
        for _ in range(1000):
            pos_cache = {}
            new_particles = []
            for i, p in enumerate(particles.copy()):
                p.move()
                if p.pos not in pos_cache:
                    new_particles.append(p)
                    pos_cache[p.pos] = len(new_particles) - 1
                else:
                    if pos_cache[p.pos] >= 0:
                        new_particles.pop(pos_cache[p.pos])
                    pos_cache[p.pos] = -1

            particles = new_particles

        return len(particles)


if __name__ == '__main__':
    day = Day(2017, 20)
    day.run(verbose=True)