import networkx as nx
from networkx import Graph
from pyvis.network import Network
import heapq
from copy import deepcopy

from Routes.Red_Blue_Route import ARCTICUNO, BOULDER_BADGE, CASCADE_BADGE, CERULEAN_CAVE, CHAMPION, CUT, FLASH, GIOVANNI_FIGHT, MARSH_BADGE, MEWTWO, MOLTRES, PALLET_TOWN, POKE_FLUTE, QUENCHED_THURST, RAINBOW_BADGE, SILPH_SCOPE, SOUL_BADGE, SS_ANNE_TICKET, THUNDER_BADGE, VOLCANO_BADGE, ZAPDOS, get_red_blue_route

class PlayerState:
    def __init__(self):
        self.visited_nodes = set()  # Keeps track of nodes visited
        self.conditions_met = set()  # Tracks conditions the player has acquired

    def has_condition(self, condition):
        return condition in self.conditions_met

    def add_condition(self, condition):
        self.conditions_met.add(condition)

def find_path_to_goal(G, start_node, goal_condition, player_state):
    # Priority queue elements: (total_cost, current_node, path, state)
    queue = []
    initial_state = deepcopy(player_state)
    heapq.heappush(queue, (0, start_node, [start_node], initial_state))
    visited = set()

    while queue:
        total_cost, current_node, path, state = heapq.heappop(queue)

        if current_node in visited:
            continue
        visited.add(current_node)

        # Visit the current node and update the player's state if needed
        new_state = deepcopy(state)
        visit_node(current_node, G, new_state)

        # Check if the goal condition is met
        if goal_condition in new_state.conditions_met:
            # Commit the updated state to the player state
            player_state.visited_nodes = new_state.visited_nodes
            player_state.conditions_met = new_state.conditions_met
            return path

        # Explore neighbors
        for neighbor in G.neighbors(current_node):
            edge_attrs = G[current_node][neighbor]
            if can_traverse(edge_attrs, new_state):
                edge_weight = edge_attrs.get('weight', 1)
                new_total_cost = total_cost + edge_weight
                new_path = path + [neighbor]
                heapq.heappush(queue, (new_total_cost, neighbor, new_path, new_state))

    return None  # No path found to the goal

def find_optimal_path_with_goals(G, start_node, end_node, goals):
    current_node = start_node
    full_path = [start_node]  # Start with the starting node in the path
    player_state = PlayerState()

    # Iterate through each goal to find paths in succession
    for goal in goals:
        path_to_goal = find_path_to_goal(G, current_node, goal, player_state)
        if path_to_goal is None:
            print(f"Failed to find a path to the goal: {goal}")
            return None
        
        # Extend the full path with all nodes in path_to_goal
        # Avoid repeating the current_node, so exclude the first node
        full_path.extend(path_to_goal[1:])  
        
        # Update current node to the goal's location
        current_node = path_to_goal[-1] 

    # Finally, find the path to the end node
    #path_to_end = find_path_to_goal(G, current_node, None, player_state)
    #if path_to_end is None:
    #    print("Failed to find a path to the end node.")
    #    return None

    # Extend the full path with all nodes in path_to_end, excluding the first node to avoid repetition
    #full_path.extend(path_to_end[1:])

    return full_path

def visit_node(node_name, G, player_state: PlayerState):
    # Add the current node to the visited set
    player_state.visited_nodes.add(node_name)

    # Grant conditions or items found at this node
    node_attrs = G.nodes[node_name]
    grants_conditions = node_attrs.get('grants_conditions', [])

    for grant in grants_conditions:
        condition = grant['condition']
        required_conditions = grant.get('required_conditions', [])

        # Only add conditions if requirements are met
        if all(cond in player_state.conditions_met for cond in required_conditions):
            if not player_state.has_condition(condition):
                player_state.add_condition(condition)
                print(f"Condition '{condition}' granted upon visiting '{node_name}'.")

def can_traverse(edge_attrs, player_state: PlayerState):
    # Check if the player meets all the conditions required to traverse this edge
    conditions = edge_attrs.get('condition')
    if not conditions:
        return True
    return all(condition in player_state.conditions_met for condition in conditions)

# Example usage
if __name__ == "__main__":

    G = get_red_blue_route()

    # Define the ordered goals that the player must acquire
    goals = [
        BOULDER_BADGE, 
        CASCADE_BADGE, 
        SS_ANNE_TICKET, 
        CUT, 
        THUNDER_BADGE, 
        FLASH,
        SILPH_SCOPE,
        RAINBOW_BADGE,
        QUENCHED_THURST,
        POKE_FLUTE,
        GIOVANNI_FIGHT,
        MARSH_BADGE,
        SOUL_BADGE,
        VOLCANO_BADGE,
        ARCTICUNO,
        ZAPDOS,
        MOLTRES,
        CHAMPION,
        MEWTWO
    ]

    # Find the optimal path while fulfilling goals in succession
    optimal_path = find_optimal_path_with_goals(G, PALLET_TOWN, CERULEAN_CAVE, goals)

    if optimal_path:
        print("Optimal Path:", " -> ".join(optimal_path))
    else:
        print("No path found to fulfill all goals.")