import numpy as np
import scipy.spatial.transform

from association_studio.common import FMT, FVT


MID_CORNERS_CN = ((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 Box:
    """Bounding box."""

    def __init__(self, vec_z: FVT) -> None:
        self._rotation = scipy.spatial.transform.Rotation.from_quat(vec_z[3:7], scalar_first=True)
        self.vec_z = vec_z
        self.unit_corners = np.array(MID_CORNERS_CN)
        self.scaled_corners = np.zeros_like(self.unit_corners)
        self.corners_buf = np.zeros_like(self.unit_corners)
        self._wlh31 = np.zeros((3, 1))

    def set_size(self, width: float, length: float, height: float) -> None:
        self.vec_z[7:] = width, length, height

    def translate(self, x: FVT) -> None:
        self.vec_z[:3] += x

    def rotate(self, wxyz: FVT) -> None:
        self._rotation *= scipy.spatial.transform.Rotation.from_quat(wxyz, scalar_first=True)
        self.vec_z[3:7] = self._rotation.as_quat(scalar_first=True, canonical=True)

    def mid_corners(self) -> FMT:
        self._wlh31[:, 0] = self.vec_z[7:]
        self.scaled_corners[:] = self.unit_corners
        self.scaled_corners *= self._wlh31
        np.dot(self._rotation.as_matrix(), self.scaled_corners, out=self.corners_buf)
        self.corners_buf += self.vec_z[:3].reshape(3, 1)
        return self.corners_buf
