"""."""

import numpy as np
import shapely

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

from kinematic_tracker.types import FMT


SPT = np.ndarray[tuple[int], np.dtype[np.object_]]
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: SPT = np.zeros(num_max, dtype=Polygon)
        self.rotations = Rotation.identity(shape=(num_max, 1))
        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
        sizes_ptr = vecs_nz[:, 7:].reshape(n, 1, 3)
        np.multiply(sizes_ptr, self.unit_vertices_14c, out=self.mid_vertices_n4c[:n])
        rotations_ptr = self.rotations[:n]
        rotations_ptr[:] = Rotation.from_quat(vecs_nz[:, 3:7].reshape(n, 1, 4), scalar_first=True)
        self.rect_coordinates_n4c[:n] = rotations_ptr.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:
        n = self.num_rect
        shapely.polygons(self.rect_coordinates_n4c[:n, :, :2], out=self.mid_polygons_n[:n])
