from __future__ import annotations
from collections import defaultdict
from tools.aoc import AOCDay
from tools.coordinate import Coordinate
from tools.grid import Grid
from typing import Any


ICE = {
    '>': Coordinate(1, 0),
    'v': Coordinate(0, 1)
}


class GraphNode(Coordinate):
    def __init__(self, x: int, y: int, z: int = None):
        self.children = {}


def build_graph(grid: Grid, part2: bool = False) -> GraphNode:
    root = GraphNode(1, 0)
    crossings = {}
    end = Coordinate(grid.maxX, grid.maxY)

    def traverse_grid(node: GraphNode, next_coord: Coordinate):
        count = 1
        n = set(grid.getNeighboursOf(next_coord, includeDefault=False, includeDiagonal=False))
        last_coord = node
        while len(n - {last_coord}) == 1:
            count += 1
            last_coord, next_coord = next_coord, (n - {last_coord}).pop()
            n = set(grid.getNeighboursOf(next_coord, includeDefault=False, includeDiagonal=False))

        if len(n) == 1:  # dead end
            if next_coord == end:
                node.children[GraphNode(*next_coord)] = count
        else:  # crossing
            if next_coord in crossings:
                if part2:
                    crossings[next_coord].children[node] = count
                node.children[crossings[next_coord]] = count
            else:
                this_node = GraphNode(*next_coord)
                crossings[next_coord] = this_node
                node.children[this_node] = count
                for c in n:
                    if part2 or grid.get(c) not in ICE or ICE[grid.get(c)] == c - next_coord:
                        traverse_grid(this_node, c)

    traverse_grid(root, Coordinate(1, 1))
    return root


def get_longest_path(root: GraphNode, looking_for: GraphNode) -> int:
    D = defaultdict(int)
    v = set()

    def calc_dist(node: GraphNode, dist: int):
        if node in v:
            return
        v.add(node)
        if D[node] < dist:
            D[node] = dist

        for c_node, c_dist in node.children.items():
            calc_dist(c_node, dist + c_dist)

        v.remove(node)

    calc_dist(root, 0)
    return D[looking_for]


class Day(AOCDay):
    inputs = [
        [
            (94, "input23_test"),
            (2326, "input23"),
        ],
        [
            (154, "input23_test"),
            (6574, "input23"),
        ]
    ]

    def parse_input(self, part2: bool = False) -> (GraphNode, Coordinate):
        grid = Grid.from_data(self.getInput(), translate={'#': False})
        root = build_graph(grid, part2=part2)
        return root, Coordinate(grid.maxX, grid.maxY)

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

    def part2(self) -> Any:
        return get_longest_path(*self.parse_input(part2=True))


if __name__ == '__main__':
    day = Day(2023, 23)
    day.run(verbose=True)