class Boid:
    def __init__(self, position, velocity):
        self.position = np.array(position, dtype=float)
        self.velocity = np.array(velocity, dtype=float)
        self.max_speed = 5.0
        self.max_force = 0.5

    def separation(self, neighbors, desired_distance=25):
        """Steer away from nearby boids"""
        steer = np.zeros(2)
        count = 0
        for other in neighbors:
            diff = self.position - other.position
            dist = np.linalg.norm(diff)
            if 0 < dist < desired_distance:
                steer += diff / dist
                count += 1
        if count > 0:
            steer /= count
            steer = self._limit(steer, self.max_force)
        return steer

    def alignment(self, neighbors):
        """Steer toward average heading of neighbors"""
        avg_velocity = np.zeros(2)
        count = 0
        for other in neighbors:
            avg_velocity += other.velocity
            count += 1
        if count > 0:
            avg_velocity /= count
            avg_velocity = self._normalize(avg_velocity) * self.max_speed
            steer = avg_velocity - self.velocity
            return self._limit(steer, self.max_force)
        return steer

    def cohesion(self, neighbors):
        """Steer toward average position of neighbors"""
        center = np.zeros(2)
        count = 0
        for other in neighbors:
            center += other.position
            count += 1
        if count > 0:
            center /= count
            desired = center - self.position
            desired = self._normalize(desired) * self.max_speed
            steer = desired - self.velocity
            return self._limit(steer, self.max_force)
        return np.zeros(2)

    def flock(self, boids):
        sep = self.separation(boids) * 1.5
        ali = self.alignment(boids) * 1.0
        coh = self.cohesion(boids) * 1.0
        self.velocity += sep + ali + coh
        self.velocity = self._limit(self.velocity, self.max_speed)
        self.position += self.velocity

    def _limit(self, vec, max_val):
        mag = np.linalg.norm(vec)
        if mag > max_val:
            vec = vec / mag * max_val
        return vec

    def _normalize(self, vec):
        mag = np.linalg.norm(vec)
        if mag > 0:
            return vec / mag
        return vec

class EnhancedBoid(Boid):
    def __init__(self, position, velocity):
        super().__init__(position, velocity)
        self.look_ahead = 50

    def avoid_obstacles(self, obstacles):
        """Steer away from obstacles"""
        steer = np.zeros(2)
        for obs_pos, obs_radius in obstacles:
            diff = self.position - obs_pos
            dist = np.linalg.norm(diff)
            if dist < obs_radius + self.look_ahead:
                # Project future position
                future_pos = self.position + self._normalize(self.velocity) * self.look_ahead
                future_dist = np.linalg.norm(future_pos - obs_pos)
                if future_dist < obs_radius:
                    desired = future_pos - obs_pos
                    desired = self._normalize(desired) * self.max_speed
                    steer = desired - self.velocity
                    steer = self._limit(steer, self.max_force * 2)
        return steer

    def stay_in_bounds(self, bounds_min, bounds_max):
        """Steer back if near boundary"""
        steer = np.zeros(2)
        margin = 50
        for i in range(2):
            if self.position[i] < bounds_min[i] + margin:
                steer[i] = self.max_force
            elif self.position[i] > bounds_max[i] - margin:
                steer[i] = -self.max_force
        return steer

from typing import List, Tuple

class AntColonyOptimizer:
    def __init__(self, n_drones, n_waypoints, alpha=1.0, beta=3.0,
                 rho=0.5, q=100):
        self.n_drones = n_drones
        self.n_waypoints = n_waypoints
        self.alpha = alpha  # Pheromone influence
        self.beta = beta    # Distance influence
        self.rho = rho      # Evaporation rate
        self.q = q          # Pheromone deposit factor

        # Initialize pheromone matrix
        self.pheromone = np.ones((n_waypoints, n_waypoints)) * 0.1

        # Distance matrix (will be set based on waypoints)
        self.distances = np.zeros((n_waypoints, n_waypoints))

    def set_distances(self, waypoints: np.ndarray):
        """Calculate distance matrix from waypoint positions"""
        for i in range(self.n_waypoints):
            for j in range(self.n_waypoints):
                self.distances[i][j] = np.linalg.norm(
                    waypoints[i] - waypoints[j]
                )

    def _probability(self, current, unvisited) -> np.ndarray:
        """Calculate transition probability to each unvisited waypoint"""
        probs = np.zeros(len(unvisited))
        for idx, next_wp in enumerate(unvisited):
            tau = self.pheromone[current][next_wp] ** self.alpha
            eta = (1.0 / self.distances[current][next_wp]) ** self.beta
            probs[idx] = tau * eta
        total = probs.sum()
        if total > 0:
            probs /= total
        return probs

    def _construct_tour(self, start) -> List[int]:
        """Construct a single drone's tour"""
        tour = [start]
        unvisited = list(range(self.n_waypoints))
        unvisited.remove(start)

        while unvisited:
            current = tour[-1]
            probs = self._probability(current, unvisited)
            next_wp = np.random.choice(unvisited, p=probs)
            tour.append(next_wp)
            unvisited.remove(next_wp)

        tour.append(start)  # Return to start
        return tour

    def _tour_distance(self, tour) -> float:
        """Calculate total distance of a tour"""
        dist = 0
        for i in range(len(tour) - 1):
            dist += self.distances[tour[i]][tour[i+1]]
        return dist

    def _update_pheromones(self, tours, distances):
        """Update pheromone levels"""
        # Evaporation
        self.pheromone *= (1 - self.rho)

        # Deposit new pheromones
        for tour, dist in zip(tours, distances):
            deposit = self.q / dist
            for i in range(len(tour) - 1):
                self.pheromone[tour[i]][tour[i+1]] += deposit
                self.pheromone[tour[i+1]][tour[i]] += deposit

    def optimize(self, iterations=100) -> Tuple[List[int], float]:
        """Run ACO optimization"""
        best_tour = None
        best_distance = float('inf')

        for iteration in range(iterations):
            tours = []
            distances = []

            for drone_id in range(self.n_drones):
                start = drone_id % self.n_waypoints
                tour = self._construct_tour(start)
                dist = self._tour_distance(tour)
                tours.append(tour)
                distances.append(dist)

                if dist < best_distance:
                    best_distance = dist
                    best_tour = tour

            self._update_pheromones(tours, distances)

            if iteration % 20 == 0:
                print(f"Iteration {iteration}: Best distance = {best_distance:.2f}")

        return best_tour, best_distance

from dataclasses import dataclass
from typing import Dict, List, Optional

@dataclass
class DroneMessage:
    sender_id: str
    receiver_id: Optional[str]  # None = broadcast
    msg_type: str
    payload: dict
    ttl: int = 5  # Time to live

class SwarmCommunicator:
    def __init__(self, drone_id: str):
        self.drone_id = drone_id
        self.neighbors: Dict[str, float] = {}  # id -> signal_strength
        self.message_queue: List[DroneMessage] = []
        self.routing_table: Dict[str, str] = {}  # dest -> next_hop

    async def broadcast(self, msg_type: str, payload: dict):
        """Broadcast message to all neighbors"""
        msg = DroneMessage(
            sender_id=self.drone_id,
            receiver_id=None,
            msg_type=msg_type,
            payload=payload
        )
        await self._send_to_neighbors(msg)

    async def send(self, receiver_id: str, msg_type: str, payload: dict):
        """Send message to specific drone via routing"""
        msg = DroneMessage(
            sender_id=self.drone_id,
            receiver_id=receiver_id,
            msg_type=msg_type,
            payload=payload
        )

        if receiver_id in self.routing_table:
            next_hop = self.routing_table[receiver_id]
            await self._send_to(next_hop, msg)
        else:
            await self._route_discovery(receiver_id, msg)

    async def _route_discovery(self, dest_id: str, msg: DroneMessage):
        """Discover route to destination"""
        discovery = DroneMessage(
            sender_id=self.drone_id,
            receiver_id=None,
            msg_type="ROUTE_REQUEST",
            payload={"dest": dest_id, "path": [self.drone_id]},
            ttl=10
        )
        await self._send_to_neighbors(discovery)

    async def handle_message(self, msg: DroneMessage):
        """Process incoming message"""
        if msg.msg_type == "ROUTE_REQUEST":
            await self._handle_route_request(msg)
        elif msg.msg_type == "POSITION_UPDATE":
            await self._handle_position_update(msg)
        elif msg.msg_type == "FORMATION_COMMAND":
            await self._handle_formation_command(msg)
        elif msg.receiver_id == self.drone_id or msg.receiver_id is None:
            await self._process_message(msg)

    async def _handle_position_update(self, msg: DroneMessage):
        """Update neighbor positions"""
        sender = msg.sender_id
        position = msg.payload["position"]
        self.neighbors[sender] = position

        # Rebroadcast to propagate positions
        if msg.ttl > 0:
            msg.ttl -= 1
            await self._send_to_neighbors(msg)

from dataclasses import dataclass
from typing import List

@dataclass
class FormationConfig:
    leader_id: str
    followers: List[str]
    offsets: List[np.ndarray]  # Relative positions

class FormationController:
    def __init__(self, config: FormationConfig):
        self.config = config
        self.leader_position = np.zeros(3)
        self.leader_velocity = np.zeros(3)
        self.follower_positions = {fid: np.zeros(3) for fid in config.followers}

    def update_leader(self, position: np.ndarray, velocity: np.ndarray):
        """Update leader state"""
        self.leader_position = position
        self.leader_velocity = velocity

    def get_follower_target(self, follower_id: str) -> np.ndarray:
        """Calculate target position for a follower"""
        idx = self.config.followers.index(follower_id)
        offset = self.config.offsets[idx]

        # Rotate offset based on leader heading
        heading = np.arctan2(self.leader_velocity[1], self.leader_velocity[0])
        rotation = np.array([
            [np.cos(heading), -np.sin(heading), 0],
            [np.sin(heading), np.cos(heading), 0],
            [0, 0, 1]
        ])

        target = self.leader_position + rotation @ offset
        return target

    def compute_control(self, follower_id: str,
                       current_pos: np.ndarray) -> np.ndarray:
        """Compute velocity command for follower"""
        target = self.get_follower_target(follower_id)
        error = target - current_pos

        # Proportional controller
        kp = 2.0
        velocity_command = error * kp

        # Limit speed
        max_speed = 5.0
        speed = np.linalg.norm(velocity_command)
        if speed > max_speed:
            velocity_command = velocity_command / speed * max_speed

        return velocity_command

class VFormationGenerator:
    def __init__(self, n_drones: int, spacing: float = 15.0,
                 angle: float = 30.0):
        self.n_drones = n_drones
        self.spacing = spacing
        self.angle = np.radians(angle)

    def generate_offsets(self) -> List[np.ndarray]:
        """Generate V-formation offset positions"""
        offsets = [np.array([0, 0, 0])]  # Leader at center

        for i in range(1, self.n_drones):
            side = 1 if i % 2 == 1 else -1
            row = (i + 1) // 2
            x = -row * self.spacing * np.cos(self.angle)
            y = side * row * self.spacing * np.sin(self.angle)
            offsets.append(np.array([x, y, 0]))

        return offsets

    def visualize(self):
        """Print ASCII visualization"""
        offsets = self.generate_offsets()
        print("V-Formation Layout:")
        print("  D0 (Leader)")
        for i in range(1, len(offsets)):
            side = "Left" if offsets[i][1] > 0 else "Right"
            print(f"    D{i} ({side})")

from enum import Enum
from typing import Dict

class TaskType(Enum):
    PATROL = "patrol"
    SURVEY = "survey"
    DELIVER = "deliver"
    RESCUE = "rescue"

class TaskAllocation:
    def __init__(self, drone_capabilities: Dict[str, list]):
        self.capabilities = drone_capabilities
        self.tasks: List[dict] = []

    def add_task(self, task_type: TaskType, priority: int,
                 position: np.ndarray):
        self.tasks.append({
            "type": task_type,
            "priority": priority,
            "position": position,
            "assigned": None
        })

    def allocate(self) -> Dict[str, dict]:
        """Allocate tasks to drones using auction-based method"""
        assignments = {}
        unassigned = [t for t in self.tasks if t["assigned"] is None]

        for task in sorted(unassigned, key=lambda t: -t["priority"]):
            best_drone = None
            best_cost = float('inf')

            for drone_id, caps in self.capabilities.items():
                if task["type"].value in caps:
                    cost = self._compute_cost(drone_id, task)
                    if cost < best_cost:
                        best_cost = cost
                        best_drone = drone_id

            if best_drone:
                task["assigned"] = best_drone
                assignments[best_drone] = task

        return assignments

    def _compute_cost(self, drone_id: str, task: dict) -> float:
        """Compute assignment cost (lower is better)"""
        # Base cost on distance and battery
        distance = np.linalg.norm(task["position"])
        priority_bonus = 1.0 / (task["priority"] + 1)
        return distance * priority_bonus

# search_swarm.py
from matplotlib.animation import FuncAnimation

class SearchArea:
    def __init__(self, width: float, height: float):
        self.width = width
        self.height = height
        self.coverage = np.zeros((100, 100))
        self.targets = []

    def add_target(self, x: float, y: float):
        self.targets.append((x, y))

    def update_coverage(self, position: np.ndarray, radius: float):
        """Mark area as covered"""
        ix = int(position[0] / self.width * 100)
        iy = int(position[1] / self.height * 100)
        r = int(radius / self.width * 100)

        for di in range(-r, r+1):
            for dj in range(-r, r+1):
                if 0 <= ix+di < 100 and 0 <= iy+dj < 100:
                    if di**2 + dj**2 <= r**2:
                        self.coverage[ix+di, iy+dj] = 1

    def check_detection(self, position: np.ndarray,
                       sensor_range: float) -> list:
        """Check if any targets are detected"""
        detected = []
        for tx, ty in self.targets:
            dist = np.sqrt((position[0]-tx)**2 + (position[1]-ty)**2)
            if dist < sensor_range:
                detected.append((tx, ty))
        return detected

class SearchSwarm:
    def __init__(self, n_drones: int, area: SearchArea):
        self.drones = []
        self.area = area
        self.detected_targets = []

        # Initialize drone positions in grid pattern
        n_cols = int(np.ceil(np.sqrt(n_drones)))
        for i in range(n_drones):
            row = i // n_cols
            col = i % n_cols
            x = (col + 1) * area.width / (n_cols + 1)
            y = (row + 1) * area.height / (n_cols + 1)
            self.drones.append({
                'pos': np.array([x, y]),
                'vel': np.zeros(2),
                'target': None
            })

    def assign_search_cells(self):
        """Assign unvisited cells to drones"""
        uncovered = np.argwhere(self.area.coverage < 0.5)

        for drone in self.drones:
            if drone['target'] is None and len(uncovered) > 0:
                # Find closest uncovered cell
                drone_pos_grid = drone['pos'] / np.array([
                    self.area.width, self.area.height]) * 100
                distances = np.linalg.norm(
                    uncovered - drone_pos_grid, axis=1)
                nearest = uncovered[np.argmin(distances)]
                drone['target'] = nearest / 100 * np.array([
                    self.area.width, self.area.height])

    def update(self, dt: float = 0.1):
        """Update swarm state"""
        self.assign_search_cells()

        for drone in self.drones:
            if drone['target'] is not None:
                # Move toward target
                diff = drone['target'] - drone['pos']
                dist = np.linalg.norm(diff)

                if dist < 1.0:
                    drone['target'] = None
                else:
                    speed = 2.0
                    drone['vel'] = diff / dist * speed
                    drone['pos'] += drone['vel'] * dt

            # Update coverage
            self.area.update_coverage(drone['pos'], radius=5.0)

            # Check for targets
            detected = self.area.check_detection(drone['pos'], 3.0)
            self.detected_targets.extend(detected)

    def get_coverage_percent(self) -> float:
        return np.mean(self.area.coverage) * 100

    def animate(self, frames=200):
        """Visualize the search"""
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

        def update_frame(frame):
            self.update()

            ax1.clear()
            ax2.clear()

            # Coverage heatmap
            ax1.imshow(self.area.coverage.T, origin='lower',
                      cmap='YlGn', extent=[0, self.area.width,
                                          0, self.area.height])
            ax1.set_title(f'Coverage: {self.get_coverage_percent():.1f}%')

            # Drone positions
            for drone in self.drones:
                ax2.plot(*drone['pos'], 'o', color='#0ea5e9', markersize=8)
                if drone['target'] is not None:
                    ax2.plot(*drone['target'], 'x', color='red')

            for tx, ty in self.area.targets:
                ax2.plot(tx, ty, '*', color='#f59e0b', markersize=12)

            ax2.set_xlim(0, self.area.width)
            ax2.set_ylim(0, self.area.height)
            ax2.set_title('Search Swarm')

        anim = FuncAnimation(fig, update_frame, frames=frames,
                           interval=50)
        plt.tight_layout()
        plt.show()

# Usage
area = SearchArea(100, 100)
area.add_target(20, 30)
area.add_target(70, 80)
area.add_target(45, 60)

swarm = SearchSwarm(n_drones=6, area=area)
swarm.animate(frames=300)