from pathlib import Path
from simulator import Simulator
from time import sleep
from trajectory_drawer import TrajectoryDrawer
from visualization import VisualizationManager
from yandex_map import YandexMap
from vision_chunk import VisionChunk
from utility import cv2_to_pil
import random

import cv2
import matplotlib.pyplot as plt
import numpy as np

import autopilot

def make_global_photo(filename):
    yandexMap = YandexMap()
    yandexMap.save_photo(filename)
    yandexMap.destroy()

def get_trajectory_points(bg_img: str) -> list[(float, float)]:
    trajectoryDrawer = TrajectoryDrawer(bg_img)
    trajectoryDrawer.show()
    trajectoryDrawer.wait()
    points = list(map(lambda p: [p[0] / trajectoryDrawer.img.shape[1], p[1] / trajectoryDrawer.img.shape[0]], trajectoryDrawer.points))
    return points

def main():
    # Скриншот местности
    # make_global_photo('map.jpg')

    # Получаем траекторию от пользователя
    points = get_trajectory_points('map.jpg')
#     points = [np.float64(0.5443937502226799), np.float64(0.4030424838774785)], [np.float64(0.18517133490120316), np.float64(0.4586935604608052)], [np.float64(0.1641887171838272), np.float64(0.5586383510594329)], [np.float64(0.587198290366127), 
# np.float64(0.5699957136274587)]
    # points = [[np.float64(0.5028362002950056), np.float64(0.508463417020251)], [np.float64(0.5079239482442035), np.float64(0.508463417020251)], [np.float64(0.5045321162780716), np.float64(0.5210964284169014)], [np.float64(0.509054558899581), np.float64(0.5200436774671806)]]
    # print(points)

    # Для каждой точки сделаем приближенный снимок
    yandexMap = YandexMap()
    chunks: list[VisionChunk] = []

    plt.ion()
    for i in range(len(points)):
        point = points[i]
        yandexMap.scroll(point[0], point[1], 5, True)
        sleep(1)
        cv2_img = yandexMap.make_screenshot(point[0], point[1], 0.2, 0.2)
        img = cv2_to_pil(cv2_img)
        chunk = VisionChunk(img)
        Path('chunks').mkdir(exist_ok=True)
        chunk.save_image(Path('.') / 'chunks' / f'chunk_{i}.png')
        plt.subplot(1, len(points), i+1)
        plt.imshow(img)
        plt.pause(0.25)
        yandexMap.scroll(point[0], point[1], 5, False)

    plt.tight_layout()

    # Выделим на каждой картинке ключевые точки
    for i in range(len(points)):
        chunk = VisionChunk.load_image(Path('chunks') / f'chunk_{i}.png')
        chunks.append(chunk)
        kp, des = chunk.compute_keypoints()

        plt.subplot(1, len(points), i+1)
        kp_coords = np.array([j.pt for j in kp])
        if len(kp_coords) > 0:
            plt.scatter(kp_coords[:, 0], kp_coords[:, 1], c='red', s=20, alpha=0.7, marker='o')
        plt.pause(0.2)
    plt.ioff()

    plt.show(block=True)

    # Начнём симуляцию полёта с первой точки
    yandexMap.scroll(points[0][0], points[0][1], 5, True)
    sleep(0.2)
    yandexMap.make_as_center(*points[0])
    sleep(1)

    vis_manager = VisualizationManager()
    width, height = yandexMap.get_size()
    # print(width, height)
    points_coords = np.array(list(map(lambda p: [
            (p[0] - points[0][0]) * width, (points[0][1] - p[1]) * height
        ], points)))
    points_coords *= 2 ** 4
    pilot = autopilot.AutoPilot(points_coords, chunks, vis_manager)
    simulator = Simulator(yandexMap)
    pilot.target_idx = 0

    chunk = simulator.get_chunk()
    command = pilot.handle(chunk)

    vis_manager.update_display()
    vis_manager.pause(1)

    vis_manager.set_target_points(points_coords)

    proc_time = np.array([])

    zoom_next_event = random.randint(5, 10)

    for i in range(10000000000):
        print(f"Image #{i}")
        if i == zoom_next_event:
            r = random.randint(0, 1)
            direction = ['up', 'down'][r]
            simulator.change_zoom(direction)
            zoom_next_event = i + random.randint(20, 40)

        chunk = simulator.get_chunk()
        command = pilot.handle(chunk)

        proc_time = np.append(proc_time, command.proccessing_time)
        
        # Save Image
        chunk.save_image(Path('images') / f"photo_{i}.png")

        vis_manager.update_display()
        vis_manager.pause(0.2)

        if command.stop:
            break

        vis_manager.set_target_index(pilot.target_idx)
        vis_manager.update_drone_trajectory(pilot.geo.x, pilot.geo.y)
        vis_manager.update_global_map(simulator.geo.x, simulator.geo.y)
        vis_manager.update_error_plot(i, pilot.geo.x, pilot.geo.y, simulator.geo.x, simulator.geo.y)

        simulator.handle(command.dangle, command.velocity)

        vis_manager.update_display()
        vis_manager.pause(0.2)

        last_proc_times = proc_time[-10:]
        print("Average FPS:", 1 / last_proc_times.mean())

    vis_manager.show_final()

if __name__ == "__main__":
    main()