import heapq
import itertools
import math
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Set, Tuple
import pathlib


def read_data():
    with open(f"data/{pathlib.Path(__file__).stem}.txt") as raw_data:
        return [line.strip() for line in raw_data.readlines()]


@dataclass(frozen=True)
class Vec:
    x: int
    y: int

    def __add__(self, delta: 'Vec') -> 'Vec':
        return Vec(self.x + delta.x, self.y + delta.y)


def deltas() -> List[Vec]:
    return [Vec(-1, 0), Vec(1, 0), Vec(0, -1), Vec(0, 1)]


@dataclass
class Node:
    pos: Vec
    height: int
    adj: Dict[Vec, Optional['Node']] = field(default_factory=lambda: {d: None for d in deltas()})
    path_value: int = math.inf

    def neighbors(self):
        return [node for node in self.adj.values() if node]

    def __hash__(self):
        return hash(self.pos)

    def __eq__(self, other):
        return isinstance(other, Node) and self.pos == other.pos

    def __lt__(self, other):
        return isinstance(other, Node) and self.path_value < other.path_value


@dataclass
class NodeMap:
    nodes: List[List['Node']]

    def __getitem__(self, item: Vec) -> Optional[Node]:
        if (0 <= item.y < len(self.nodes)) and (0 <= item.x < len(self.nodes[0])):
            return self.nodes[item.y][item.x]
        return None

    def solve_adj(self):
        """ figures out what neighbors each node can access """
        for node in self.all_nodes():
            for delta, adj_node in [(delta, self[node.pos + delta]) for delta in deltas()]:
                if adj_node and adj_node.height <= node.height + 1:
                    node.adj[delta] = adj_node

    def dijkstra(self, start: Vec, end: Vec):
        """ Solves the shortest path between start and end. If path is possible, nodes forming the shortest path will have their path values updated with distance to start """
        start = self[start]
        end = self[end]

        unvisited: Set[Node] = {node for node in self.all_nodes()}
        for node in unvisited:
            node.path_value = math.inf
        start.path_value = 0

        next_node = [start]

        while next_node:
            current = heapq.heappop(next_node)
            if current == end:
                break
            if current not in unvisited:
                continue

            new_path_value = current.path_value + 1

            for node in current.neighbors():
                if node in unvisited:
                    node.path_value = min(node.path_value, new_path_value)
                    heapq.heappush(next_node, node)

            unvisited.remove(current)

    def all_nodes(self):
        yield from itertools.chain(*self.nodes)


def print_nodes(nodes: NodeMap):
    right = Vec(1, 0)
    left = Vec(-1, 0)
    up = Vec(0, -1)
    down = Vec(0, 1)
    for row in nodes.nodes:
        for node in row:
            if node.adj[up]:
                print("    ↑    ", end='')
            else:
                print("         ", end='')
        print()
        for node in row:
            if node.adj[left]:
                print("←──", end='')
            else:
                print("   ", end='')
            print(f"{node.path_value:^3}", end='')
            if node.adj[right]:
                print("──→", end='')
            else:
                print("   ", end='')
        print()
        for node in row:
            if node.adj[down]:
                print("    ↓    ", end='')
            else:
                print("         ", end='')
        print()


def parse_node_map(data: List[str]) -> Tuple[NodeMap, Vec, Vec]:
    nodes = [[Node(Vec(0, 0), 0) for _ in line] for line in data]
    start = None
    end = None

    for idx_row, row in enumerate(data):
        for idx_col, ch in enumerate(row):
            vec = Vec(idx_col, idx_row)
            match ch:
                case 'S':
                    height = 0
                    start = vec
                case 'E':
                    height = 25
                    end = vec
                case s:
                    height = ord(s) - ord('a')
            nodes[idx_row][idx_col] = Node(vec, height)

    node_map = NodeMap(nodes)
    node_map.solve_adj()

    return node_map, start, end


def part1(data: List[str]) -> int:
    node_map, start, end = parse_node_map(data)
    node_map.dijkstra(start, end)
    return node_map[end].path_value


def part2(data: List[str]) -> int:
    node_map, _, end = parse_node_map(data)
    best_path = math.inf
    for start_point in [node.pos for node in node_map.all_nodes() if node.height == 0]:
        node_map.dijkstra(start_point, end)
        best_path = min(node_map[end].path_value, best_path)
    return best_path


def main():
    data = read_data()
    print(f"Part 1: {part1(data)}")
    print(f"Part 2: {part2(data)}")


if __name__ == "__main__":
    main()