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 CERULEAN_CAVE, PALLET_TOWN, get_red_blue_route

class PlayerState:
    def __init__(self):
        self.visited_nodes = {}
        self.conditions_met = set()
        self.goals = set()  # Tracks unmet conditions required for future traversal

def find_optimal_path(G, start_node, end_node):
    # Priority queue elements: (total_cost, current_node, player_state, path)
    queue = []
    initial_state = PlayerState()
    heapq.heappush(queue, (0, start_node, [start_node], initial_state))
    visited_states = {}

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

        # If end node is reached
        if current_node == end_node:
            return path
        
        if current_node in player_state.visited_nodes:
            player_state.visited_nodes[current_node] += 1
        else:
            player_state.visited_nodes[current_node] = 1

        # Visit the current node and update the player's state
        visit_node(current_node, G, player_state)

        # Explore neighbors
        for neighbor in G.neighbors(current_node):
            edge_attrs = G[current_node][neighbor]
            if can_traverse(edge_attrs, player_state):
                visit_count = player_state.visited_nodes.get(neighbor, 0)
                edge_weight = max(1, visit_count)

                # Apply weight reduction if this neighbor helps meet a goal
                if any(condition in player_state.goals for condition in player_state.conditions_met):
                    edge_weight = max(1, edge_weight - 1)  # Reduce weight to incentivize this path

                new_total_cost = total_cost + edge_weight

                # Copy path, but keep player_state as is (state should persist across paths)
                new_path = path + [neighbor]
                heapq.heappush(queue, (new_total_cost, neighbor, new_path, deepcopy(player_state)))
            else:
                unmet_conditions = edge_attrs.get('condition', [])
                for condition in unmet_conditions:
                    if condition not in player_state.conditions_met:
                        player_state.goals.add(condition)

    return None  # No path found

def can_traverse(edge_attrs, player_state: PlayerState):
    conditions = edge_attrs.get('condition')
    if not conditions:
        return True
    return any(condition in conditions for condition in player_state.conditions_met)

def visit_node(node_name, G, player_state: PlayerState):
    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', [])
        
        if all(cond in player_state.conditions_met for cond in required_conditions):
            if condition not in player_state.conditions_met:
                player_state.conditions_met.add(condition)
                print(f"Condition '{condition}' granted upon visiting '{node_name}'.")
            
            if condition in player_state.goals:
                player_state.goals.remove(condition)
        else:
            print(f"Cannot obtain '{condition}' from '{node_name}' because required conditions are not met: {required_conditions}")

def get_available_moves(current_node, G, player_state):
    available_moves = []
    for neighbor in G.neighbors(current_node):
        edge_attrs = G[current_node][neighbor]
        if can_traverse(edge_attrs, player_state):
            available_moves.append(neighbor)
    return available_moves


def main():
    player_state = {
        'visited_nodes': set(),
        'conditions_met': set()
    }

    G = get_red_blue_route()

    net = Network(notebook=True)
    optimal_path = find_optimal_path(G, PALLET_TOWN, CERULEAN_CAVE)

    if optimal_path is None:
        return

    # Create a set of edges in the optimal path for quick lookup
    optimal_edges = set(zip(optimal_path, optimal_path[1:]))

    # Add nodes to the network
    for node, attrs in G.nodes(data=True):
        net.add_node(
            node,
            label=node,
            title=f"Type: {attrs['node_type']}",
            color='lightblue' if attrs['node_type'] == 'location' else 'lightgreen'
        )

    # Add edges, highlighting those in the optimal path
    for source, target, attrs in G.edges(data=True):
        condition = attrs['condition']
        edge_label = f"Condition: {condition}" if condition else "No condition"
        edge_color = 'blue' if (source, target) in optimal_edges else 'gray'
        net.add_edge(
            source,
            target,
            title=edge_label,
            label=edge_label,
            color=edge_color
        )

    net.show('optimal_pokemon_map.html')

    net2 = Network(notebook=True)
    # Add nodes with special styling for conditional grants
    for node, attrs in G.nodes(data=True):
        color = 'lightblue'
        title = f"Type: {attrs.get('node_type', 'Unknown')}"
        if 'grants_condition' in attrs:
            title += f"<br>Grants: {attrs['grants_condition']}"
            if 'required_conditions' in attrs:
                title += f" (Requires: {attrs['required_conditions']})"
                color = 'orange'  # Highlight nodes with conditional grants
            else:
                color = 'green'  # Nodes that grant conditions without prerequisites
        net2.add_node(node, label=node, title=title, color=color)

    # Add edges
    for source, target, attrs in G.edges(data=True):
        condition = attrs.get('condition')
        edge_label = f"Condition: {condition}" if condition else "No condition"
        color = 'red' if condition else 'gray'
        net2.add_edge(source, target, title=edge_label, label=edge_label, color=color)

    net2.show('pokemon_map_with_conditions.html')


if __name__ == "__main__":
    try:
        main()
    finally:
        pass