#!/usr/bin/env python

"""
tobii_camera_calibration.py

Author: 
    - Théo de la Hogue, theo.de-la-hogue@enac.fr

This program:
    - Captures board pictures with a full displayed board inside
    - Outputs camera calibration data into a camera.json file

Reference:
    - https://automaticaddison.com/how-to-perform-pose-estimation-using-an-aruco-marker/
"""
import os
import time

from TobiiGlassesPro2 import TobiiController, TobiiVideo
from ArUcoMarkers import ArUcoBoard, ArUcoTracker, ArUcoCamera

import cv2 as cv

# tobii glasses ip address
ip_address = '192.168.1.10'

# manage export folder
current_folder = os.path.dirname(__file__)
export_folder = os.path.join(current_folder, '_export')
if not os.path.exists(export_folder):
    os.makedirs(export_folder)
    print(f'\'_export\' folder created')

# create tobii controller
tobii_controller = TobiiController.TobiiController(ip_address, 'ArGaze', 1)

# create tobii video thread
tobii_video_thread = TobiiVideo.TobiiVideoThread(tobii_controller)
tobii_video_thread.start()

# create aruco camera
aruco_camera = ArUcoCamera.ArUcoCamera()

# create aruco board
aruco_board = ArUcoBoard.ArUcoBoard('DICT_4X4_50', 7, 5, 5, 3) # 7 columns, 5 rows, square size (cm), marker size (cm)

# create aruco tracker
aruco_tracker = ArUcoTracker.ArUcoTracker('DICT_4X4_50', 6, aruco_camera) # aruco dictionaries, marker length (cm), camera

# start tobii glasses streaming
tobii_controller.start_streaming()

print("Camera calibration starts")
print("Waiting for calibration board...")

frame_width = 0
frame_height = 0

expected_markers_number = len(aruco_board.get_ids())
expected_corners_number = (aruco_board.get_size()[0] - 1 ) * (aruco_board.get_size()[1] - 1)

# capture frame with a full displayed board
while True:

    frame, frame_width, frame_height, frame_time, frame_pts = tobii_video_thread.read()

    # track all markers in the board
    aruco_tracker.track_board(frame, aruco_board, expected_markers_number)

    # draw only markers
    aruco_tracker.draw(frame)

    # draw current calibration data count
    cv.putText(frame, f'Capture: {aruco_camera.get_calibration_data_count()}', (50, 50), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv.LINE_AA)
    cv.imshow('Tobii Camera Calibration', frame)

    # if all board corners are detected
    if aruco_tracker.get_board_corners_number() == expected_corners_number:

        # draw board corners to notify a capture is done
        aruco_tracker.draw_board(frame)

        # append data
        aruco_camera.store_calibration_data(aruco_tracker.get_board_corners(), aruco_tracker.get_board_corners_ids())

        cv.imshow(f'Tobii Camera Calibration', frame)

        time.sleep(2)

    # quit on 'Esc' command
    key = cv.waitKey(1)
    if key == 27:
        cv.destroyAllWindows()
        break

# stop tobii objects
tobii_video_thread.stop()

tobii_controller.stop_streaming()
tobii_controller.close()

print('\nCalibrating camera...')
aruco_camera.calibrate(aruco_board, frame_width, frame_height)

print('\nCalibration succeeded!')
print(f'\nRMS:\n{aruco_camera.get_rms()}')
print(f'\nCamera matrix:\n{aruco_camera.get_K()}')
print(f'\nDistortion coefficients:\n{aruco_camera.get_D()}')

aruco_camera.save_calibration_file(os.join(export_folder,'tobii_camera.json'))

print(f'\nCalibration data exported into tobii_camera.json file')