#!/usr/bin/env python3
"""
Тестовый скрипт для проверки системы отслеживания траектории беспилотника
"""

import numpy as np
from PIL import Image
import cv2
from autopilot import AutoPilot, RandomPilot
from simulator import Simulator, SimMode
import matplotlib.pyplot as plt
import time

def create_test_image(width=640, height=480, pattern_type="random"):
    """Создает тестовое изображение с различными паттернами"""
    if pattern_type == "random":
        # Создаем изображение с случайными точками
        img = np.random.randint(0, 255, (height, width, 3), dtype=np.uint8)
        # Добавляем несколько четких объектов
        cv2.rectangle(img, (100, 100), (200, 200), (255, 0, 0), -1)
        cv2.circle(img, (400, 300), 50, (0, 255, 0), -1)
        cv2.line(img, (50, 400), (550, 400), (0, 0, 255), 5)
    elif pattern_type == "grid":
        # Создаем сетку
        img = np.zeros((height, width, 3), dtype=np.uint8)
        for i in range(0, width, 50):
            cv2.line(img, (i, 0), (i, height), (255, 255, 255), 1)
        for i in range(0, height, 50):
            cv2.line(img, (0, i), (width, i), (255, 255, 255), 1)
    elif pattern_type == "circles":
        # Создаем изображение с кругами
        img = np.zeros((height, width, 3), dtype=np.uint8)
        for i in range(5):
            x = np.random.randint(50, width-50)
            y = np.random.randint(50, height-50)
            radius = np.random.randint(20, 80)
            color = (np.random.randint(0, 255), np.random.randint(0, 255), np.random.randint(0, 255))
            cv2.circle(img, (x, y), radius, color, -1)
    
    return Image.fromarray(img)

def simulate_drone_movement(base_image, dx, dy, rotation=0):
    """Симулирует движение БПЛА, создавая смещенное изображение"""
    img_array = np.array(base_image)
    height, width = img_array.shape[:2]
    
    # Создаем матрицу трансформации
    cos_rot = np.cos(rotation)
    sin_rot = np.sin(rotation)
    
    # Матрица поворота и смещения
    M = np.float32([
        [cos_rot, -sin_rot, dx],
        [sin_rot, cos_rot, dy]
    ])
    
    # Применяем трансформацию
    transformed = cv2.warpAffine(img_array, M, (width, height), 
                                borderMode=cv2.BORDER_REPLICATE)
    
    return Image.fromarray(transformed)

def test_drone_tracking():
    """Тестирует систему отслеживания БПЛА"""
    print("Запуск тестирования системы отслеживания БПЛА...")
    
    # Создаем автопилот
    autopilot = AutoPilot(initial_x=0.0, initial_y=0.0, initial_altitude=10.0)
    
    # Создаем базовое изображение
    base_image = create_test_image(pattern_type="circles")
    
    # Траектория движения (dx, dy, rotation)
    trajectory = [
        (10, 0, 0),      # Движение вперед
        (10, 5, 0.1),    # Движение вперед с небольшим поворотом
        (0, 10, 0),      # Движение вправо
        (-5, 5, -0.1),   # Движение назад-вправо с поворотом
        (5, -5, 0.05),   # Движение вперед-влево
        (0, -10, 0),     # Движение влево
        (-10, 0, 0),     # Движение назад
    ]
    
    # Массивы для хранения истории позиций
    x_history = []
    y_history = []
    angle_history = []
    
    # Обрабатываем каждое изображение в траектории
    for i, (dx, dy, rotation) in enumerate(trajectory):
        print(f"\n--- Кадр {i+1} ---")
        
        # Создаем смещенное изображение
        current_image = simulate_drone_movement(base_image, dx, dy, rotation)
        
        # Обрабатываем изображение
        autopilot.handle(current_image)
        
        # Получаем состояние БПЛА
        state = autopilot.get_drone_state()
        x_history.append(state['x'])
        y_history.append(state['y'])
        angle_history.append(state['angle_degrees'])
        
        # Обновляем базовое изображение для следующего кадра
        base_image = current_image
        
        # Небольшая пауза для визуализации
        time.sleep(0.5)
    
    # Закрываем окно OpenCV
    cv2.destroyAllWindows()
    
    # Визуализируем траекторию
    plot_trajectory(x_history, y_history, angle_history)

def plot_trajectory(x_history, y_history, angle_history):
    """Визуализирует траекторию движения БПЛА"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    # График траектории
    ax1.plot(x_history, y_history, 'b-o', linewidth=2, markersize=8)
    ax1.plot(x_history[0], y_history[0], 'go', markersize=12, label='Начало')
    ax1.plot(x_history[-1], y_history[-1], 'ro', markersize=12, label='Конец')
    ax1.set_xlabel('X координата (м)')
    ax1.set_ylabel('Y координата (м)')
    ax1.set_title('Траектория движения БПЛА')
    ax1.grid(True)
    ax1.legend()
    ax1.axis('equal')
    
    # График угла поворота
    ax2.plot(range(len(angle_history)), angle_history, 'r-o', linewidth=2, markersize=8)
    ax2.set_xlabel('Номер кадра')
    ax2.set_ylabel('Угол поворота (градусы)')
    ax2.set_title('Изменение угла поворота БПЛА')
    ax2.grid(True)
    
    plt.tight_layout()
    plt.show()

def test_single_frame():
    """Тестирует обработку одного кадра"""
    print("Тестирование обработки одного кадра...")
    
    autopilot = AutoPilot()
    
    # Создаем два похожих изображения
    img1 = create_test_image(pattern_type="circles")
    img2 = simulate_drone_movement(img1, 15, 10, 0.05)
    
    # Обрабатываем первое изображение
    autopilot.handle(img1)
    
    # Обрабатываем второе изображение
    autopilot.handle(img2)
    
    # Получаем состояние
    state = autopilot.get_drone_state()
    print(f"Позиция БПЛА: ({state['x']:.2f}, {state['y']:.2f})")
    print(f"Угол БПЛА: {state['angle_degrees']:.1f}°")
    
    cv2.destroyAllWindows()

def test_trajectory_tracking():
    """Тестирует систему отслеживания траектории"""
    print("Запуск теста системы отслеживания траектории...")
    
    # Создаем пилотов
    operator_pilot = RandomPilot(velocity=1.0)
    autonome_pilot = AutoPilot(initial_x=0.0, initial_y=0.0)
    
    # Создаем симулятор
    simulator = Simulator(operator_pilot, autonome_pilot)
    
    try:
        # Запускаем симуляцию
        simulator.loop()
        
        print("Симуляция завершена!")
        print(f"Всего точек траектории: {len(simulator.trajectory_x)}")
        print(f"Режим оператора: {sum(1 for mode in simulator.trajectory_modes if mode == SimMode.OPERATOR)} точек")
        print(f"Автономный режим: {sum(1 for mode in simulator.trajectory_modes if mode == SimMode.AUTONOME)} точек")
        
    except Exception as e:
        print(f"Ошибка во время симуляции: {e}")
    finally:
        # Закрываем браузер
        simulator.driver.quit()

if __name__ == "__main__":
    print("Выберите тест:")
    print("1. Тест одного кадра")
    print("2. Тест полной траектории")
    
    choice = input("Введите номер теста (1 или 2): ").strip()
    
    if choice == "1":
        test_single_frame()
    elif choice == "2":
        test_trajectory_tracking()
    else:
        print("Неверный выбор. Запускаю тест одного кадра...")
        test_single_frame() 