import re

import collections

from solutions import BaseSolution


class Solution(BaseSolution):
    input_file = "20.txt"

    def __str__(self):
        return "Day 20: Particle Swarm"

    def _get_particle(self, line):
        return [
            list(map(int, coordinate.split(",")))
            for coordinate in re.findall(r".=<([^>]+)>", line)
        ]

    def solve(self, puzzle_input, ticks=10**4):
        particles = [self._get_particle(line) for line in puzzle_input.splitlines()]
        distances = [[] for _ in range(len(particles))]
        for _ in range(ticks):
            for i, particle in enumerate(particles):
                p, v, a = particle
                v[0] += a[0]
                v[1] += a[1]
                v[2] += a[2]
                p[0] += v[0]
                p[1] += v[1]
                p[2] += v[2]
                distances[i].append(sum(map(abs, p)))
        d = sorted(
            map(lambda d: (d[0], sum(d[1]) / len(d[1])), enumerate(distances)),
            key=lambda x: x[1],
        )
        return d[0][0]

    def solve_again(self, puzzle_input, ticks=10**3):
        particles = [self._get_particle(line) for line in puzzle_input.splitlines()]
        for _ in range(ticks):
            positions = collections.defaultdict(list)
            for particle in particles:
                p, v, a = particle
                v[0] += a[0]
                v[1] += a[1]
                v[2] += a[2]
                p[0] += v[0]
                p[1] += v[1]
                p[2] += v[2]
                k = "-".join(map(str, p))
                positions[k].append(particle)
            for duplicates in positions.values():
                if len(duplicates) > 1:
                    for d in duplicates:
                        particles.remove(d)
        return len(particles)


if __name__ == "__main__":
    solution = Solution()
    solution.show_results()