"""The module for computing a Mahalanobis-based metric imitating the GIoU association metric.

Classes:
    MetricIouLikeTwoYaw: the association metrics between detections and Kalman filters.
"""

from typing import Sequence

import numpy as np

from scipy.spatial.transform import Rotation

from kinematic_tracker.types import FMT, FT

from .metric_driver_base import MetricDriverBase


class MetricIouLikeTwoYaw(MetricDriverBase):
    """Computes a Mahalanobis-based metric imitating the GIoU-yaw association metric.

    This class inherits from MetricDriverBase and implements the metric computation.
    The computation is performed between cuboids (3D bounding boxes) rotated by
    yaw angle (angle around z axis).
    """

    def __init__(
        self,
        num_reports_max: int,
        num_targets_max: int,
        num_z: int,
        ind_pos_w_rz_size: Sequence[int],
    ) -> None:
        """Initialize the MetricGIoUAligned instance.

        Args:
            num_reports_max: maximum number of detection reports to handle.
            num_targets_max: maximum number of targets (Kalman filters) to handle.
            num_z: number of variables in the measurement vector.
            ind_pos_w_rz_size: location of variables in the measurement vector.
                          By default, the indices are 0,1,2,3,6,7,8,9, i.e.
                          the center of the cuboid is taken as first three variables,
                          the scalar part of the quaternion follows the center,
                          the z-projection of the quaternion follows the scalar part,
                          and sizes (dimensions) of the cuboids are taken from
                          the last three variables.
        """
        assert num_z > 7, 'Number of variables should be 8 or greater.'
        super().__init__(num_reports_max, num_targets_max, num_z)
        self.ind_pos_w_rz_size = np.array(ind_pos_w_rz_size, dtype=int)
        assert self.ind_pos_w_rz_size.ndim == 1, 'Indices should be one-dimensional.'
        assert self.ind_pos_w_rz_size.size == 8, 'There should be 8 fancy indices.'
        self.euler_angles_t3 = np.zeros((num_targets_max, 3))
        self.xx_tn = np.zeros((num_targets_max, 10))
        self.xx_rn = np.zeros((num_reports_max, 10))
        self.xx_t4 = np.zeros((num_targets_max, 4))
        self.xx_idx = np.array((0, 1, 2, 3, 6, 7, 8, 9), int)
        self.square_size_t3 = np.zeros((num_targets_max, 3))
        self.pos = 0.2136
        self.size = 0.25
        self.precision_dia_tn = np.zeros((num_targets_max, 10))
        qua = 0.1345
        self.precision_dia_tn[:, 3:7] = qua

    def set_target_tz(self, filters: FT) -> None:
        for t, kf in enumerate(filters):
            np.dot(kf.measurementMatrix, kf.statePre[:, 0], out=self.vec_z[:, 0])
            self.xx_tn[t, self.xx_idx] = self.vec_z[self.ind_pos_w_rz_size, 0]

    def compute_metric(self, det_rz: FMT, filters: FT) -> FMT:
        """Compute GIoU association metric between detections and Kalman filter states.

        Args:
            det_rz: Detection measurements, either as a 2D numpy array.
            filters: Sequence of OpenCV KalmanFilter objects representing tracked targets.

        Returns:
            The matrix of scores between each detection-target pair,
            with shape (num_detections, num_targets).
        """
        num_r = det_rz.shape[0]
        num_t = len(filters)
        metric_rt = self.get_rect_chunk(num_r, num_t)
        self.xx_rn[:num_r, self.xx_idx] = det_rz[:, self.ind_pos_w_rz_size]
        self.set_target_tz(filters)
        inv_rot_trk_t = Rotation.from_quat(self.xx_tn[:num_t, 3:7], scalar_first=True).inv()
        for r_num, det_n in enumerate(self.xx_rn[:num_r, :]):
            self.square_size_t3[:num_t] = self.xx_tn[:num_t, 7:]
            self.square_size_t3[:num_t] += det_n[7:]
            np.square(self.square_size_t3[:num_t], out=self.square_size_t3[:num_t])
            self.square_size_t3[:num_t] /= 4.0

            rotation_det = Rotation.from_quat(det_n[3:7], scalar_first=True)
            rel_rotation_t = rotation_det * inv_rot_trk_t
            self.euler_angles_t3[:num_t] = rel_rotation_t.as_euler('xyz')
            self.euler_angles_t3[:num_t] *= 2  # xxx the roll and pitch are zeros
            self.xx_t4[:num_t] = Rotation.from_euler('xyz', self.euler_angles_t3[:num_t]).as_quat(
                True, scalar_first=True
            )
            self.xx_tn[:num_t, :] -= det_n
            self.xx_tn[:num_t, 3:7] = self.xx_t4[:num_t]
            self.xx_tn[:num_t, 3] -= 1
            np.square(self.xx_tn[:num_t], out=self.xx_tn[:num_t])

            self.precision_dia_tn[:num_t, :3] = self.pos
            self.precision_dia_tn[:num_t, :3] /= self.square_size_t3[:num_t]
            self.precision_dia_tn[:num_t, 7:] = self.size
            self.precision_dia_tn[:num_t, 7:] /= self.square_size_t3[:num_t]

            self.xx_tn[:num_t] *= self.precision_dia_tn[:num_t]
            np.sum(self.xx_tn[:num_t], axis=1, out=metric_rt[r_num, :num_t])
            np.sqrt(metric_rt[r_num, :num_t], out=metric_rt[r_num, :num_t])
            metric_rt[r_num, :num_t] *= -1
            np.exp(metric_rt[r_num, :num_t], out=metric_rt[r_num, :num_t])

        return metric_rt
