"""The module for computing Generalized Intersection over Union (GIoU) association metric.

This module provides functionality for computing the similarity score between detection reports
and Kalman filter states (targets) using the Generalized Intersection over Union (GIoU) metric.
The metric is computed between the 3D bounding boxes which are oriented by yaw angle.

Classes:
    MetricGIoU: GIoU association metrics between detections and Kalman filters.
"""

from typing import Sequence

import numpy as np

from kinematic_tracker.geometry.giou_yaw_aux import GiouYawAux
from kinematic_tracker.geometry.rectangle_buffers import RectangleBuffers
from kinematic_tracker.types import FMT, FT

from .metric_driver_base import MetricDriverBase


class MetricGIoUYaw(MetricDriverBase):
    """Computes GIoU association metric between detections and Kalman filter states.

    This class inherits from MetricDriverBase and implements GIoU-based metric computation
    for data association between detection reports and Kalman filter states. The computation
    is performed between cuboids (3D bounding boxes) rotated by yaw angle (around z axis
    by default).

    Args:
        num_reports_max: Maximum number of detection reports to handle.
        num_targets_max: Maximum number of targets (Kalman filters) to handle.
        num_z: Dimension of the measurement vector.
    """

    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.report_buf = RectangleBuffers(num_reports_max)
        self.target_buf = RectangleBuffers(num_targets_max)
        self.xx_tz = np.zeros((num_targets_max, 10))
        self.xx_rz = np.zeros((num_reports_max, 10))
        self.xx_idx = np.array((0, 1, 2, 3, 6, 7, 8, 9), int)
        self.volumes_t = np.zeros(num_targets_max)
        self.aux = GiouYawAux(num_targets_max)

    def set_det_rz(self, det_rz: FMT) -> None:
        num_r = det_rz.shape[0]
        self.xx_rz[:num_r, self.xx_idx] = det_rz[:, self.ind_pos_w_rz_size]
        self.report_buf.set_nu_scenes(self.xx_rz[:num_r])

    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_tz[t, self.xx_idx] = self.vec_z[self.ind_pos_w_rz_size, 0]
        num_t = len(filters)
        self.target_buf.set_nu_scenes(self.xx_tz[:num_t])
        self.volumes_t[:num_t] = self.xx_tz[:num_t, 7:].prod(axis=1)

    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:
            np.ndarray: A matrix of GIoU scores between each detection-target pair,
                with shape (num_detections, num_targets).
        """
        self.set_det_rz(det_rz)
        self.set_target_tz(filters)
        num_r = det_rz.shape[0]
        num_t = len(filters)
        metric_rt = self.get_rect_chunk(num_r, num_t)
        trk_tz = self.xx_tz[:num_t]
        trk_poly_t = self.target_buf.mid_polygons_n[:num_t]
        volumes_t = self.volumes_t[:num_t]
        self.aux.set_second_4_vertices(self.target_buf.rect_coordinates_n4c[:num_t, :, :2])

        for r_num in range(num_r):
            det_z = self.xx_rz[r_num]
            det_poly = self.report_buf.mid_polygons_n[r_num]
            rect_4c = self.report_buf.rect_coordinates_n4c[r_num, :, :2]
            self.aux.set_first_4_vertices(rect_4c, num_t)
            self.aux.compute(det_z, det_poly, trk_tz, trk_poly_t, volumes_t, metric_rt[r_num])
        return metric_rt
