from collections import deque

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


class Node:
    def __init__(self, node_id: int, pos: Coordinate) -> None:
        self.node_id = node_id
        self.pos = pos
        self.others: dict[Node, int] = {}

    def __repr__(self):
        return f"Node({self.node_id}, {self.pos})"


def get_dist(grid: Grid, start: Coordinate, end: Coordinate) -> int:
    q = deque([(0, start)])
    seen = set()
    while q:
        dist, pos = q.popleft()
        if pos == end:
            return dist
        if pos in seen:
            continue
        seen.add(pos)
        for n in grid.getNeighboursOf(pos, includeDiagonal=False, includeDefault=False):
            q.append((dist + 1, n))


def get_min_dist(
    nodes: dict[int, Node], return_to_zero: bool = False, start: int = 0, dist: int = 0, seen: set = None
) -> int:
    if seen is None:
        seen = {start}

    if len(seen) == len(nodes):
        return dist + nodes[start].others[nodes[0]] if return_to_zero else dist

    min_dist = int(1e9)
    for n, n_dist in nodes[start].others.items():
        if n.node_id in seen or n.node_id == start:
            continue
        this_dist = get_min_dist(
            nodes, return_to_zero=return_to_zero, start=n.node_id, dist=dist + n_dist, seen=seen | {n.node_id}
        )
        if this_dist < min_dist:
            min_dist = this_dist

    return min_dist


class Day(AOCDay):
    inputs = [
        [
            (14, "input24_test"),
            (502, "input24"),
        ],
        [
            (724, "input24"),
        ],
    ]

    def parse_input(self) -> dict[int, Node]:
        grid = Grid.from_data(self.getInput(), translate={"[0-9]": int}, default="#")
        nodes = {}
        for c in grid.getActiveCells():
            if isinstance(grid.get(c), int):
                nodes[grid.get(c)] = Node(grid.get(c), c)
                grid.set(c, ".")

        for node_a in nodes.values():
            for node_b in nodes.values():
                if node_b.node_id == node_a.node_id:
                    continue

                if node_a.node_id in node_b.others:
                    continue

                dist = get_dist(grid, node_a.pos, node_b.pos)
                node_a.others[node_b] = dist
                node_b.others[node_a] = dist

        return nodes

    def part1(self) -> Any:
        return get_min_dist(self.parse_input())

    def part2(self) -> Any:
        return get_min_dist(self.parse_input(), True)


if __name__ == "__main__":
    day = Day(2016, 24)
    day.run(verbose=True)