from shapely.geometry import Polygon,Point
from typing import Tuple, List, Dict, Set
from scipy.spatial import cKDTree
import sys
import numpy as np
from agent_setup import AgentSetup


class SpawnManager:
    def __init__(self,spawn_area:Polygon,min_spacing:float=0.55):
        """
        self.spawn_area:     geometry where agents may spawn
        self.min_spacing:    minimum spacing for spawn_points 
        self.spawn_coords:   all spawn points available
        self.filled_coords:  spawn points currently filled by agents
        self.spawned_agents: all agents already spawned in sim
        self.agent_pos:      connects agent_id to spawn_index
        self.rng:            random number generator object
        """
        
        self.spawn_area = spawn_area
        self.min_spacing = min_spacing
        self.spawn_coords: np.ndarray = np.array([])
        self.filled_coords: np.ndarray = np.array([])
        self.spawned_agents:Set[int] = set()
        self.agent_pos:Dict[int,int] = {}
        self.rng = np.random.default_rng()

    def generate_coords(self,max_samples:int=1000)->None:
        buffered_spawn = self.spawn_area.buffer(-self.min_spacing)
        min_x,min_y,max_x,max_y = buffered_spawn.bounds
        points = []
        while len(points) ==0:
            x = self.rng.uniform(min_x,max_x)
            y = self.rng.uniform(min_y,max_y)
            if buffered_spawn.contains(Point(x,y)):
                points.append([x,y])
        for _ in range(max_samples):
            idx = self.rng.integers(0,len(points))
            base = points[idx]
            for _ in range(25):
                angle = self.rng.uniform(0,2*np.pi)
                radius = self.rng.uniform(self.min_spacing,2*self.min_spacing)
                x = base[0]+radius*np.cos(angle)
                y = base[1]+radius*np.sin(angle)
                if not buffered_spawn.contains(Point(x,y)):
                    continue
                if len(points)>0:
                    tree = cKDTree(points)
                    distance,_ = tree.query([[x,y]],k=1)
                    if distance[0]<self.min_spacing:
                        continue
                points.append([x,y])
                break
        self.spawn_coords = np.array(points)
        self.filled_coords = np.zeros(len(points),dtype=bool)
        return self.spawn_coords
    
    def get_coords(self)->Tuple[float,float]|None:
        free_idx = np.where(~self.filled_coords)[0]
        if len(free_idx) == 0:
            return None
        idx = self.rng.choice(free_idx)
        return tuple(self.spawn_coords[idx])
    

    def spawn_agent(self,all_agents:List[AgentSetup])->AgentSetup|None:
        if len(self.spawned_agents) >= len(all_agents):
            return None
        spawn_point = self.get_coords()
        if not spawn_point:
            return None
        
        free_agents = [agent for agent in all_agents \
            if agent.id not in self.spawned_agents
            ]
        if not free_agents:
            return None
        agent = self.rng.choice(free_agents)
        self.spawned_agents.add(agent.id)
        distances = np.linalg.norm(self.spawn_coords-spawn_point,axis=1)
        spawn_idx = np.argmin(distances)
        self.filled_coords[spawn_idx] = True
        self.agent_pos[agent.id] = spawn_idx
        return agent

    def unfill_coords(self,agent_id:int)->None:
        if agent_id in self.agent_pos:
            spawn_idx = self.agent_pos[agent_id]
            self.filled_coords[spawn_idx] = False
            del self.agent_pos[agent_id]
            self.spawned_agents.discard(agent_id)

    def get_agent_pos(self,agent_id:int)->Tuple[float,float]:
        if agent_id in self.agent_pos:
            spawn_idx = self.agent_pos[agent_id]
            return tuple(self.spawn_coords[spawn_idx])
        return (0,0)
    
    def check_spawn_complete(self,all_agents:List[AgentSetup])->bool:
        return (len(self.spawned_agents)<len(all_agents) and\
             np.sum(~self.filled_coords) > 0)