import heapq
import networkx as nx
from typing import List, Set

FLY_OUT_OF_BATTLE = 'Fly out of battle'

class State:
    def __init__(self, location, conditions, cost, path, visited_required_nodes):
        self.location = location
        self.conditions = conditions  # A frozenset of conditions
        self.cost = cost
        self.path = path  # List of locations visited in order
        self.visited_required_nodes = visited_required_nodes  # A frozenset of required nodes visited

    def __lt__(self, other):
        return self.cost < other.cost  # For priority queue

def heuristic(state, goal_conditions, required_nodes):
    # Since we don't have actual distances, we can use the number of badges remaining as the heuristic
    remaining_conditions = goal_conditions - state.conditions
    remaining_nodes = required_nodes - state.visited_required_nodes
    return len(remaining_conditions) + len(remaining_nodes)

def is_goal_state(state, goal_location, goals, required_nodes):
    return (
        state.location == goal_location and
        goals.issubset(state.conditions) and
        required_nodes.issubset(state.visited_required_nodes)
    )

class PokemonGameDesc:
    def __init__(self):
        self.game_name: str = ""
        self.towns_and_cities: Set[str] = set()
        self.badges: Set[str] = set()
        self.items: Set[str] = set()
        self.hms: Set[str] = set()
        self.starting_town: str
        self.end_goal: str
        self.flying_badge: str
        self.additional_goals: List[str] = []
        self.one_way_routes: List[str] = []
        self.must_visit: Set[str] = set()
        self.graph: nx.Graph = nx.Graph()

    def astar_search(self):
        from collections import deque

        self.goals = set(self.badges + self.additional_goals)

        # Priority queue for open states
        open_list = []
        heapq.heappush(open_list, (0, State(
            location=self.starting_town,
            conditions=frozenset(),  # Start with no conditions
            cost=0,
            path=[self.starting_town],
            visited_required_nodes=frozenset([self.starting_town]) if self.starting_town in self.must_visit else frozenset()
        )))

        # Closed set to keep track of visited states
        closed_set = {}

        while open_list:
            _, current_state = heapq.heappop(open_list)

            # Check if we've reached the goal location with all required conditions
            if is_goal_state(current_state, self.end_goal, self.goals, self.must_visit):
                return current_state.path, current_state.cost, current_state.conditions

            # Check if we've already visited this state with equal or better conditions
            state_key = (current_state.location, current_state.conditions, current_state.visited_required_nodes)
            if state_key in closed_set and closed_set[state_key] <= current_state.cost:
                continue  # Skip this state

            closed_set[state_key] = current_state.cost

            # Expand neighbors via normal moves
            for neighbor in self.graph.neighbors(current_state.location):
                edge_data = self.graph.get_edge_data(current_state.location, neighbor)
                edge_condition = edge_data.get('condition', [])

                if edge_condition is None:
                    edge_requires = set()
                else:
                    edge_requires = set(edge_condition)

                # Check if we have the required conditions to traverse this edge
                if not edge_requires.issubset(current_state.conditions):
                    continue  # Can't traverse this edge

                # Update conditions based on grants at the neighbor node
                neighbor_data = self.graph.nodes[neighbor]
                new_conditions = set(current_state.conditions)

                # Check if the neighbor grants any conditions
                grants = neighbor_data.get('grants_conditions', [])
                for grant in grants:
                    required_for_grant = set(grant.get('required_conditions', []))
                    if required_for_grant.issubset(new_conditions):
                        # We can acquire the condition
                        new_conditions.add(grant['condition'])

                # Update visited required nodes
                new_visited_required_nodes = set(current_state.visited_required_nodes)
                if neighbor in self.must_visit:
                    new_visited_required_nodes.add(neighbor)

                new_state = State(
                    location=neighbor,
                    conditions=frozenset(new_conditions),
                    cost=current_state.cost + 1,  # Assuming uniform cost; adjust if needed
                    path=current_state.path + [neighbor],
                    visited_required_nodes=frozenset(new_visited_required_nodes)
                )

                estimated_total_cost = new_state.cost + heuristic(new_state, self.goals, self.must_visit)

                heapq.heappush(open_list, (estimated_total_cost, new_state))

            # Expand neighbors via FLY if applicable
            if FLY_OUT_OF_BATTLE in current_state.conditions and current_state.location in self.towns_and_cities:
                for fly_target in self.towns_and_cities:
                    if fly_target != current_state.location and fly_target in current_state.path:
                        # You can fly to this location
                        new_conditions = set(current_state.conditions)
                        neighbor_data = self.graph.nodes[fly_target]
                        grants = neighbor_data.get('grants_conditions', [])
                        for grant in grants:
                            required_for_grant = set(grant.get('required_conditions', []))
                            if required_for_grant.issubset(new_conditions):
                                new_conditions.add(grant['condition'])

                        # Update visited required nodes
                        new_visited_required_nodes = set(current_state.visited_required_nodes)
                        if fly_target in self.must_visit:
                            new_visited_required_nodes.add(fly_target)

                        fly_state = State(
                            location=fly_target,
                            conditions=frozenset(new_conditions),
                            cost=current_state.cost + 1,  # Adjust cost if flying is different
                            path=current_state.path + [fly_target],
                            visited_required_nodes=frozenset(new_visited_required_nodes)
                        )
                        estimated_total_cost = fly_state.cost + heuristic(fly_state, self.goals, self.must_visit)
                        heapq.heappush(open_list, (estimated_total_cost, fly_state))

        return None  # No path found 


__all__ = ["PokemonGameDesc"]