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

The metric uses the concept of Mahalanobis distance between two measurements $a_i$ and $b_j$:

$$D_ij = sqrt{(a_i - b_j) C^-1 (a_i - b_j)},$$

where the covariance $C$ is assumed to be diagonal matrix, and we omit the measurement
dimension indices for brevity. In order to imitate the IoU-based metrics, we modulate
the covariance with average size of the objects. The size of the objects should be a part
of measurement vectors. Instead of covariance, we parametrize its inverse (precision matrix)
$P$

$$D_ij = sqrt{(a_i - b_j)^2 * P_ij}.$$

The precision matrix $P_ij$ depends on the pair of the measurement vectors $ij$
and is diagonal. The measurement vectors $a$ and $b$ have three positional components,
two yaw quaternion components and three size components

$$
a = p_x, p_y, p_z, w, r_z, s_x, s_y, s_z.
$$

The yaw quaternion is vector of size 2 composed of scalar and z-component of the quaternion,
i.e. $w = cos(yaw / 2), r_z = sin(yaw / 2)$.

The diagonal precision matrix is scaled with inverse square sum sizes in its
position- and size parts with position and size calibration constants.
The quaternion part is given by a quaternion calibration constant.

Let's define a scale vector of three elements

$$scale_ij = 1 / (si_x + sj_x)**2, 1 / (si_y + sj_y)**2, 1 / (si_z + sj_z)**2.$$

The precision matrix expressed in terms of the scale vectors

$$
P_ij = pos_c * scale_ij, quat_c, quat_c, size_c * scale_ij,
$$

where $pos_c$, $quat_c$, and $size_c$, are calibration constants.

The square difference $(a - b)^2$ is computed element-wise for the positions and
sizes, while the quaternion part is computed as a wrapped double geodesic distance
$d_ij$ between two quaternions:

$$
(a - b)^2_ij = (pi_x - pj_x)^2, (pi_y - pj_y)^2, (pi_z - pj_z)^2,
               d_ij^2, d_ij^2,
               (si_x - sj_x)^2, (si_y - sj_y)^2, (si_z - sj_z)^2,
$$

where $d_ij$ is the wrapped double geodesic distance

$$
d_ij = 4 * arccos(|w_i * w_j + rz_i * rz_j|)
d_ij = d_ij if d_ij < pi else 2 pi - d_ij.
$$

The reason for wrapping is to imitate the behavior of IoU metric with
respect to rotation. Namely, the IoU metrics give the same values for
yaw angle $theta$ and $theta + pi$.

Using the Mahalanobis distance $D_ij$, the association metric $A_ij$ is
computed using the negative exponent (not Gaussian) of the distance

$$
A_ij = exp(-D_ij).
$$

The choice of exponent function instead of the bell shaped dependency
is for the sake of better mimicking the IoU metric. Namely, the IoU
metrics decrease faster at small shifts than the Gaussian.

We calibrate the metric using three reference pairs of points:

  - for translation: a = (3,0,0,1,0,1,1,1), b = (0,0,0,1,0,1,1,1)
  - for rotation: a = (0,0,0,1,0,4,2,1), b = (0,0,0,sqrt{2}/2,sqrt{2}/2,4,2,1)
  - for scale: a = (0,0,0,1,0,1,1,1), b = (0,0,0,1,0,2,2,1)

The condition for the calibration constants $pos_c$, $quat_c$ and $size_c$ is

$$
A_ij = GIoU_ij.
$$

Classes:
    MetricIouLikeTwoYaw: a driver for the two-yaw association metric.
"""

from typing import Sequence

import numpy as np

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.det_r8 = np.zeros((num_reports_max, 8))
        self.trk_t8 = np.zeros((num_targets_max, 8))
        self.tmp_m2 = np.zeros((max(num_targets_max, num_reports_max), 2))
        self.pos = 0.2136 * 4
        self.size = 1.0
        num_rt_max = num_reports_max * num_targets_max
        self.mask_rt = np.zeros(num_rt_max, bool)
        self.square_size_rt3 = np.zeros(num_rt_max * 3)
        self.square_diff_rt8 = np.zeros((num_rt_max * 8))
        self.precision_dia_rt8 = np.zeros((num_rt_max * 8))
        self.precision_dia_rt8[:] = 0.00682 * 2

    def get_det_t8(self, det_rz: FMT) -> FMT:
        det_r8 = self.det_r8[: det_rz.shape[0]]
        det_r8[:] = det_rz[:, self.ind_pos_w_rz_size]
        return det_r8

    def get_target_t8(self, filters: FT) -> FMT:
        for t, kf in enumerate(filters):
            np.dot(kf.measurementMatrix, kf.statePre, out=self.vec_z1)
            self.trk_t8[t] = self.vec_z1[self.ind_pos_w_rz_size, 0]
        return self.trk_t8[: len(filters), :]

    def normalize_yaw_quat_inplace(self, wrz_m2: FMT) -> None:
        """Normalize the yaw quaternion inplace.

        A _yaw quaternion_ is a pair of numbers (scalar, rz).
        """
        tmp_m2 = self.tmp_m2[: wrz_m2.shape[0], :]
        tmp_m2[:, :] = wrz_m2[:, :]
        np.square(tmp_m2, out=tmp_m2)
        np.sum(tmp_m2, axis=1, out=tmp_m2[:, 0])
        np.sqrt(tmp_m2[:, 0], out=tmp_m2[:, 0])
        np.divide(wrz_m2, tmp_m2[:, :1], out=wrz_m2)

    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)
        if num_r == 0:
            return metric_rt
        det_r8 = self.get_det_t8(det_rz)
        trk_t8 = self.get_target_t8(filters)
        self.normalize_yaw_quat_inplace(det_r8[:, 3:5])
        self.normalize_yaw_quat_inplace(trk_t8[:, 3:5])

        num_rt = num_r * num_t
        dist_yaw_rt = metric_rt
        mask_rt = self.mask_rt[:num_rt].reshape(num_r, num_t)
        square_size_rt3 = self.square_size_rt3[: num_rt * 3].reshape(num_r, num_t, 3)
        square_diff_rt8 = self.square_diff_rt8[: num_rt * 8].reshape(num_r, num_t, 8)
        precision_dia_rt8 = self.precision_dia_rt8[: num_rt * 8].reshape(num_r, num_t, 8)

        det_r8[:, 3:5].dot(trk_t8[:, 3:5].T, out=dist_yaw_rt)
        np.abs(dist_yaw_rt, out=dist_yaw_rt)
        dist_yaw_rt -= 1.0e-14
        np.acos(dist_yaw_rt, out=dist_yaw_rt)
        dist_yaw_rt *= 4.0
        np.greater(dist_yaw_rt, np.pi, out=mask_rt)
        dist_yaw_rt[mask_rt] *= -1
        dist_yaw_rt[mask_rt] += 2 * np.pi

        np.add(
            det_r8[:, 5:8].reshape(num_r, 1, 3),
            trk_t8[:, 5:8].reshape(1, num_t, 3),
            out=square_size_rt3,
        )
        np.square(square_size_rt3, out=square_size_rt3)

        np.subtract(det_r8.reshape(num_r, 1, 8), trk_t8.reshape(1, num_t, 8), out=square_diff_rt8)
        square_diff_rt8[:, :, 3:5] = dist_yaw_rt.reshape(num_r, num_t, 1)
        np.square(square_diff_rt8, out=square_diff_rt8)

        np.divide(self.pos, square_size_rt3, out=precision_dia_rt8[:, :, :3])
        np.divide(self.size, square_size_rt3, out=precision_dia_rt8[:, :, 5:])
        np.multiply(square_diff_rt8, precision_dia_rt8, out=square_diff_rt8)
        np.sum(square_diff_rt8, axis=2, out=metric_rt)
        np.sqrt(metric_rt, out=metric_rt)
        metric_rt *= -1
        np.exp(metric_rt, out=metric_rt)
        return metric_rt
