from time import sleep
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
from matplotlib.path import Path
import matplotlib.patches as patches
import numpy as np

class TrajectoryDrawer:
    def __init__(self, image_path):
        self.image_path = image_path
        self.points = []
        self.fig, self.ax = plt.subplots(figsize=(12, 8))
        
        # Загрузка фонового изображения
        self.img = mpimg.imread(image_path)
        self.ax.imshow(self.img)
        
        # Настройка графика
        self.ax.set_title('Нарисуйте траекторию движения (кликните для точек, Enter для завершения)')
        self.ax.set_xlabel('X координата')
        self.ax.set_ylabel('Y координата')
        self.closed = False
        
        # Подключение обработчиков событий
        self.fig.canvas.mpl_connect('button_press_event', self.on_click)
        self.fig.canvas.mpl_connect('key_press_event', self.on_key)
        
        self.line, = self.ax.plot([], [], 'ro-', linewidth=2, markersize=6)
        
    def on_click(self, event):
        if event.inaxes == self.ax:
            # Добавляем точку только если клик в области графика
            x, y = event.xdata, event.ydata
            self.points.append((x, y))
            self.update_plot()
            
    def on_key(self, event):
        if event.key == 'enter':
            self.complete_trajectory()
        elif event.key == 'escape':
            self.clear_trajectory()
        elif event.key == 'e':
            if len(self.points):
                self.points.append(self.points[0])
                self.update_plot()
            
            
    def update_plot(self):
        if self.points:
            x_coords = [p[0] for p in self.points]
            y_coords = [p[1] for p in self.points]
            self.line.set_data(x_coords, y_coords)
            self.fig.canvas.draw()
            
    def complete_trajectory(self, *args):
        if len(self.points) > 1:
            print("Траектория завершена!")
            print(f"Количество точек: {len(self.points)}")
            print("Координаты точек:")
            for i, (x, y) in enumerate(self.points):
                print(f"Точка {i+1}: ({x:.2f}, {y:.2f})")
                
            # Можно сохранить координаты в файл
            self.destroy()
            self.closed = True
            
    def clear_trajectory(self):
        self.points = []
        self.line.set_data([], [])
        self.fig.canvas.draw()
        print("Траектория очищена")
        
    def save_trajectory(self):
        # Сохранение координат в файл
        with open('trajectory_coordinates.txt', 'w') as f:
            for i, (x, y) in enumerate(self.points):
                f.write(f"{i+1},{x:.4f},{y:.4f}\n")
        print("Координаты сохранены в trajectory_coordinates.txt")

    def destroy(self):
        plt.close(self.fig)
        
    def show(self):
        plt.tight_layout()
        plt.show()

    def wait(self):
        while not self.closed:
            sleep(0.1)

# Использование
if __name__ == "__main__":
    # Укажите путь к вашему изображению
    image_path = "map.jpg"  # Замените на путь к вашему изображению
    
    drawer = TrajectoryDrawer(image_path)
    drawer.show()