main.py

import sys
import random
import pygame
import networkx as nx
import numpy as np

from src.world.road_generator import generate_road_network
from src.world.city import City
from src.world.influence_map import InfluenceMapManager
from src.ai.pathfinding import Pathfinder
from src.ai.prediction import TrajectoryPredictor
from src.ai.coordination import RoleAllocator
from src.agents.player import PlayerVehicle
from src.agents.cop import CopVehicle
from src.render.camera import Camera
from src.render.renderer import Renderer

def find_furthest_spawn_edge(city, player_node, excluded_nodes=None):
    """
    Finds a node furthest from the player_node and returns a valid spawn edge starting at that node.
    """
    graph = city.graph
    try:
        lengths = nx.single_source_dijkstra_path_length(graph, player_node, weight='length')
        sorted_nodes = sorted(lengths.items(), key=lambda x: x[1], reverse=True)
    except Exception:
        sorted_nodes = [(n, 0) for n in graph.nodes()]
        random.shuffle(sorted_nodes)

    if excluded_nodes:
        sorted_nodes = [item for item in sorted_nodes if item[0] not in excluded_nodes]

    # Pick from the furthest top 15 nodes
    candidates = [node for node, dist in sorted_nodes[:15]]
    if not candidates:
        candidates = list(graph.nodes())

    spawn_node = random.choice(candidates)
    neighbors = list(graph.neighbors(spawn_node))
    if neighbors:
        return (spawn_node, neighbors[0])
    return (spawn_node, spawn_node)

def main():
    # Initialize Pygame
    pygame.init()
    pygame.display.set_caption("Interpose - Real-Time Pursuit-Evasion")
    
    screen_width, screen_height = 1024, 768
    screen = pygame.display.set_mode((screen_width, screen_height))
    clock = pygame.time.Clock()
    
    # Simulation seed
    current_seed = random.randint(0, 100000)
    
    # Game state variables
    show_debug = True
    active_strategy = "GREEDY"             # "GREEDY", "PREDICTIVE_INTERCEPT", or "COORDINATED"
    prediction_model = "NAIVE"             # "NAIVE" or "ROAD_AWARE"
    current_time = 0.0
    player_history = []                    # sliding history buffer: [(timestamp, edge, progress, speed, heading, pos)]
    cop_escalation_mode = True             # True: progressive difficulty scaling, False: 4 cops at start
    
    # Game rules
    game_over = False
    game_result = None                     # "BUSTED" or "VICTORY"
    time_left = 60.0                       # Survive for 60 seconds
    
    # Roadblock variables
    roadblock_cooldown = 0.0
    roadblock_lifetimes = {}               # sorted_edge -> time_left
    
    # Setup world state
    def reset_simulation(seed):
        nonlocal current_seed, current_time, player_history, game_over, game_result, time_left
        nonlocal roadblock_cooldown, roadblock_lifetimes, cop_escalation_mode
        current_seed = seed
        current_time = 0.0
        player_history = []
        game_over = False
        game_result = None
        time_left = 60.0
        
        roadblock_cooldown = 0.0
        roadblock_lifetimes.clear()
        
        # Generate new city graph
        graph = generate_road_network(rows=8, cols=8, spacing=130, cycle_prob=0.35, remove_node_prob=0.1, seed=current_seed)
        city = City(graph)
        pathfinder = Pathfinder(graph)
        predictor = TrajectoryPredictor(city)
        role_allocator = RoleAllocator(city)
        influence_map_manager = InfluenceMapManager(city)
        
        # Spawn player
        edges = list(graph.edges())
        if not edges:
            raise ValueError("Graph has no edges to spawn agents on.")
        player_edge = random.choice(edges)
        player = PlayerVehicle(city, player_edge, speed=100.0)
        
        # Spawn cops
        cops = []
        player_node = player_edge[0]
        initial_cops_count = 1 if cop_escalation_mode else 4
        excluded = {player_node}
        
        for _ in range(initial_cops_count):
            cop_edge = find_furthest_spawn_edge(city, player_node, excluded_nodes=excluded)
            excluded.add(cop_edge[0])
            cop = CopVehicle(city, cop_edge, player, pathfinder, predictor, speed=75.0)
            cops.append(cop)
            
        # Camera centered on player
        camera = Camera(screen_width, screen_height, start_pos=player.pos, zoom=0.8, lerp_speed=6.0)
        
        # Renderer
        renderer = Renderer(screen, city, camera)
        renderer.show_debug = show_debug
        
        return city, pathfinder, predictor, role_allocator, influence_map_manager, player, cops, camera, renderer

    # Initial boot
    city, pathfinder, predictor, role_allocator, influence_map_manager, player, cops, camera, renderer = reset_simulation(current_seed)
    
    # Main simulation loop
    running = True
    while running:
        # Limit to 60fps and calculate delta-time in seconds
        dt = clock.tick(60) / 1000.0
        # Prevent huge dt spikes when dragging window
        dt = min(dt, 0.1)

        # 1. Event Handling
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
                
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    running = False
                elif event.key == pygame.K_h:
                    show_debug = not show_debug
                    renderer.show_debug = show_debug
                elif event.key == pygame.K_r:
                    city, pathfinder, predictor, role_allocator, influence_map_manager, player, cops, camera, renderer = reset_simulation(random.randint(0, 100000))
                elif event.key == pygame.K_m:
                    cop_escalation_mode = not cop_escalation_mode
                    city, pathfinder, predictor, role_allocator, influence_map_manager, player, cops, camera, renderer = reset_simulation(current_seed)
                elif event.key == pygame.K_b:
                    player.autoplay = not player.autoplay
                elif event.key == pygame.K_1:
                    active_strategy = "GREEDY"
                elif event.key == pygame.K_2:
                    active_strategy = "PREDICTIVE_INTERCEPT"
                elif event.key == pygame.K_3:
                    active_strategy = "COORDINATED"
                elif event.key == pygame.K_p:
                    prediction_model = "ROAD_AWARE" if prediction_model == "NAIVE" else "NAIVE"
                elif event.key == pygame.K_EQUALS:
                    camera.zoom_in()
                elif event.key == pygame.K_MINUS:
                    camera.zoom_out()
                    
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 4:  # Scroll Up
                    camera.zoom_in()
                elif event.button == 5:  # Scroll Down
                    camera.zoom_out()

        # 2. Simulation Updates (skip if game over)
        keys = pygame.key.get_pressed()
        
        # Initialize empty maps in case loop is bypassed
        danger_map = {}
        reachability_map = {}
        
        if not game_over:
            # Handle player inputs (driving & direction queueing)
            player.handle_input(keys)
            player.update(dt)
            
            # Record player state in history buffer
            current_time += dt
            player_history.append((
                current_time,
                player.current_edge,
                player.progress,
                player.speed,
                player.get_heading(),
                player.pos.copy()
            ))
            
            # Prune player history older than 3.0 seconds
            while player_history and player_history[0][0] < current_time - 3.0:
                player_history.pop(0)
                
            # Progressive Cop Spawning
            if cop_escalation_mode:
                player_node = city.find_nearest_node(player.pos[0], player.pos[1])
                if len(cops) < 2 and current_time >= 15.0:
                    cop_edge = find_furthest_spawn_edge(city, player_node, excluded_nodes={c.current_edge[0] for c in cops})
                    cop = CopVehicle(city, cop_edge, player, pathfinder, predictor, speed=75.0)
                    cops.append(cop)
                if len(cops) < 3 and current_time >= 30.0:
                    cop_edge = find_furthest_spawn_edge(city, player_node, excluded_nodes={c.current_edge[0] for c in cops})
                    cop = CopVehicle(city, cop_edge, player, pathfinder, predictor, speed=75.0)
                    cops.append(cop)
                if len(cops) < 4 and current_time >= 45.0:
                    cop_edge = find_furthest_spawn_edge(city, player_node, excluded_nodes={c.current_edge[0] for c in cops})
                    cop = CopVehicle(city, cop_edge, player, pathfinder, predictor, speed=75.0)
                    cops.append(cop)

            # Update active roadblocks lifetimes
            for edge in list(roadblock_lifetimes.keys()):
                roadblock_lifetimes[edge] -= dt
                if roadblock_lifetimes[edge] <= 0.0:
                    city.remove_roadblock(*edge)
                    del roadblock_lifetimes[edge]
            
            # Update roadblock deployment cooldown
            roadblock_cooldown = max(0.0, roadblock_cooldown - dt)
                
            # Check survival timer win condition
            time_left -= dt
            if time_left <= 0.0:
                time_left = 0.0
                game_over = True
                game_result = "VICTORY"
                
            # Check collision/capture loss condition
            for cop in cops:
                dist = np.hypot(player.pos[0] - cop.pos[0], player.pos[1] - cop.pos[1])
                if dist < 22.0:
                    game_over = True
                    game_result = "BUSTED"
                    break
                    
            if not game_over:
                cops_positions = [cop.pos for cop in cops]
                
                # Get delayed player state for coordinate tracking and maps
                p_edge, p_progress, p_speed, p_heading, p_pos = cops[0].get_delayed_player_state(player_history, current_time)
                
                # Recompute threat and reachability influence maps
                influence_map_manager.update(p_pos, p_edge, p_heading, p_speed, cops_positions)
                danger_map = influence_map_manager.danger_map
                reachability_map = influence_map_manager.reachability_map
                player.danger_map = danger_map
                
                # Dynamic role allocation under coordinated strategy
                if active_strategy == "COORDINATED":
                    # Predict future intercept node
                    if prediction_model == "NAIVE":
                        pred_node = predictor.predict_naive(p_edge, p_progress, p_speed, p_heading, 2.0)
                    else:
                        pred_node = predictor.predict_road_aware(reachability_map, danger_map, fallback_node=p_edge[1])
                        
                    # Assign roles dynamically to the cop squad
                    role_allocator.assign_roles(cops, p_pos, pred_node, p_heading)
                    
                    # Coordinated Roadblock Deployment
                    if current_time >= 15.0 and roadblock_cooldown == 0.0 and len(city.active_roadblocks) < 2:
                        for cop in cops:
                            if getattr(cop, 'role', None) == "FLANKER" and getattr(cop, 'flanking_target_node', None) is not None:
                                flank_node = cop.flanking_target_node
                                if pred_node is not None and city.graph.has_edge(pred_node, flank_node):
                                    pos_flank = city.get_node_pos(flank_node)
                                    dist = np.hypot(cop.pos[0] - pos_flank[0], cop.pos[1] - pos_flank[1])
                                    if dist <= 30.0:
                                        sorted_player_edge = (min(player.current_edge), max(player.current_edge))
                                        sorted_target_edge = (min(pred_node, flank_node), max(pred_node, flank_node))
                                        
                                        # Verify player is not currently on the target edge, and it's not already blocked
                                        if sorted_player_edge != sorted_target_edge and sorted_target_edge not in city.active_roadblocks:
                                            city.add_roadblock(pred_node, flank_node)
                                            roadblock_lifetimes[sorted_target_edge] = 15.0
                                            roadblock_cooldown = 25.0
                                            break
                else:
                    # Clear roles if strategy is individual
                    for cop in cops:
                        cop.role = None
                        cop.flanking_target_node = None
                
                # Plan paths and update cop positions
                for cop in cops:
                    cop.plan_path(
                        player_history=player_history,
                        current_time=current_time,
                        strategy=active_strategy,
                        prediction_model=prediction_model,
                        cops_positions=cops_positions,
                        prediction_horizon=2.0,
                        danger_map=danger_map,
                        reachability_map=reachability_map
                    )
                    cop.update(dt)
        else:
            # Keep static snapshot of influence maps for game-over screen rendering
            danger_map = influence_map_manager.danger_map
            reachability_map = influence_map_manager.reachability_map
            
        # Center camera on player
        camera.update(player.pos, dt)

        # 3. Drawing
        renderer.draw(
            player, cops, active_strategy, prediction_model, time_left, 
            game_over, game_result, danger_map, reachability_map,
            cop_escalation_mode=cop_escalation_mode, roadblock_cooldown=roadblock_cooldown
        )
        pygame.display.flip()

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()