import math
import os
from io import BytesIO
from time import sleep
from PIL import Image
import cv2
import numpy as np
from geolocation import Geolocation
from vision_chunk import VisionChunk
from yandex_map import YandexMap


class Simulator:
    def __init__(self, yandex_map: YandexMap = None):
        self.yandex_map = yandex_map
        self.geo = Geolocation(0, 0, 1, 0)
        
        # Параметры камеры (в градусах)
        self.pitch = 0.0  # тангаж (-10 до 10)
        self.roll = 0.0   # крен (-10 до 10)
        
        os.makedirs('./images', exist_ok=True)
    
    def set_pitch(self, pitch: float):
        """Установить тангаж камеры (градусы, -10 до 10)"""
        self.pitch = max(-10, min(10, pitch))
    
    def set_roll(self, roll: float):
        """Установить крен камеры (градусы, -10 до 10)"""
        self.roll = max(-10, min(10, roll))
    
    def _calculate_camera_angles(self, velocity: float, dangle: float):
        """Автоматический расчёт тангажа и крена на основе скорости и поворота"""
        # Тангаж: чем больше скорость, тем больше тангаж (до 10°)
        self.pitch = min(10, velocity / 10)
        
        # Крен: чем больше угол поворота, тем больше крен
        self.roll = max(-10, min(10, math.degrees(dangle) * 2))
    
    def _get_perspective_points(self, image_width: int, image_height: int, 
                                yaw_deg: float) -> tuple:
        """
        Вычисляет 4 точки для перспективной трансформации.
        Учитывает тангаж, крен, поворот и масштаб.
        """
        center_x = image_width / 2
        center_y = image_height / 2
        
        # Базовый размер области захвата (квадрат)
        base_size = min(image_width, image_height) * 0.7
        half_size = base_size / 2
        
        # Исходные углы квадрата (до применения перспективы)
        corners = np.float32([
            [center_x - half_size, center_y - half_size],  # top-left
            [center_x + half_size, center_y - half_size],  # top-right
            [center_x + half_size, center_y + half_size],  # bottom-right
            [center_x - half_size, center_y + half_size],  # bottom-left
        ])
        
        # Применяем смещения для имитации тангажа (pitch)
        # Положительный тангаж - дрон наклонён вперёд, камера смотрит дальше
        pitch_offset = self.pitch * 3  # коэффициент для видимого эффекта
        corners[0][1] -= pitch_offset  # top-left y
        corners[1][1] -= pitch_offset  # top-right y
        corners[2][1] += pitch_offset  # bottom-right y
        corners[3][1] += pitch_offset  # bottom-left y
        
        # Применяем смещения для имитации крена (roll)
        # Положительный крен - дрон наклонён вправо
        roll_offset = self.roll * 3
        corners[0][0] += roll_offset  # top-left x
        corners[1][0] -= roll_offset  # top-right x
        corners[2][0] -= roll_offset  # bottom-right x
        corners[3][0] += roll_offset  # bottom-left x
        
        # Поворот вокруг центра (yaw)
        angle_rad = math.radians(yaw_deg)
        cos_a = math.cos(angle_rad)
        sin_a = math.sin(angle_rad)
        
        rotated_corners = []
        for corner in corners:
            # Смещаем к началу координат
            x = corner[0] - center_x
            y = corner[1] - center_y
            
            # Поворачиваем
            new_x = x * cos_a - y * sin_a
            new_y = x * sin_a + y * cos_a
            
            # Возвращаем обратно
            rotated_corners.append([new_x + center_x, new_y + center_y])
        
        return np.float32(rotated_corners)
    
    def _apply_perspective_transform(self, image: Image.Image) -> Image.Image:
        """
        Применяет перспективную трансформацию к изображению.
        Возвращает квадратное изображение 700x700.
        """
        img_array = np.array(image)
        h, w = img_array.shape[:2]
        
        # Получаем исходные точки с учётом перспективы
        yaw_deg = -math.degrees(self.geo.angle)
        src_points = self._get_perspective_points(w, h, yaw_deg)
        
        # Целевые точки - квадрат 700x700
        dst_points = np.float32([
            [0, 0],
            [700, 0],
            [700, 700],
            [0, 700]
        ])
        
        # Вычисляем матрицу трансформации
        matrix = cv2.getPerspectiveTransform(src_points, dst_points)
        
        # Применяем трансформацию
        transformed = cv2.warpPerspective(img_array, matrix, (700, 700))
        
        return Image.fromarray(transformed)
    
    def update_trajectory(self, dx: float, dy: float):
        """Обновляет координаты дрона"""
        self.geo.x += dx * self.geo.z
        self.geo.y += dy * self.geo.z
    
    def handle(self, dangle: float, velocity: float = 50) -> None:
        """
        Управление движением дрона.
        dangle - изменение угла курса (радианы)
        velocity - скорость движения
        """
        from selenium.webdriver.common.by import By
        from selenium.webdriver.common.action_chains import ActionChains
        
        # Автоматический расчёт углов камеры
        self._calculate_camera_angles(velocity, dangle)
        
        html = self.yandex_map.driver.find_element(By.TAG_NAME, 'html')
        action = ActionChains(self.yandex_map.driver)
        action.move_to_element_with_offset(html, 200, 200)
        action.click_and_hold()
        
        self.geo.angle += dangle
        velocity = max(velocity, 10)
        
        dx = math.cos(math.pi / 2 + self.geo.angle) * velocity / self.geo.z
        dy = math.sin(math.pi / 2 + self.geo.angle) * velocity / self.geo.z
        
        self.update_trajectory(dx, dy)
        
        action.move_by_offset(-dx, dy)
        action.release()
        action.perform()
    
    def set_zoom(self, zoom_level: float):
        """Программное изменение масштаба"""
        self.geo.z = zoom_level
    
    def get_chunk(self) -> VisionChunk:
        """Получить текущий снимок с камеры дрона"""
        png = self.yandex_map.driver.get_screenshot_as_png()
        im = Image.open(BytesIO(png))
        
        # Применяем перспективную трансформацию
        transformed_im = self._apply_perspective_transform(im)
        
        return VisionChunk(transformed_im)