import math
import os
from io import BytesIO
from time import sleep
from PIL import Image
import cv2
import numpy as np
from position import Position
from vision_chunk import VisionChunk
from yandex_map import YandexMap
import constants
import utility

class Simulator:
    def __init__(self, yandex_map: YandexMap = None):
        self.yandex_map = yandex_map
        # Используем новый конструктор с yaw, pitch, roll
        self.pos = Position(x=0, y=0, z=1, yaw=0, pitch=0, roll=0)

        os.makedirs('./images', exist_ok=True)

    def set_pitch(self, pitch_deg: float):
        """Установить тангаж камеры (градусы, -10 до 10)"""
        clamped_pitch = max(-10, min(10, pitch_deg))
        self.pos.pitch = math.radians(clamped_pitch)

    def set_roll(self, roll_deg: float):
        """Установить крен камеры (градусы, -10 до 10)"""
        clamped_roll = max(-10, min(10, roll_deg))
        self.pos.roll = math.radians(clamped_roll)

    def _apply_perspective_transform(self, image: Image.Image) -> Image.Image:
        """
        Применяет перспективную трансформацию к изображению.
        Возвращает квадратное изображение 700x700.
        """
        img_array = np.array(image)
        print(img_array.shape)
        h, w, _ = img_array.shape

        # Применяем трансформацию
        pos = self.pos.copy()
        pos.x = 0
        pos.y = 0
        K = utility.calc_camera_matrix(w, h)
        K = constants.K
        img_array = img_array[:constants.CHUNK_WIDTH, :constants.CHUNK_WIDTH]
        transformed = cv2.warpPerspective(img_array, pos.get_homography_matrix(K), (constants.CHUNK_WIDTH, constants.CHUNK_WIDTH))

        return Image.fromarray(transformed)

    def update_trajectory(self, dx: float, dy: float):
        """Обновляет координаты дрона"""
        self.pos.x += dx * self.pos.z
        self.pos.y += dy * self.pos.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

        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()

        # Обновляем yaw в объекте Position
        self.pos.yaw += dangle
        velocity = max(velocity, 10)

        # Вычисляем смещение на основе текущего yaw
        dx = math.cos(math.pi / 2 + self.pos.yaw) * velocity / self.pos.z
        dy = math.sin(math.pi / 2 + self.pos.yaw) * velocity / self.pos.z

        self.update_trajectory(dx, dy)

        action.move_by_offset(-dx, dy)
        action.release()
        action.perform()

    def set_zoom(self, zoom_level: float):
        """Программное изменение масштаба"""
        self.pos.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)

    def get_orientation_info(self) -> dict:
        """Получить полную информацию об ориентации дрона"""
        return {
            'position': (self.pos.x, self.pos.y, self.pos.z),
            'yaw_deg': math.degrees(self.pos.yaw),
            'pitch_deg': math.degrees(self.pos.pitch),
            'roll_deg': math.degrees(self.pos.roll),
            'rotation_matrix': self.pos.get_rotation_matrix()
        }