import numpy as np
from PIL import Image
import cv2
from autopilot import AutoPilot
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()

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