#!/usr/bin/env python

import argparse
import os

from argaze import DataStructures, GazeFeatures
from argaze.ArUcoMarkers import ArUcoMarkersDictionary
from argaze.AreaOfInterest import *

import cv2 as cv
import numpy

from ivy.std_api import *

def main():
    """
    Define AOI scene from a ArUco marker and bind to Ivy default bus to receive live look at pointer data.
    """

    # Manage arguments
    parser = argparse.ArgumentParser(description=main.__doc__.split('-')[0])
    parser.add_argument('-y', '--ivy_bus', metavar='IVY_BUS', type=str, default='0.0.0.0:2010', help='Ivy bus ip and port')
    parser.add_argument('-a', '--aoi_scene', metavar='AOI_SCENE', type=str, default='aoi3D_scene.obj', help='obj aoi scene filepath')
    parser.add_argument('-d', '--dictionary', metavar='DICT', type=str, default='DICT_ARUCO_ORIGINAL', help='aruco marker dictionnary (DICT_4X4_50, DICT_4X4_100, DICT_4X4_250, DICT_4X4_1000, DICT_5X5_50, DICT_5X5_100, DICT_5X5_250, DICT_5X5_1000, DICT_6X6_50, DICT_6X6_100, DICT_6X6_250, DICT_6X6_1000, DICT_7X7_50, DICT_7X7_100, DICT_7X7_250, DICT_7X7_1000, DICT_ARUCO_ORIGINAL,DICT_APRILTAG_16h5, DICT_APRILTAG_25h9, DICT_APRILTAG_36h10, DICT_APRILTAG_36h11)')
    parser.add_argument('-m', '--marker_size', metavar='MKR', type=float, default=6, help='aruco marker size (cm)')
    parser.add_argument('-i', '--marker_id', metavar='MARKER_ID', type=int, default=0, help='marker id to display')
    args = parser.parse_args()

    # Enable Ivy bus
    IvyInit(os.path.basename(__file__))
    IvyStart(args.ivy_bus)

    def on_looking_message(*args):

        look_at = numpy.fromstring(args[2].replace('[','').replace(']',''), dtype=float, count=2, sep=', ')

        visu_gaze_pixel = aoi2D_visu_scene[args[1]].looked_pixel(look_at)

        cv.circle(visu_frame, visu_gaze_pixel, 4, (0, 0, 255), -1)

    IvyBindMsg(on_looking_message, 'looking (.*) at (.*)')

    # Create AOIs 3D scene
    aoi3D_scene = AOI3DScene.AOI3DScene()
    aoi3D_scene.load(args.aoi_scene)
    print(f'AOIs names: {aoi3D_scene.keys()}')

    # Create a visual scan visualisation frame
    visu_width = 1920
    visu_height = 1080
    visu_ratio = visu_height
    visu_frame = numpy.full((visu_height, visu_width, 3), 255, dtype=numpy.uint8)

    cv.imshow('Scene', visu_frame)

    # Project 3D scene on the reference frame
    # TODO : center projection on a reference AOI
    ref_aoi = 'Scene_Plan'

    # TODO: pass the reference AOI in argument
    aoi3D_scene.rotation = numpy.array([[-numpy.pi, 0.0, 0.0]])
    aoi3D_scene.translation = numpy.array([[19.0, 8.0, 25.0]])

    # Edit a projection matrix for the reference frame
    K0 = numpy.array([[visu_ratio, 0.0, visu_width/2], [0.0, visu_ratio, visu_height/2], [0.0, 0.0, 1.0]])

    aoi2D_visu_scene = aoi3D_scene.project(K0)

    # Create aruco markers dictionary
    aruco_markers_dict = ArUcoMarkersDictionary.ArUcoMarkersDictionary(args.dictionary)

    # Create aruco marker
    marker_box = aoi2D_visu_scene['Marker_Plan'].bounding_box().astype(int)
    marker_size = marker_box[2] - marker_box[0]
    marker = aruco_markers_dict.create_marker(args.marker_id, int(marker_size[0]))
    print(f'Creating Aruco marker {args.marker_id} from the {args.dictionary} dictionary')

    def draw_scene():

        # Clear frame
        visu_frame[:] = 255

        # Display AOI 2D scene
        for name, aoi in aoi2D_visu_scene.items():
            aoi.draw(visu_frame, (0, 0, 0))

        # Display aruco marker
        visu_frame[marker_box[0][1]:marker_box[2][1], marker_box[0][0]:marker_box[2][0], :] = marker

    # On mouse over : redraw scene and draw target
    def on_mouse_event(event, x, y, flags, param):

        draw_scene()

        # Draw target
        cv.circle(visu_frame, (x, y), 40, (0, 255, 255), -1)

    cv.setMouseCallback('Scene', on_mouse_event)
    
    # Screen display loop
    try:

        draw_scene()
    
        while True:

            # Close window using 'Esc' key
            if cv.waitKey(1) == 27:
                break

            cv.imshow('Scene', visu_frame)

    # Exit on 'ctrl+C' interruption
    except KeyboardInterrupt:
        pass

    # Stop frame display
    cv.destroyAllWindows()

if __name__ == '__main__':

    main()