#!/usr/bin/env python3
"""
Модуль для управления общим окном визуализации
"""

import matplotlib.axes
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
from enum import Enum
import cv2
from PIL import Image
import matplotlib

# Настройки matplotlib
matplotlib.use('TkAgg')
plt.rcParams['figure.raise_window'] = False

class SimMode(Enum):
    OPERATOR = 1
    AUTONOME = 2

class VisualizationManager:
    """
    Менеджер для управления общим окном визуализации
    """
    
    def __init__(self, window_title="Drone Autopilot Visualization"):
        self.window_title = window_title
        self.fig = None
        self.ax_error_plot = None  # График погрешности позиции
        self.ax_global_map = None
        self.ax_detection = None
        self.ax_matches = None
        self.ax_chunk_matches = None
        self.ax_motion_vectors = None
        
        # Данные для глобальной карты
        self.trajectory_x = []
        self.trajectory_y = []
        self.current_x = 0.0
        self.current_y = 0.0
        
        self.target_idx = 0
        self.target_pts = []
        
        # Данные для траектории БПЛА (его собственное видение)
        self.drone_trajectory_x = []
        self.drone_trajectory_y = []
        
        # Данные для графика погрешности
        self.error_times = []
        self.position_errors = []
        
        # Данные для детекции
        self.current_frame = None
        self.keypoints = []
        self.matches = []

        self._setup_window()
    
    def _setup_window(self):
        """Настраивает общее окно с несколькими областями"""
        plt.ion()
        self.fig = plt.figure(figsize=(16, 10))
        self.fig.canvas.manager.window.title(self.window_title)
        
        # Открываем окно на полный экран
        self.fig.canvas.manager.window.state('zoomed')
        
        # Создаем сетку 3x3 с разными размерами колонок
        gs = self.fig.add_gridspec(2, 3, hspace=0.3, wspace=0.3, width_ratios=[1, 0.7, 1])
        
        # График погрешности позиции (левый верхний угол)
        self.ax_error_plot = self.fig.add_subplot(gs[0, 0])
        self.ax_error_plot.set_title('Погрешность позиции от времени')
        self.ax_error_plot.set_xlabel('Время (кадры)')
        self.ax_error_plot.set_ylabel('Погрешность (пиксели)')
        self.ax_error_plot.grid(True, alpha=0.3)
        
        # Глобальная карта (левый средний угол)
        self.ax_global_map = self.fig.add_subplot(gs[1, 0])
        self.ax_global_map.set_title('Global Map - Траектория полета беспилотника')
        self.ax_global_map.set_xlabel('X координата')
        self.ax_global_map.set_ylabel('Y координата')
        self.ax_global_map.grid(True, alpha=0.3)
        self.ax_global_map.axhline(y=0, color='k', linestyle='-', alpha=0.3)
        self.ax_global_map.axvline(x=0, color='k', linestyle='-', alpha=0.3)
        
        # Сопоставление точек (правый верхний угол)
        self.ax_matches = self.fig.add_subplot(gs[0, 2])
        self.ax_matches.set_title('Feature Matching')
        self.ax_matches.axis('off')
        
        # Сопоставление точек (средний средний угол)
        self.ax_chunk_matches = self.fig.add_subplot(gs[1, 2])
        self.ax_chunk_matches.set_title('Chunk Matching')
        self.ax_chunk_matches.axis('off')
        
        # Визуализация движения ключевых точек (левый нижний угол)
        self.ax_motion_vectors = self.fig.add_subplot(gs[1, 1])
        self.ax_motion_vectors.set_title('Motion Vectors - Движение ключевых точек')
        self.ax_motion_vectors.axis('off')
        
        # Визуализация движения ключевых точек на основе матрицы гомографии
        self.ax_motion_gomography = self.fig.add_subplot(gs[0, 1])
        self.ax_motion_gomography.set_title('Keypoint Detection')
        self.ax_motion_gomography.axis('off')
        
        # Настройки окна
        self.fig.canvas.manager.window.attributes('-topmost', False)
        
        plt.tight_layout()
        plt.show(block=False)

    def set_target_points(self, target_pts):
        """ Обновление списка координат целевых точек """
        self.target_pts = target_pts

    def set_target_index(self, target_idx):
        """ Обновление номера целевой точки """
        self.target_idx = target_idx
    
    def update_global_map(self, x: float, y: float):
        """Обновляет глобальную карту"""
        self.current_x = x
        self.current_y = y
        self.trajectory_x.append(x)
        self.trajectory_y.append(y)
        
        self.ax_global_map.clear()
        self.ax_global_map.set_title('Global Map - Траектория полета беспилотника')
        self.ax_global_map.set_xlabel('X координата')
        self.ax_global_map.set_ylabel('Y координата')
        self.ax_global_map.grid(True, alpha=0.3)
        self.ax_global_map.axhline(y=0, color='k', linestyle='-', alpha=0.3)
        self.ax_global_map.axvline(x=0, color='k', linestyle='-', alpha=0.3)
        
        if len(self.trajectory_x) > 1:
            # Рисуем траекторию оператора (синий цвет)
            self.ax_global_map.plot(self.trajectory_x, self.trajectory_y, 'b-', linewidth=2, label='Режим оператора')

            # Рисуем траекторию БПЛА (пунктирная линия, тонкая)
            if len(self.drone_trajectory_x) > 1:
                self.ax_global_map.plot(self.drone_trajectory_x, self.drone_trajectory_y, 
                                      'g--', linewidth=1, alpha=0.7, label='Данные по одометрии')

            # Рисуем текущую позицию (черная)
            self.ax_global_map.plot(self.current_x, self.current_y, 'ko', markersize=6, label='Текущая позиция')

            # Рисуем ориентиры
            for i in range(len(self.target_pts)):
                if i != self.target_idx:
                    pt = self.target_pts[i]
                    self.ax_global_map.plot(pt[0], pt[1], 'go', markersize=8)

            # Рисуем текущую целевую точку
            if self.target_idx < len(self.target_pts):
                pt = self.target_pts[self.target_idx]
                self.ax_global_map.plot(pt[0], pt[1], 'yo', markersize=8, label='Цель (0, 0)')

            self.ax_global_map.legend()

        # Автоматически масштабируем оси
        if len(self.trajectory_x) > 0:
            margin = 50
            x_min, x_max = min(self.trajectory_x), max(self.trajectory_x)
            y_min, y_max = min(self.trajectory_y), max(self.trajectory_y)

            for pt in self.target_pts:
                x_min = min(x_min, pt[0])
                x_max = max(x_max, pt[0])
                y_min = min(y_min, pt[1])
                y_max = max(y_max, pt[1])
            
            # Учитываем также траекторию БПЛА при масштабировании
            if len(self.drone_trajectory_x) > 0:
                x_min = min(x_min, min(self.drone_trajectory_x))
                x_max = max(x_max, max(self.drone_trajectory_x))
                y_min = min(y_min, min(self.drone_trajectory_y))
                y_max = max(y_max, max(self.drone_trajectory_y))
            
            x_min = min(x_min, 0)
            x_max = max(x_max, 0)
            y_min = min(y_min, 0)
            y_max = max(y_max, 0)
            
            self.ax_global_map.set_xlim(x_min - margin, x_max + margin)
            self.ax_global_map.set_ylim(y_min - margin, y_max + margin)
    
    def update_drone_trajectory(self, drone_x: float, drone_y: float):
        """Обновляет траекторию БПЛА (его собственное видение позиции)"""
        self.drone_trajectory_x.append(drone_x)
        self.drone_trajectory_y.append(drone_y)
    
    def update_error_plot(self, frame_count: int, drone_x: float, drone_y: float, true_x: float, true_y: float):
        """Обновляет график погрешности позиции"""
        # Вычисляем погрешность как расстояние между реальной и предполагаемой позицией
        error = np.sqrt((drone_x - true_x)**2 + (drone_y - true_y)**2)
        
        self.error_times.append(frame_count)
        self.position_errors.append(error)
        
        self.ax_error_plot.clear()
        self.ax_error_plot.set_title('Погрешность позиции от времени')
        self.ax_error_plot.set_xlabel('Время (кадры)')
        self.ax_error_plot.set_ylabel('Погрешность (метры)')
        self.ax_error_plot.grid(True, alpha=0.3)
        
        if len(self.error_times) > 1:
            self.ax_error_plot.plot(self.error_times, self.position_errors, 'b-', linewidth=2)
            
            # Автоматически масштабируем оси
            if len(self.position_errors) > 0:
                margin = 0.1
                error_min, error_max = min(self.position_errors), max(self.position_errors)
                if error_max > error_min:
                    self.ax_error_plot.set_ylim(0, error_max + margin)
                else:
                    self.ax_error_plot.set_ylim(0, 1)

    def update_matches(self, img1: np.ndarray, img2: np.ndarray, 
                      kp1, kp2, matches, transformation_info=None):
        """Обновляет визуализацию сопоставления точек"""
        self.ax_matches.clear()
        self.ax_matches.set_title('Feature Matching')
        
        if img1 is not None and img2 is not None and matches:
            # Рисуем сопоставления
            img_matches = cv2.drawMatches(img1, kp1, img2, kp2, matches, None,
                                         flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
            
            # Конвертируем BGR в RGB
            if len(img_matches.shape) == 3 and img_matches.shape[2] == 3:
                img_matches_rgb = cv2.cvtColor(img_matches, cv2.COLOR_BGR2RGB)
            else:
                img_matches_rgb = img_matches
            
            self.ax_matches.imshow(img_matches_rgb)
            
            # Добавляем информацию о трансформации
            if transformation_info:
                tx, ty = transformation_info['translation']
                angle = transformation_info['rotation']
                
                info_text = f"Translation: ({tx:.2f}, {ty:.2f})"
                info_text2 = f"Rotation: {angle:.2f} rad ({np.degrees(angle):.1f}°)"
                
                self.ax_matches.text(10, 30, info_text, fontsize=8, color='green', 
                                   bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
                self.ax_matches.text(10, 90, info_text2, fontsize=8, color='green',
                                   bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
        
        self.ax_matches.axis('off')
    
    def update_chunk_matches(self, img1: np.ndarray, img2: np.ndarray, 
                      kp1, kp2, matches, transformation_info=None):
        """Обновляет визуализацию сопоставления точек"""
        self.ax_chunk_matches.clear()
        self.ax_chunk_matches.set_title('Chunk Matching')
        
        if img1 is not None and img2 is not None and matches:
            # Рисуем сопоставления
            img_matches = cv2.drawMatches(img1, kp1, img2, kp2, matches, None,
                                         flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS)
            
            # Конвертируем BGR в RGB
            if len(img_matches.shape) == 3 and img_matches.shape[2] == 3:
                img_matches_rgb = cv2.cvtColor(img_matches, cv2.COLOR_BGR2RGB)
            else:
                img_matches_rgb = img_matches
            
            self.ax_chunk_matches.imshow(img_matches_rgb)
            
            # Добавляем информацию о трансформации
            if transformation_info:
                tx, ty = transformation_info['translation']
                angle = transformation_info['rotation']
                
                info_text = f"Translation: ({tx:.2f}, {ty:.2f})"
                info_text2 = f"Rotation: {angle:.2f} rad ({np.degrees(angle):.1f}°)"
                
                self.ax_chunk_matches.text(10, 30, info_text, fontsize=8, color='green', 
                                   bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
                self.ax_chunk_matches.text(10, 90, info_text2, fontsize=8, color='green',
                                   bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
        
        self.ax_chunk_matches.axis('off')
    
    def _update_motion_vectors(self, axes: matplotlib.axes.Axes, current_frame: np.ndarray, prev_keypoints, current_keypoints, matches=None):
        """Обновляет визуализацию движения ключевых точек между кадрами"""
        axes.clear()
        axes.set_title('Motion Vectors - Движение ключевых точек')
        
        if current_frame is not None:
            # Конвертируем BGR в RGB для matplotlib
            if len(current_frame.shape) == 3 and current_frame.shape[2] == 3:
                frame_rgb = cv2.cvtColor(current_frame, cv2.COLOR_BGR2RGB)
            else:
                frame_rgb = current_frame
            
            # Показываем текущий кадр
            axes.imshow(frame_rgb)
            
            # Если есть совпадения, рисуем векторы движения
            if matches is not None and len(matches) > 0:
                # Получаем координаты ключевых точек
                prev_pts = np.array([prev_keypoints[m.queryIdx].pt for m in matches])
                curr_pts = np.array([current_keypoints[m.trainIdx].pt for m in matches])
                
                # Вычисляем векторы движения
                motion_vectors = curr_pts - prev_pts
                
                # Вычисляем длину и направление векторов
                vector_lengths = np.linalg.norm(motion_vectors, axis=1)
                vector_angles = np.arctan2(motion_vectors[:, 1], motion_vectors[:, 0])
                
                # Нормализуем длины для цветовой карты (0-1)
                if len(vector_lengths) > 0:
                    max_length = np.max(vector_lengths)
                    if max_length > 0:
                        normalized_lengths = vector_lengths / max_length
                    else:
                        normalized_lengths = np.zeros_like(vector_lengths)
                else:
                    normalized_lengths = np.array([])
                
                # Рисуем векторы с цветовой индикацией
                for i, (start_pt, end_pt, length, angle, norm_length) in enumerate(
                    zip(prev_pts, curr_pts, vector_lengths, vector_angles, normalized_lengths)):
                    
                    # Цвет зависит от направления (угол -> HSV)
                    hue = (angle + np.pi) / (2 * np.pi)  # Нормализуем угол к 0-1
                    saturation = 1.0
                    value = 0.8 + 0.2 * norm_length  # Яркость зависит от длины
                    
                    # Конвертируем HSV в RGB
                    import matplotlib.colors as mcolors
                    rgb = mcolors.hsv_to_rgb([hue, saturation, value])
                    
                    # Толщина линии зависит от длины вектора
                    linewidth = max(1, min(5, 2 + 3 * norm_length))
                    
                    # Рисуем вектор
                    axes.arrow(
                        start_pt[0], start_pt[1], 
                        end_pt[0] - start_pt[0], end_pt[1] - start_pt[1],
                        head_width=3, head_length=5, 
                        fc=rgb, ec=rgb, alpha=0.8, linewidth=linewidth
                    )
                    
                    # Добавляем текст с информацией о движении
                    # if length > 5:  # Показываем только для значительных движений
                    #     axes.text(
                    #         end_pt[0] + 5, end_pt[1] + 5, 
                    #         f'{length:.1f}px', fontsize=6, color='white',
                    #         bbox=dict(boxstyle="round,pad=0.2", facecolor="black", alpha=0.7)
                    #     )
                
                # Добавляем легенду с цветовой схемой
                # legend_text = "Цвет: направление, Яркость: скорость"
                # axes.text(
                #     10, 30, legend_text, fontsize=8, color='white',
                #     bbox=dict(boxstyle="round,pad=0.3", facecolor="black", alpha=0.8)
                # )
                
                # Статистика движения
                # if len(vector_lengths) > 0:
                #     avg_speed = np.mean(vector_lengths)
                #     max_speed = np.max(vector_lengths)
                #     stats_text = f"Средняя скорость: {avg_speed:.1f}px\nМаксимальная: {max_speed:.1f}px"
                #     axes.text(
                #         10, 60, stats_text, fontsize=8, color='white',
                #         bbox=dict(boxstyle="round,pad=0.3", facecolor="black", alpha=0.8)
                #     )
        
        axes.axis('off')
    
    def update_motion_vectors(self, current_frame: np.ndarray, prev_keypoints, current_keypoints, matches=None):
        self._update_motion_vectors(self.ax_motion_vectors, current_frame, prev_keypoints, current_keypoints, matches)

    def update_motion_vectors(self, current_frame: np.ndarray, prev_keypoints, current_keypoints, matches=None):
        self._update_motion_vectors(self.ax_motion_vectors, current_frame, prev_keypoints, current_keypoints, matches)

    def update_motion_gomography(self, current_frame: np.ndarray, prev_keypoints, current_keypoints, matches=None):
        self._update_motion_vectors(self.ax_motion_gomography, current_frame, prev_keypoints, current_keypoints, matches)
    
    def update_homography_grid(self, current_frame: np.ndarray, homography_matrix: np.ndarray, grid_step: int = 80):
        """
        Визуализирует движение точек по сетке на основе матрицы гомографии
        """
        self.ax_motion_gomography.clear()
        self.ax_motion_gomography.set_title('Движение точек по сетке')
        
        if current_frame is None or homography_matrix is None:
            self.ax_motion_gomography.axis('off')
            return
            
        # Конвертируем BGR в RGB для matplotlib
        if len(current_frame.shape) == 3 and current_frame.shape[2] == 3:
            frame_rgb = cv2.cvtColor(current_frame, cv2.COLOR_BGR2RGB)
        else:
            frame_rgb = current_frame
        
        # Показываем текущий кадр
        self.ax_motion_gomography.imshow(frame_rgb)
        
        # Получаем размеры изображения и центр
        height, width = current_frame.shape[:2]
        center_x, center_y = width // 2, height // 2
        
        # Создаем сетку точек с заданным шагом
        grid_points = []
        for y in range(grid_step, height, grid_step):
            for x in range(grid_step, width, grid_step):
                grid_points.append([x, y])
        
        if len(grid_points) == 0:
            self.ax_motion_gomography.axis('off')
            return
            
        # Конвертируем в numpy массив и отцентрируем координаты относительно центра изображения
        grid_points = np.array(grid_points, dtype=np.float32)
        grid_points_centered = []
        for pt in grid_points:
            # Отцентрируем координаты точно так же, как в detect_and_match_keypoints
            centered_x = pt[0] - center_x
            centered_y = center_y - pt[1]  # Инвертируем Y (изображение Y направлен вниз)
            grid_points_centered.append([centered_x, centered_y])
        
        grid_points_centered = np.array(grid_points_centered, dtype=np.float32)
        grid_points_homogeneous = np.column_stack([grid_points_centered, np.ones(len(grid_points_centered))])
        
        # Применяем матрицу гомографии
        transformed_points_homogeneous = homography_matrix @ grid_points_homogeneous.T
        transformed_points_homogeneous = transformed_points_homogeneous.T
        
        # Нормализуем по третьей координате (перспективное преобразование)
        transformed_points_centered = transformed_points_homogeneous[:, :2] / transformed_points_homogeneous[:, 2:3]
        
        # Конвертируем обратно в координаты изображения
        transformed_points = []
        for pt in transformed_points_centered:
            # Обратное преобразование от центрированных координат к координатам изображения
            img_x = pt[0] + center_x
            img_y = center_y - pt[1]  # Инвертируем Y обратно
            transformed_points.append([img_x, img_y])
        
        transformed_points = np.array(transformed_points, dtype=np.float32)
        
        # Фильтруем точки, которые остались в пределах изображения
        valid_indices = (
            (transformed_points[:, 0] >= 0) & 
            (transformed_points[:, 0] < width) & 
            (transformed_points[:, 1] >= 0) & 
            (transformed_points[:, 1] < height)
        )
        
        if np.sum(valid_indices) == 0:
            self.ax_motion_gomography.axis('off')
            return
            
        # Получаем валидные исходные и трансформированные точки
        valid_source_points = grid_points[valid_indices]
        valid_transformed_points = transformed_points[valid_indices]
        
        # Вычисляем векторы движения
        motion_vectors = valid_transformed_points - valid_source_points
        
        # Вычисляем длину и направление векторов
        vector_lengths = np.linalg.norm(motion_vectors, axis=1)
        
        if len(vector_lengths) > 0:
            # Нормализуем длины для цветовой карты (0-1)
            max_length = np.max(vector_lengths)
            if max_length > 0:
                normalized_lengths = vector_lengths / max_length
            else:
                normalized_lengths = np.zeros_like(vector_lengths)
                
            # Рисуем векторы движения
            for i, (start_pt, end_pt, length, norm_length) in enumerate(
                zip(valid_source_points, valid_transformed_points, vector_lengths, normalized_lengths)):
                
                # Пропускаем очень маленькие векторы
                if length < 1.0:
                    continue
                    
                # Цвет зависит от длины вектора (синий -> красный)
                if norm_length < 0.5:
                    color = [0, norm_length * 2, 1 - norm_length * 2]  # Синий -> Голубой
                else:
                    color = [(norm_length - 0.5) * 2, 1 - (norm_length - 0.5) * 2, 0]  # Голубой -> Красный
                
                # Толщина линии зависит от длины вектора
                linewidth = max(1, min(4, 1 + 3 * norm_length))
                
                # Рисуем вектор
                self.ax_motion_gomography.arrow(
                    start_pt[0], start_pt[1], 
                    end_pt[0] - start_pt[0], end_pt[1] - start_pt[1],
                    head_width=3, head_length=5, 
                    fc=color, ec=color, alpha=0.8, linewidth=linewidth
                )
                
                # Рисуем исходную точку сетки
                self.ax_motion_gomography.plot(start_pt[0], start_pt[1], 'o', 
                                            color='green', markersize=3, alpha=0.7)
                
                # Рисуем целевую точку
                self.ax_motion_gomography.plot(end_pt[0], end_pt[1], 's', 
                                            color='red', markersize=2, alpha=0.7)
            
            # Добавляем информацию о статистике
            avg_length = np.mean(vector_lengths)
            max_length = np.max(vector_lengths)
            total_points = len(vector_lengths)
            
            # info_text = f"Точек сетки: {total_points}\nСреднее движение: {avg_length:.1f}px\nМакс. движение: {max_length:.1f}px"
            # self.ax_motion_gomography.text(
            #     10, 30, info_text, fontsize=8, color='white',
            #     bbox=dict(boxstyle="round,pad=0.3", facecolor="black", alpha=0.8)
            # )
            
            # Добавляем легенду
            # legend_text = "Зеленые точки: исходные\nКрасные квадраты: целевые\nЦвет стрелок: скорость"
            # self.ax_motion_gomography.text(
            #     10, 90, legend_text, fontsize=7, color='white',
            #     bbox=dict(boxstyle="round,pad=0.3", facecolor="black", alpha=0.8)
            # )
        
        self.ax_motion_gomography.axis('off')

    def update_display(self):
        """Обновляет отображение всех областей"""
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
        plt.pause(0.2)
    
    def close(self):
        """Закрывает окно"""
        plt.close(self.fig)
    
    def show_final(self):
        """Показывает финальное состояние окна"""
        plt.ioff()
        print("Симуляция завершена. Окно визуализации остается открытым для анализа.")
        plt.pause(100000)

    def pause(self, duration: float):
        plt.pause(duration) 
