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


class Day(AOCDay):
    inputs = [
        [
            (40, "input8_test"),
            (140008, "input8"),
        ],
        [
            (25272, "input8_test"),
            (9253260633, "input8"),
        ]
    ]

    def parse_input(self) -> tuple[dict[float, Coordinate], dict[Coordinate, int]]:
        boxes = [Coordinate(*map(int, line.split(","))) for line in self.getInput()]
        dists = {}
        for i, box in enumerate(boxes):
            for box2 in boxes[i + 1 :]:
                dists[box.getDistanceTo(box2)] = (box, box2)

        members = {box: i for i, box in enumerate(boxes)}
        return dists, members

    def solve(self, part2: bool = False) -> int:
        dists, members = self.parse_input()
        circuits = {i: [box] for box, i in members.items()}
        connections_needed = 10 if self.is_test() else 1000
        connections_done = 0
        last_pair = None

        for _, (box0, box1) in sorted(dists.items()):
            connections_done += 1
            if (not part2 and connections_done > connections_needed) or len(circuits) == 1:
                break

            if members[box0] == members[box1]:
                continue

            merger = members[box1]
            for box in circuits[merger]:
                members[box] = members[box0]
                circuits[members[box0]].append(box)
            del circuits[merger]
            last_pair = (box0, box1)

        if part2:
            return last_pair[0].x * last_pair[1].x
        else:
            ans = 1
            for _, v in sorted(circuits.items(), key=lambda k: len(k[1]), reverse=True)[:3]:
                ans *= len(v)

            return ans

    def part1(self) -> Any:
        return self.solve()


    def part2(self) -> Any:
        return self.solve(part2=True)


if __name__ == '__main__':
    day = Day(2025, 8)
    day.run(verbose=True)