import heapq
import math
from typing import List, Tuple, Dict, Optional, Callable

# Define directions for 4-way or 8-way movement
FOUR_WAY = [(0, 1), (1, 0), (0, -1), (-1, 0)]                    # Up, Right, Down, Left
EIGHT_WAY = [(0, 1), (1, 0), (0, -1), (-1, 0), 
             (1, 1), (1, -1), (-1, 1), (-1, -1)]             # Includes diagonals

class Node:
    """
    Represents a node in the grid.
    """
    def __init__(self, x: int, y: int, cost: float = 1.0, obstacle: bool = False):
        self.x = x
        self.y = y
        self.cost = cost      # Cost to enter this node (usually 1, higher for rough terrain)
        self.obstacle = obstacle

    def __lt__(self, other):
        # Required for priority queue (heapq compares objects)
        return False

def heuristic(a: Tuple[int, int], b: Tuple[int, int], heuristic_type: str = "manhattan") -> float:
    """
    Calculate heuristic distance between two points.
    Supported: 'manhattan', 'euclidean', 'octile'
    """
    x1, y1 = a
    x2, y2 = b
    dx = abs(x1 - x2)
    dy = abs(y1 - y2)

    if heuristic_type == "manhattan":
        return dx + dy
    elif heuristic_type == "euclidean":
        return math.sqrt(dx * dx + dy * dy)
    elif heuristic_type == "octile":
        # Best for 8-way movement
        return max(dx, dy) + (math.sqrt(2) - 1) * min(dx, dy)
    else:
        return dx + dy  # default to manhattan


def a_star_grid(
    grid: List[List[int]], 
    start: Tuple[int, int], 
    goal: Tuple[int, int],
    heuristic_type: str = "manhattan",
    allow_diagonal: bool = True
) -> Optional[List[Tuple[int, int]]]:
    """
    A* algorithm on a 2D grid.

    Parameters:
        grid: 2D list where 0 = walkable, 1 = obstacle
        start: (x, y) starting position
        goal: (x, y) target position
        heuristic_type: 'manhattan', 'euclidean', or 'octile'
        allow_diagonal: True for 8-way movement, False for 4-way

    Returns:
        List of tuples as path from start to goal (including both), or None if no path
    """
    rows = len(grid)
    cols = len(grid[0]) if rows > 0 else 0

    # Validate start and goal
    if not (0 <= start[0] < rows and 0 <= start[1] < cols and grid[start[0]][start[1]] == 0):
        print("Invalid start position")
        return None
    if not (0 <= goal[0] < rows and 0 <= goal[1] < cols and grid[goal[0]][goal[1]] == 0):
        print("Invalid goal position")
        return None

    directions = EIGHT_WAY if allow_diagonal else FOUR_WAY

    # Priority queue: (f_score, g_score, x, y)
    open_set = []
    heapq.heappush(open_set, (0 + heuristic(start, goal, heuristic_type), 0, start[0], start[1]))

    # Track the best path
    came_from: Dict[Tuple[int, int], Tuple[int, int]] = {}

    # Cost from start to node
    g_score: Dict[Tuple[int, int], float] = {(start[0], start[1]): 0}

    # To avoid revisiting
    closed_set: set[Tuple[int, int]] = set()

    while open_set:
        _, current_g, x, y = heapq.heappop(open_set)
        current = (x, y)

        if current == goal:
            # Reconstruct and return path
            path = []
            while current in came_from:
                path.append(current)
                current = came_from[current]
            path.append(start)
            path.reverse()
            return path

        if current in closed_set:
            continue

        closed_set.add(current)

        # Explore neighbors
        for dx, dy in directions:
            nx, ny = x + dx, y + dy

            if 0 <= nx < rows and 0 <= ny < cols and grid[nx][ny] == 0:
                neighbor = (nx, ny)

                if neighbor in closed_set:
                    continue

                # Diagonal move cost is sqrt(2), otherwise 1
                move_cost = 1.414 if abs(dx) + abs(dy) == 2 else 1.0
                tentative_g = g_score[current] + move_cost

                if neighbor not in g_score or tentative_g < g_score[neighbor]:
                    came_from[neighbor] = current
                    g_score[neighbor] = tentative_g
                    f_score = tentative_g + heuristic(neighbor, goal, heuristic_type)
                    heapq.heappush(open_set, (f_score, tentative_g, nx, ny))

    return None  # No path found


# =============== Example Usage ===============
if __name__ == "__main__":
    # 0 = free space, 1 = wall/obstacle
    grid = [
        [0, 0, 0, 0, 0, 0],
        [0, 1, 1, 1, 1, 0],
        [0, 0, 0, 0, 1, 0],
        [1, 1, 1, 0, 1, 0],
        [0, 0, 1, 0, 0, 0],
        [0, 0, 0, 0, 0, 0]
    ]

    start = (0, 0)
    goal = (5, 5)

    path = a_star_grid(grid, start, goal, heuristic_type="octile", allow_diagonal=True)

    if path:
        print(f"Path found with {len(path)} steps:")
        for x, y in path:
            print(f"({x}, {y})", end=" → ")
        print("GOAL")

        # Visual representation
        for i in range(len(grid)):
            for j in range(len(grid[0])):
                if (i, j) == start:
                    print("S", end=" ")
                elif (i, j) == goal:
                    print("G", end=" ")
                elif (i, j) in path:
                    print("·", end=" ")
                elif grid[i][j] == 1:
                    print("█", end=" ")
                else:
                    print(".", end=" ")
            print()
    else:
        print("No path found!")

Path found with 13 steps:
(0, 0) → (1, 1) → (2, 2) → ... → (5, 5) → GOAL

S · · . . . 
. █ █ █ █ . 
. . · · █ . 
█ █ █ · █ . 
. . █ · · · 
. . . . . G