import numpy as np
import random
import heapq
import numpy as np
import random
from queue import Queue
class PathNotFoundException(Exception):
    pass

def create_city_network_input(grid_size=30):  # TODO: make this work for varying shapes
    WHITE, BLACK, START, END = 0, 1, 2, 3

    def create_path(grid, start, end):
        queue = Queue()
        queue.put(start)
        visited = set([start])
        parent = {start: None}
        
        while not queue.empty():
            current = queue.get()
            if current == end:
                break
            
            for dx, dy in [(0, 1), (1, 0), (0, -1), (-1, 0)]:
                next_pos = (current[0] + dx, current[1] + dy)
                if (0 <= next_pos[0] < grid_size and 0 <= next_pos[1] < grid_size 
                    and next_pos not in visited):
                    queue.put(next_pos)
                    visited.add(next_pos)
                    parent[next_pos] = current
        
        if end not in parent:
            return False
        
        # Create the path
        current = end
        while current != start:
            if grid[current] == WHITE:
                grid[current] = BLACK
            current = parent[current]
        
        return True

    def create_network():
        grid = np.zeros((grid_size, grid_size), dtype=int)
        
        # Place start and end points with minimum distance
        min_distance = grid_size // 2
        while True:
            start = (random.randint(0, grid_size-1), random.randint(0, grid_size-1))
            end = (random.randint(0, grid_size-1), random.randint(0, grid_size-1))
            if abs(start[0] - end[0]) + abs(start[1] - end[1]) >= min_distance:
                break
        
        grid[start] = START
        grid[end] = END
        
        # Ensure a path exists between start and end
        create_path(grid, start, end)
        
        # Create additional sparse paths
        for _ in range(grid_size * 3):  # Increased number of paths
            x, y = random.randint(0, grid_size-1), random.randint(0, grid_size-1)
            length = random.randint(3, grid_size // 2)
            direction = random.choice([(0, 1), (1, 0), (1, 1), (1, -1)])
            
            for i in range(length):
                nx, ny = x + i * direction[0], y + i * direction[1]
                if 0 <= nx < grid_size and 0 <= ny < grid_size:
                    if grid[nx, ny] == WHITE:
                        grid[nx, ny] = BLACK
        
        return grid

    return create_network(), {}

def city_network_shortest_path(inputs, key):
    grid, _ = inputs
    
    WHITE, BLACK, START, END = 0, 1, 2, 3
    PATH = 4

    grid = np.array(grid)

    def find_start_end():
        start, end = None, None
        for i in range(len(grid)):
            for j in range(len(grid[0])):
                if grid[i, j] == START:
                    start = (i, j)
                elif grid[i, j] == END:
                    end = (i, j)
                if start and end:
                    return start, end
        raise ValueError("Start or end point not found in the grid")

    def get_neighbors(pos):
        directions = [(0, -1), (0, 1), (-1, 0), (1, 0)]
        for dx, dy in directions:
            nx, ny = pos[0] + dx, pos[1] + dy
            if 0 <= nx < len(grid) and 0 <= ny < len(grid[0]):
                yield (nx, ny)

    def heuristic(a, b):
        return abs(a[0] - b[0]) + abs(a[1] - b[1])

    def cost(next):
        if grid[next] == WHITE:
            return float('inf')
        return 1

    def a_star(start, goal):
        frontier = []
        counter = 0
        heapq.heappush(frontier, (0, counter, start))
        came_from = {start: None}
        cost_so_far = {start: 0}
        
        while frontier:
            current = heapq.heappop(frontier)[2]
            
            if current == goal:
                return reconstruct_path(came_from, start, goal)
            
            for next in get_neighbors(current):
                next_cost = cost(next)
                if next_cost != float('inf'):
                    new_cost = cost_so_far[current] + next_cost
                    if next not in cost_so_far or new_cost < cost_so_far[next]:
                        cost_so_far[next] = new_cost
                        priority = new_cost + heuristic(goal, next)
                        counter += 1
                        heapq.heappush(frontier, (priority, counter, next))
                        came_from[next] = current
        
        raise PathNotFoundException("No path found from start to end")

    def reconstruct_path(came_from, start, goal):
        current = goal
        path = []
        while current != start:
            path.append(current)
            current = came_from[current]
        path.append(start)
        path.reverse()
        return path

    start, end = find_start_end()
    try:
        path = a_star(start, end)
    except PathNotFoundException:
        # Re-raise the exception to be consistent with the original behavior
        raise

    for pos in path[1:-1]:
        grid[pos] = PATH

    return grid

# Test function to check the frequency of unsolvable grids
def test_pathfinding(num_tests=100):
    unsolvable = 0
    for _ in range(num_tests):
        grid, _ = create_city_network_input()
        try:
            city_network_shortest_path((grid, None), None)
        except PathNotFoundException:
            unsolvable += 1
    
    print(f"Unsolvable grids: {unsolvable / num_tests * 100:.2f}%")

test_pathfinding()