#!/usr/bin/env python

""" """

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

from dataclasses import dataclass, field

from argaze import DataStructures

import json
import numpy
import cv2
import cv2.aruco as aruco

K0 = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 0.]])
"""Define default optic intrinsic parameters matrix."""

D0 = numpy.array([0.0, 0.0, 0.0, 0.0, 0.0])
"""Define default optic distorsion coefficients vector."""

@dataclass
class OpticParameters():
    """Define optic parameters outputed by optic calibrator."""

    rms: float = field(default=0)
    """Root Mean Square error of calibration."""

    dimensions: numpy.array = field(default_factory=lambda : numpy.array([0, 0]))
    """Image dimensions in pixels from which the calibration have been done."""

    K: numpy.array = field(default_factory=lambda : K0)
    """Intrinsic parameters matrix (focal lengths and principal point)."""

    D: numpy.array = field(default_factory=lambda : D0)
    """Distorsion coefficients vector."""

    @classmethod
    def from_json(self, json_filepath):
        """Load optical parameters from .json file."""

        with open(json_filepath) as calibration_file:

            return OpticParameters(**json.load(calibration_file))

    def to_json(self, json_filepath):
        """Save optical parameters into .json file."""

        with open(json_filepath, 'w', encoding='utf-8') as calibration_file:

            json.dump(self, calibration_file, ensure_ascii=False, indent=4, cls=DataStructures.JsonEncoder)

    def __str__(self) -> str:
        """String display"""

        output = f'\trms: {self.rms}\n'
        output += f'\tdimensions: {self.dimensions}\n'
        output += f'\tK: {self.K}\n'
        output += f'\tD: {self.D}\n'

        return output

    def draw(self, image: numpy.array, width:float, height:float, z:float, color=(0, 0, 255)):
        """Draw grid to display K and D"""

        # Edit 3D grid
        grid_3D = []
        for x in range(-int(width/2), int(width/2)):
            for y in range(-int(height/2), int(height/2)):
                grid_3D.append([x, y, z])

        # Project 3d grid
        grid_2D, _ = cv2.projectPoints(numpy.array(grid_3D).astype(float), numpy.array([0., 0., 0.]), numpy.array([0., 0., 0.]), numpy.array(self.K), -numpy.array(self.D))

        # Draw projection
        for point in grid_2D:

            # Ignore point out out field
            try:
                cv2.circle(image, point.astype(int)[0], 1, color, -1)
            except:
                pass

class ArUcoOpticCalibrator():
    """Handle optic calibration process."""

    def __init__(self,):

        # Calibration data
        self.__corners_set_number = 0
        self.__corners_set = []
        self.__corners_set_ids = []

    def calibrate(self, board, dimensions:tuple = (0, 0)) -> OpticParameters:
        """Retrieve K and D parameters from stored calibration data.

        Parameters:
            dimensions: camera image dimensions

        Returns:
            Optic parameters
        """

        if self.__corners_set_number > 0:

            rms, K, D, r, t = aruco.calibrateCameraCharuco(self.__corners_set, self.__corners_set_ids, board.model, dimensions, None, None)

            return OpticParameters(rms, dimensions, K, D)

    def reset_calibration_data(self):
        """Clear all calibration data."""

        self.__corners_set_number = 0
        self.__corners_set = []
        self.__corners_set_ids = []

    def store_calibration_data(self, corners, corners_identifiers):
        """Store calibration data."""

        self.__corners_set_number += 1
        self.__corners_set.append(corners)
        self.__corners_set_ids.append(corners_identifiers)

    @property
    def calibration_data_count(self) -> int:
        """Get how much calibration data are stored."""

        return self.__corners_set_number