"""."""

import numpy as np

from scipy.spatial.transform import Rotation
from shapely import Polygon

from kinematic_tracker.types import FMT


MID_VERTICES_C4 = ((0.5, 0.5, -0.5, -0.5), (-0.5, 0.5, 0.5, -0.5), (0.0, 0.0, 0.0, 0.0))


class RectangleBuffers:
    def __init__(self, num_max: int) -> None:
        self.unit_vertices_14c = np.array(MID_VERTICES_C4).T.reshape(1, 4, 3)
        self.num_max = num_max
        self.mid_vertices_n4c = np.zeros((num_max, 4, 3))
        self.rect_coordinates_n4c = np.zeros((num_max, 4, 3))
        self.mid_polygons_n = np.zeros(num_max, dtype=object)
        self.num_rect = 0

    def set_nu_scenes(self, vecs_nz: FMT) -> None:
        assert vecs_nz.ndim == 2 and vecs_nz.shape[1] == 10
        self.num_rect = n = vecs_nz.shape[0]
        if n == 0:
            return
        np.multiply(
            vecs_nz[:, 7:].reshape(n, 1, 3), self.unit_vertices_14c, out=self.mid_vertices_n4c[:n]
        )
        rotations = Rotation.from_quat(vecs_nz[:, 3:7].reshape(n, 1, 4), scalar_first=True)
        self.rect_coordinates_n4c[:n] = rotations.apply(self.mid_vertices_n4c[:n])
        self.rect_coordinates_n4c[:n] += vecs_nz[:, :3].reshape(n, 1, 3)
        self._mid_polygons()

    def _mid_polygons(self) -> None:
        for rect_num in range(self.num_rect):
            self.mid_polygons_n[rect_num] = Polygon(self.rect_coordinates_n4c[rect_num, :, :2])
