#!/usr/bin/env python3

import sys
import typing

import matplotlib.pyplot as plt
import networkx as nx

input_file = sys.argv[1]

with open(input_file) as fd:
    lines = [line.rstrip() for line in fd.readlines()]

height = len(lines)
width = len(lines[0])

directions = [
    (-1, 0),  # ^ North
    (0, 1),  # > East
    (1, 0),  # v South
    (0, -1),  # < West
]

# Parse input
g = nx.DiGraph()
for i in range(height):
    for j in range(width):
        char = lines[i][j]
        if char == "#":
            continue
        for d, direction in enumerate(directions):
            cur = (i, j, d)
            # Start/end
            if char == "E":
                g.add_edge(cur, "end", weight=0)
            elif char == "S" and d == 1:
                g.add_edge("start", cur, weight=0)
            # Rotate
            g.add_edge(cur, (i, j, (d + 1) % len(directions)), weight=1000)
            g.add_edge(cur, (i, j, (d - 1) % len(directions)), weight=1000)
            # Advance
            ii, jj = i + direction[0], j + direction[1]
            if lines[ii][jj] == "#":
                continue
            g.add_edge(cur, (ii, jj, d), weight=1)

# Part 1
score = nx.shortest_path_length(g, "start", "end", weight="weight")
print(f"{score=}")

# Part 2
paths = nx.all_shortest_paths(g, "start", "end", weight="weight")
best_orientations = set()
for path in paths:
    best_orientations |= set(path)
path_edges = list(zip(path, path[1:]))  # Will be one random best path
best_places = set(bo[:2] for bo in best_orientations - {"start", "end"})
print(f"{len(best_places)=}")

# Draw graph
if len(g.nodes) > 1000:
    sys.exit(0)

node_colors = ["blue" if node in best_orientations else "cyan" for node in g.nodes()]
edge_colors = ["red" if edge in path_edges else "black" for edge in g.edges()]
node_pos: dict[typing.Any, tuple[float, float]] = dict()
for node in g.nodes():
    pos: tuple[float, float]
    if node == "start":
        pos = height - 1, 0
    elif node == "end":
        pos = 0, width - 1
    else:
        i, j, d = node
        direction = directions[d]
        pos = i + direction[0] / 3, j + direction[1] / 3
    node_pos[node] = pos[1], pos[0] * -1

nx.draw_networkx_nodes(g, node_pos, node_color=node_colors)
nx.draw_networkx_edges(g, node_pos, edge_color=edge_colors)
# nx.draw_networkx_labels(g, node_pos)
# nx.draw_networkx_edge_labels(
#     g, node_pos, edge_labels={(u, v): d["weight"] for u, v, d in g.edges(data=True)}
# )

plt.show()