174 lines
6.2 KiB
Python
174 lines
6.2 KiB
Python
import pygame
|
|
import random
|
|
import math
|
|
from collections import deque
|
|
|
|
# --- Simulation Classes ---
|
|
|
|
class Task:
|
|
__slots__ = ('source', 'destination', 'item', 'amount')
|
|
def __init__(self, source, destination, item, amount):
|
|
self.source = source
|
|
self.destination = destination
|
|
self.item = item
|
|
self.amount = amount
|
|
|
|
class Storage:
|
|
def __init__(self, name, pos):
|
|
self.name = name
|
|
self.position = pos # (x, y) position on screen
|
|
self.inventory = {"item": 1000} # starting inventory
|
|
|
|
def add_item(self, item, amount):
|
|
self.inventory[item] = self.inventory.get(item, 0) + amount
|
|
|
|
def remove_item(self, item, amount):
|
|
if self.inventory.get(item, 0) >= amount:
|
|
self.inventory[item] -= amount
|
|
return True
|
|
return False
|
|
|
|
def request_transfer(self, destination, item, amount, task_queue):
|
|
if self.inventory.get(item, 0) >= amount:
|
|
# Remove items immediately, simulating reservation.
|
|
self.remove_item(item, amount)
|
|
task = Task(self, destination, item, amount)
|
|
task_queue.append(task)
|
|
|
|
class Drone:
|
|
__slots__ = ('id', 'x', 'y', 'task', 'speed')
|
|
def __init__(self, id, pos):
|
|
self.id = id
|
|
self.x, self.y = pos
|
|
self.task = None
|
|
self.speed = 2.0 # pixels per frame
|
|
|
|
def assign_task(self, task):
|
|
self.task = task
|
|
# Set drone's position to the source storage position.
|
|
self.x, self.y = task.source.position
|
|
|
|
def update(self):
|
|
if self.task:
|
|
dest_x, dest_y = self.task.destination.position
|
|
dx = dest_x - self.x
|
|
dy = dest_y - self.y
|
|
dist = math.hypot(dx, dy)
|
|
if dist < self.speed:
|
|
# Arrived: set position exactly to destination and deliver item.
|
|
self.x, self.y = dest_x, dest_y
|
|
self.task.destination.add_item(self.task.item, self.task.amount)
|
|
self.task = None
|
|
else:
|
|
# Move a step toward the destination.
|
|
self.x += self.speed * dx / dist
|
|
self.y += self.speed * dy / dist
|
|
|
|
# --- Pygame Simulation ---
|
|
|
|
def main():
|
|
pygame.init()
|
|
screen_width, screen_height = 800, 600
|
|
screen = pygame.display.set_mode((screen_width, screen_height))
|
|
pygame.display.set_caption("Drone Simulation")
|
|
clock = pygame.time.Clock()
|
|
font = pygame.font.SysFont("Arial", 16)
|
|
|
|
# Global simulation state
|
|
task_queue = deque()
|
|
storages = []
|
|
drones = []
|
|
|
|
# Unique id counters
|
|
drone_id_counter = 0
|
|
storage_id_counter = 0
|
|
|
|
# Create some initial storages at random positions
|
|
for _ in range(5):
|
|
pos = (random.randint(50, screen_width - 50), random.randint(50, screen_height - 50))
|
|
storages.append(Storage(f"Storage-{storage_id_counter}", pos))
|
|
storage_id_counter += 1
|
|
|
|
# Create some initial drones at random positions
|
|
for _ in range(20):
|
|
pos = (random.randint(0, screen_width), random.randint(0, screen_height))
|
|
drones.append(Drone(drone_id_counter, pos))
|
|
drone_id_counter += 1
|
|
|
|
running = True
|
|
while running:
|
|
# --- 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_t:
|
|
# Request a random transfer if at least two storages exist.
|
|
if len(storages) >= 2:
|
|
src, dest = random.sample(storages, 2)
|
|
amount = random.randint(1, 10)
|
|
src.request_transfer(dest, "item", amount, task_queue)
|
|
elif event.key == pygame.K_d:
|
|
# Add a new drone at a random position.
|
|
pos = (random.randint(0, screen_width), random.randint(0, screen_height))
|
|
drones.append(Drone(drone_id_counter, pos))
|
|
drone_id_counter += 1
|
|
elif event.key == pygame.K_s:
|
|
# Add a new storage at a random position.
|
|
pos = (random.randint(50, screen_width - 50), random.randint(50, screen_height - 50))
|
|
storages.append(Storage(f"Storage-{storage_id_counter}", pos))
|
|
storage_id_counter += 1
|
|
elif event.key == pygame.K_i:
|
|
# Add 50 items to each storage.
|
|
for s in storages:
|
|
s.add_item("item", 50)
|
|
|
|
# --- Update Simulation ---
|
|
for drone in drones:
|
|
if drone.task is None and task_queue:
|
|
# Assign a task from the queue to the idle drone.
|
|
task = task_queue.popleft()
|
|
drone.assign_task(task)
|
|
drone.update()
|
|
|
|
# --- Drawing ---
|
|
screen.fill((30, 30, 30)) # dark background
|
|
|
|
# Draw storages (blue circles) with name and inventory.
|
|
for s in storages:
|
|
pygame.draw.circle(screen, (0, 100, 255), s.position, 20)
|
|
info = font.render(f"{s.name}: {s.inventory.get('item', 0)}", True, (255, 255, 255))
|
|
screen.blit(info, (s.position[0] - 30, s.position[1] - 40))
|
|
|
|
# Draw drones.
|
|
# Idle drones are drawn in green; drones with tasks (in transit) are red.
|
|
for d in drones:
|
|
color = (255, 0, 0) if d.task else (0, 255, 0)
|
|
pygame.draw.circle(screen, color, (int(d.x), int(d.y)), 5)
|
|
|
|
# Debug info: show task queue length, number of drones, storages.
|
|
debug_info = font.render(
|
|
f"Task Queue: {len(task_queue)} | Drones: {len(drones)} | Storages: {len(storages)}",
|
|
True, (255, 255, 0))
|
|
screen.blit(debug_info, (10, 10))
|
|
|
|
# Display instructions.
|
|
instructions = [
|
|
"Controls:",
|
|
"T - Request Transfer",
|
|
"D - Add Drone",
|
|
"S - Add Storage",
|
|
"I - Add Items to all Storages"
|
|
]
|
|
for i, line in enumerate(instructions):
|
|
instr_text = font.render(line, True, (200, 200, 200))
|
|
screen.blit(instr_text, (10, 30 + i * 20))
|
|
|
|
pygame.display.flip()
|
|
clock.tick(60)
|
|
|
|
pygame.quit()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|