#!/usr/bin/env python

""" """

__author__ = "Théo de la Hogue"
__credits__ = []
__copyright__ = "Copyright 2023, Ecole Nationale de l'Aviation Civile (ENAC)"
__license__ = "BSD"

import argparse
import time
import itertools

from argaze.ArUcoMarkers import ArUcoCamera, ArUcoMarkersGroup
from argaze.utils import UtilsFeatures

import cv2
import numpy

def main():
    """
    Load a MOVIE and an ArUcoCamera CONFIGURATION to detect ArUco markers inside a selected movie frame then, export detected ArUco markers group as .obj file into an OUTPUT folder.
    """

    # Manage arguments
    parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0])
    parser.add_argument('movie', metavar='MOVIE', type=str, default=None, help='movie path')
    parser.add_argument('configuration', metavar='CONFIGURATION', type=str, default=None, help='ArUco camera configuration')

    parser.add_argument('-s','--start', metavar='START', type=float, default=0., help='start time in second')
    parser.add_argument('-o', '--output', metavar='OUTPUT', type=str, default='.', help='export folder path')
    args = parser.parse_args()

    # Load movie
    video_capture = cv2.VideoCapture(args.movie)

    video_fps = video_capture.get(cv2.CAP_PROP_FPS)
    image_width = int(video_capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    image_height = int(video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Load ArUco camera
    aruco_camera = ArUcoCamera.ArUcoCamera.from_json(args.configuration)

    # Create empty ArUco scene
    aruco_markers_group = None

    # Create a window
    cv2.namedWindow(aruco_camera.name, cv2.WINDOW_AUTOSIZE)

    # Enable exit signal handler
    exit = UtilsFeatures.ExitSignalHandler()

    # Init image selection
    current_image_index = -1
    _, current_image = video_capture.read()
    next_image_index = int(args.start * video_fps)
    refresh = False

    while not exit.status():

        # Select a new image and detect markers once
        if next_image_index != current_image_index or refresh:

            video_capture.set(cv2.CAP_PROP_POS_FRAMES, next_image_index)

            success, video_image = video_capture.read()

            if success:

                # Refresh once
                refresh = False

                current_image_index = video_capture.get(cv2.CAP_PROP_POS_FRAMES) - 1
                current_image_time = video_capture.get(cv2.CAP_PROP_POS_MSEC)

                # Detect markers
                detection_time, projection_time, exceptions = aruco_camera.watch(video_image)

                # Estimate each markers pose
                aruco_camera.aruco_detector.estimate_markers_pose(aruco_camera.aruco_detector.detected_markers)

                # Build aruco scene from detected markers
                aruco_markers_group = ArUcoMarkersGroup.ArUcoMarkersGroup(aruco_camera.aruco_detector.marker_size, aruco_camera.aruco_detector.dictionary, aruco_camera.aruco_detector.detected_markers)

                # Get camera image
                camera_image = aruco_camera.image()

                # Write detected markers
                cv2.putText(camera_image, f'Detecting markers {list(aruco_camera.aruco_detector.detected_markers.keys())}', (20, aruco_camera.size[1]-40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                
                # Write timing
                cv2.putText(camera_image, f'Frame at {int(current_image_time)}ms', (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                cv2.putText(camera_image, f'Detection {int(detection_time)}ms', (20, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
                cv2.putText(camera_image, f'Projection {int(projection_time)}ms', (20, 120), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1, cv2.LINE_AA)
        
                # Write documentation
                cv2.putText(camera_image, f'<- previous image', (aruco_camera.size[0]-500, aruco_camera.size[1]-120), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
                cv2.putText(camera_image, f'-> next image', (aruco_camera.size[0]-500, aruco_camera.size[1]-80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)
                cv2.putText(camera_image, f'Ctrl+s: export ArUco markers', (aruco_camera.size[0]-500, aruco_camera.size[1]-40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1, cv2.LINE_AA)

                # Copy image
                current_image = camera_image.copy()

        # Keep last image
        else:

            video_image = current_image.copy()

        key_pressed = cv2.waitKey(10)

        #if key_pressed != -1:
        #    print(key_pressed)

        # Select previous image with left arrow
        if key_pressed == 2:
            next_image_index -= 1

        # Select next image with right arrow
        if key_pressed == 3:
            next_image_index += 1

        # Clip image index
        if next_image_index < 0:
            next_image_index = 0

        # r: reload configuration
        if key_pressed == 114:
            
            aruco_camera = ArUcoCamera.ArUcoCamera.from_json(args.configuration)
            refresh = True
            print('Configuration reloaded')

        # Save selected marker edition using 'Ctrl + s'
        if key_pressed == 19:

            if aruco_markers_group:

                aruco_markers_group.to_obj(f'{args.output}/{int(current_image_time)}-aruco_markers_group.obj')
                print(f'ArUco markers saved into {args.output}')

            else:

                print(f'No ArUco markers to export')

        # Close window using 'Esc' key
        if key_pressed == 27:
            break

        # Display video
        cv2.imshow(aruco_camera.name, video_image)

    # Close movie capture
    video_capture.release()

    # Stop image display
    cv2.destroyAllWindows()

if __name__ == '__main__':

    main()