"""."""

import numpy as np

from kinematic_tracker.geometry.rectangle_buffers import RectangleBuffers
from kinematic_tracker.geometry.slow.giou_yaw_shapely import giou_3d_shapely
from kinematic_tracker.types import FMT, FT


class NuScenesGiou:
    def __init__(self, num_reports_max: int, num_targets_max: int) -> None:
        self.report_buf = RectangleBuffers(num_reports_max)
        self.target_buf = RectangleBuffers(num_targets_max)
        self.vec_tz = np.zeros((num_targets_max, 10))

    def set_det_rz(self, det_rz: FMT) -> None:
        self.report_buf.set_nu_scenes(det_rz)

    def set_target_tz(self, filters: FT) -> None:
        for t, kf in enumerate(filters):
            np.dot(kf.measurementMatrix, kf.statePre[:, 0], out=self.vec_tz[t])
        num_t = len(filters)
        self.target_buf.set_nu_scenes(self.vec_tz[:num_t])

    def compute_giou(self, det_rz: FMT, filters: FT, metric_rt: FMT) -> None:
        self.set_det_rz(det_rz)
        self.set_target_tz(filters)
        for r_num in range(self.report_buf.num_rect):
            det_z = det_rz[r_num]
            det_poly = self.report_buf.mid_polygons_n[r_num]
            for t_num in range(self.target_buf.num_rect):
                trk_z = self.vec_tz[t_num]
                trk_poly = self.target_buf.mid_polygons_n[t_num]
                metric_rt[r_num, t_num] = giou_3d_shapely(det_z, det_poly, trk_z, trk_poly)
