import numpy as np
import shapely

from shapely import Polygon

from kinematic_tracker.geometry.rectangle_buffers import SPT
from kinematic_tracker.types import F3T, FMT, FVT


class GiouYawAux:
    def __init__(self, num_max: int) -> None:
        self.merged_p8c = np.zeros((num_max, 8, 2))
        self.multi_points_p = np.zeros(num_max, object)
        self.convex_hulls_p = np.zeros(num_max, object)
        self.hull_areas_p = np.zeros(num_max)
        self.intersection_areas_p = np.zeros(num_max)
        self.intersection_heights_p = np.zeros(num_max)
        self.union_heights_p = np.zeros(num_max)
        self.intersection_volumes_p = np.zeros(num_max)
        self.union_volumes_p = np.zeros(num_max)
        self.convex_volumes_p = np.zeros(num_max)
        self.za_m_zb_p = np.zeros(num_max)
        self.averaged_heights_p = np.zeros(num_max)
        self.stacked_4p = np.zeros((4, num_max))

    def comp_heights(self, cuboid: FVT, cuboids: FMT) -> tuple[FVT, FVT]:
        n = cuboids.shape[0]
        intersections = self.intersection_heights_p[:n]
        unions = self.union_heights_p[:n]
        za_m_zb = self.za_m_zb_p[:n]
        za_m_zb[:] = cuboid[2]
        za_m_zb -= cuboids[:, 2]
        stacked_4p = self.stacked_4p[:, :n]
        stacked_4p[3] = cuboids[:, -1]
        averaged_heights = self.averaged_heights_p[:n]
        averaged_heights[:] = cuboid[-1]
        averaged_heights += stacked_4p[3, :]
        averaged_heights /= 2.0
        stacked_4p[0] = averaged_heights
        stacked_4p[0] += za_m_zb
        stacked_4p[1] = averaged_heights
        stacked_4p[1] -= za_m_zb
        stacked_4p[2] = cuboid[-1]
        np.amin(stacked_4p, axis=0, out=intersections)
        np.maximum(intersections, 0, out=intersections)
        np.amax(stacked_4p, axis=0, out=unions)
        return intersections, unions

    def set_first_4_vertices(self, rect_4c: FMT, n: int) -> None:
        self.merged_p8c[:n, :4] = rect_4c

    def set_second_4_vertices(self, rect_t4c: F3T) -> None:
        n = rect_t4c.shape[0]
        self.merged_p8c[:n, 4:] = rect_t4c

    def comp_convex_hull_areas(self, n: int) -> FMT:
        shapely.multipoints(self.merged_p8c[:n], out=self.multi_points_p[:n])
        shapely.convex_hull(self.multi_points_p[:n], out=self.convex_hulls_p[:n])
        shapely.area(self.convex_hulls_p[:n], out=self.hull_areas_p[:n])
        return self.hull_areas_p[:n]

    def compute(
        self, cuboid: FVT, poly: Polygon, cuboids: FMT, polygons: SPT, volumes: FVT, out: FVT
    ) -> None:
        n = polygons.shape[0]
        shapely.area(shapely.intersection(poly, polygons), out=self.intersection_areas_p[:n])
        ch_areas = self.comp_convex_hull_areas(n)
        intersection_heights, union_heights = self.comp_heights(cuboid, cuboids)

        intersection_volumes = self.intersection_volumes_p[:n]
        intersection_volumes[:] = self.intersection_areas_p[:n]
        intersection_volumes *= intersection_heights

        union_volumes = self.union_volumes_p[:n]
        union_volumes[:] = cuboid[-3:].prod()
        union_volumes += volumes[:]
        union_volumes -= intersection_volumes
        convex_vol = self.convex_volumes_p[:n]
        np.multiply(ch_areas, union_heights, out=convex_vol)

        tmp_vol = union_volumes

        out[:] = intersection_volumes
        out /= tmp_vol
        tmp_vol /= convex_vol
        out += tmp_vol
        out /= 2.0
