autopilot/simulator.py

import math
from io import BytesIO
from time import sleep

from PIL import Image
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.action_chains import ActionChains

from autopilot import Pilot
from visualization import VisualizationManager, SimMode

import os

class Simulator:
    driver: webdriver.Chrome
    mode: SimMode

    operatorPilot: Pilot
    autonomePilot: Pilot

    angle: float
    
    # Менеджер визуализации
    viz_manager: VisualizationManager
    current_x: float
    current_y: float

    def __init__(self, operatorPilot: Pilot, autonomePilot: Pilot, viz_manager: VisualizationManager = None):
        self.mode = SimMode.OPERATOR
        self.operatorPilot = operatorPilot
        self.autonomePilot = autonomePilot

        # Инициализация переменных для траектории
        self.current_x = 0.0
        self.current_y = 0.0
        
        # Создаем менеджер визуализации
        self.viz_manager = viz_manager
        
        # Передаем менеджер визуализации в автопилот, если он поддерживает это
        if hasattr(self.autonomePilot, 'viz_manager'):
            self.autonomePilot.viz_manager = self.viz_manager
        
        # Создаем папку для изображений, если её нет
        os.makedirs('./images', exist_ok=True)

        options = webdriver.ChromeOptions()
        # options.add_experimental_option("detach", True)
        self.driver = webdriver.Chrome(options)
        self.driver.get("https://yandex.ru/maps/43/kazan/?ll=49.103814%2C55.794258&z=14")
        sleep(2)

        action = ActionChains(self.driver)

        # Закрытие левой панели
        action.click(self.driver.find_element(By.CLASS_NAME, 'sidebar-toggle-button'))
        action.move_to_element_with_offset(self.driver.find_element(By.XPATH, "//div[@class='rounded-controls']/div[@class='rounded-controls__child'][5]//button"), 5, 5)
        action.perform()

        # Режим спутника
        action.click(self.driver.find_element(By.CLASS_NAME, '_key_satellite'))
        action.perform()

        self.angle = 0

    def rotate_image_like_drone(self, image: Image.Image, angle: float) -> Image.Image:
        """
        Поворачивает картинку как будто съемка ведется с летящего дрона.
        Выделяет концентрический квадрат, поворачивает его и извлекает результат.
        """
        # Получаем размеры изображения
        width, height = image.size
        square_size = min(width, height)
        cropped_image = image.crop((0, 0, square_size, square_size))
        cropped_image = cropped_image.rotate(angle / math.pi * 180, expand=True)
        
        # Определяем размер концентрического квадрата (80% от минимальной стороны)
        local_square_size = int(square_size / 2 ** 0.5)
        
        # Вычисляем координаты для центрирования квадрата
        left = (cropped_image.width - local_square_size) // 2
        top = (cropped_image.height - local_square_size) // 2
        right = left + local_square_size
        bottom = top + local_square_size
        
        # Вырезаем концентрический квадрат
        final_image = cropped_image.crop((left, top, right, bottom))

        return final_image

    def update_trajectory(self, dx: float, dy: float):
        """
        Обновляет траекторию полета беспилотника
        """
        # Обновляем текущие координаты
        self.current_x += dx
        self.current_y += dy

    def update_map(self):
        """
        Обновляет карту траектории полета
        """
        self.viz_manager.update_global_map(self.current_x, self.current_y, self.mode)

    def loop(self):

        html = self.driver.find_element(By.TAG_NAME, 'html')
        action = ActionChains(self.driver)

        # Добавляем начальную точку в траекторию
        self.update_trajectory(0, 0)
        self.viz_manager.update_global_map(self.current_x, self.current_y, self.mode)

        for i in range(1000):
            signal = None
            if self.mode == SimMode.OPERATOR:
                signal = self.operatorPilot.act()
                if signal is None:
                    self.mode = SimMode.AUTONOME
                    print("Режим возвращения домой!")

            if self.mode == SimMode.AUTONOME:
                signal = self.autonomePilot.act()

            if signal is None:
                break

            dangle, velocity = signal

            # Сдвиг камеры
            action = ActionChains(self.driver)
            action.move_to_element_with_offset(html, 200, 200)
            action.click_and_hold()

            self.angle += dangle
            dx = math.cos(self.angle) * velocity
            dy = math.sin(self.angle) * velocity
            action.move_by_offset(-dx, dy)
            action.release()
            action.perform()

            # Обновляем траекторию
            self.update_trajectory(dx, dy)

            # Загрузка скриншота
            png = self.driver.get_screenshot_as_png()
            im = Image.open(BytesIO(png))
            im = im.crop([0, 80, im.width-80, im.height-60])
            
            # Применяем поворот как будто съемка с дрона
            rotated_im = self.rotate_image_like_drone(im, math.pi / 2 - self.angle)
            
            # Передаем изображение в AutoPilot для анализа
            self.autonomePilot.handle(rotated_im)
            
            rotated_im.save(f"./images/{i}.png")
            
            # Обновляем визуализацию каждые несколько итераций для производительности
            if i % 5 == 0:
                self.update_map()
                self.viz_manager.update_display()
            
            self.viz_manager.pause(0.25)

        # Финальное обновление карты
        self.update_map()
        self.viz_manager.update_display()
        print("last position: ", self.driver.current_url)
        print(f"Финальная позиция: ({self.current_x:.2f}, {self.current_y:.2f})")
        
        # Показываем карту до закрытия, но не поднимаем на передний план
        self.viz_manager.show_final()
        print("Симуляция завершена. Окно визуализации остается открытым для анализа.")