import logging

logger = logging.getLogger(__name__)


with open('2024/data/input16.data') as f:
# with open('2024/data/input16.example') as f:
    # Dont need edges, these are irrelevant to us
    content = [list(line.strip()[1:-1]) for line in f.readlines()[1:-1]]

directions = {'N': (-1, 0), 'E': (0, 1), 'S': (1, 0), 'W': (0, -1)}

start = next((row, col, 'E') for row in range(len(content)) for col in range(len(content[0])) if content[row][col] == 'S')
end = next((row, col) for row in range(len(content)) for col in range(len(content[0])) if content[row][col] == 'E')

# Python it up^^ (outer loop goes over each cell in grid and continues with non-walls,
# for each valid cell go over all directions and applies costs dependent on turning)
graph = { (row, col, dirName): {(row+nextdirRow, col+nextdirCol, nextDir): (0 if dirName == nextDir else 1000) + 1
    for nextDir, (nextdirRow, nextdirCol) in directions.items()
    if 0 <= row+nextdirRow < len(content) and 0 <= col+nextdirCol < len(content[0]) and content[row+nextdirRow][col+nextdirCol] != '#'}
    for row in range(len(content)) for col in range(len(content[0])) if content[row][col] != '#' for dirName in directions }

# Algo
def dijsktra(graph, start, end):
    import heapq    # Priority Queue, Do not use own for shorter Code^^

    pq, costs, predecessors = [(0, start)], {node: float('inf') for node in graph}, {node: None for node in graph}
    costs[start] = 0
    while pq:
        cost, node = heapq.heappop(pq)
        logger.debug(f"Procesing node {node} with current cost {cost}")
        if node[:2] == end[:2]:  # Stop when reaching the end position, dont care about direction
            path = []
            while node:
                path.append(node)
                node = predecessors[node]
            logger.debug(f"Final cost: {cost}")
            return cost, path[::-1]  # Reverse path to get start-to-end order
        if cost > costs[node]: continue

        for neighbor, edge_cost in graph[node].items():
            if (new_cost := cost + edge_cost) < costs[neighbor]:
                costs[neighbor], predecessors[neighbor] = new_cost, node
                heapq.heappush(pq, (new_cost, neighbor))
                logger.debug(f"Relaxing edge {node} -> {neighbor}, edge cost: {edge_cost}, new cost: {new_cost}")
    return float('inf'), []  # No path found

def drawGraph(content, path=None):
    # Create copy of grid to overlay graph -> No inPlace Stuff
    grid = [row[:] for row in content]
    
    if path:
        symbols = {'N': '^', 'E': '>', 'S': 'v', 'W': '<'}
        for row, col, direction in path:
            grid[row][col] = symbols[direction]
    return '\n'.join(''.join(row) for row in grid)

costShortest, shortestPath = dijsktra(graph, start, end)
print("Shortest path cost:", costShortest)
print(drawGraph(content, shortestPath))