import numpy as np
import trimesh
from sklearn import linear_model
import copy
import math

def distance_3d(a, b):
    return np.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2 + (a[2] - b[2]) ** 2)

class CONVEX_HULL:
    def __init__(self, points) -> None:
        self.average_z = 0
        self.highest_z = -np.inf
        self.convex_hull = self.polygon_scan(points)
        self.calc_average_z()
        self.calc_highest_z()

    def distance(self, a, b):
        return np.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2)

    def cross_product(self, a, b, c):
        return (b[1]-a[1])*(c[0]-b[0]) - (b[0]-a[0])*(c[1]-b[1])

    def dot_product(self, a, b, c):
        ab = [b[0] - a[0], b[1] - a[1]]
        bc = [c[0] - b[0], c[1] - b[1]]
        return ab[0] * bc[0] + ab[1] * bc[1]

    def angle(self, a, b, c):
        cp = self.cross_product(a, b, c)
        dp= self.dot_product(a, b, c)
        angle_radians = math.atan2(cp, dp)
        angle_degrees = math.degrees(angle_radians)

        return angle_degrees

    def new_point_orientation(self, pivot, p1, p2):
        orientation = self.cross_product(pivot, p1, p2)
        if orientation == 0:
            return 0
        elif(orientation > 0):
            return 1
        else:
            return -1

    def polygon_scan(self, points_in):
        points = copy.deepcopy(points_in)
        n = len(points)
        if n < 3:
            print("Convec Hull required at least 3 points.")
            return
        else:
            if(n==3 and self.new_point_orientation[points[0],points[1],points[2]]==0):
                print("3 points are in one line.")
                return
            else:
                pivot_index, pivot_point = min(enumerate(points), key=lambda point: point[1]) # 优先根据y值排序，若y相同，则找到相同y下x值最小的点
                start_point = copy.deepcopy(pivot_point)
                start_point[1] -= 1
                convex_hull = [start_point, pivot_point]
                closed = False

                while not closed:
                    max_x = -np.inf
                    max_x_idx = 0
                    for p_idx in range(len(points)):
                        point = points[p_idx]
                        x = self.angle(convex_hull[-2], convex_hull[-1], point)
                        if(x > max_x and x != 180 and x != 0):
                            max_x = x
                            max_x_idx = p_idx
                    if(len(points) == 0):
                        closed = True
                    else:
                        new_point = points[max_x_idx]
                        if(new_point[0] == pivot_point[0] and new_point[1] == pivot_point[1]):
                            closed = True
                        else:
                            convex_hull.append(new_point)
                            points.pop(max_x_idx)

                convex_hull.pop(0)
                return convex_hull

    def is_point_inside_convex_hull(self, point_in):
        point = copy.deepcopy(point_in)
        new_convex_hull_points = copy.deepcopy(self.convex_hull)
        new_convex_hull_points.append(point)
        new_convex_hull = self.polygon_scan(new_convex_hull_points)
        ans = (new_convex_hull == self.convex_hull)
        return ans

    def are_points_inside_convex_hull(self, points_in):
        points = copy.deepcopy(points_in)
        ans = True
        for point in points:
            p_ans =  self.is_point_inside_convex_hull(point)
            ans = ans * p_ans
        return ans

    def get_convex_hull(self):
        return self.convex_hull
    
    def calc_average_z(self):
        for p in self.convex_hull:
            self.average_z += p[2]/len(self.convex_hull)

    def get_average_z(self):
        return self.average_z

    def calc_highest_z(self):
        for p in self.convex_hull:
            if(p[2]>self.highest_z):
                self.highest_z = p[2]

    def get_highest_z(self):
        return self.highest_z

    

class Object:
    def __init__(self, obj_path, usd_path=None, scale=[1, 1, 1], eular_angles=[0, 0, 0], max_num_planes=2,
                 label=None, name=None, ) -> None:
        # trimesh
        self.obj_path = obj_path
        self.usd_path = usd_path
        self.scale = scale
        self.eular_angles = eular_angles
        self.label = label
        self.name = name
        self.bb = None

        # get placable points
        self.placable_points = []
        
        # ransac
        self.max_num_planes = max_num_planes
        self.X = []
        self.y = []
        self.plane_point_groups = []
        self.plane_coef_groups = []

        # convex_hull
        self.convex_hulls = []
        self.convex_hulls_z = []
        
    def _load_mesh(self):
        # Load mesh
        self._mesh = trimesh.load(self.obj_path)
        self._mesh.apply_scale(self.scale)

        # Get Bounding Box
        self.bb = self._mesh.bounds

        # Rotate Mesh
        self._rotate_mesh(self.eular_angles)

    def _find_placable_areas(self):
        # Find placable points
        self._get_placable_points()
        self._convert_placable_points_to_sklearn_X_y()
        self._get_plane_point_groups()
        self._get_convex_hulls()

    def _rotate_mesh(self, rot_rpy):
        center = [0, 0, 0]
        angle = rot_rpy[0]
        direction = [1, 0, 0]
        rot_matrix_r = trimesh.transformations.rotation_matrix(angle, direction, center)
        self._mesh.apply_transform(rot_matrix_r)
        angle = rot_rpy[1]
        direction = [0, 1, 0]
        rot_matrix_p = trimesh.transformations.rotation_matrix(angle, direction, center)
        self._mesh.apply_transform(rot_matrix_p)
        angle = rot_rpy[2]
        direction = [0, 0, 1]
        rot_matrix_y = trimesh.transformations.rotation_matrix(angle, direction, center)
        self._mesh.apply_transform(rot_matrix_y)

    def _get_placable_points(self):
        gravity=np.array([0, 0, -1.0])
        support_facet_indices = np.argsort(self._mesh.facets_area)
        for idx in support_facet_indices:
            if(np.isclose(self._mesh.facets_normal[idx].dot(gravity), 1.0, atol=0.15)):
                facet_boundary = self._mesh.facets_boundary[idx]
                for bondary_vertices_indices in facet_boundary:
                    for boundary_vertex_index in bondary_vertices_indices:
                        point = self._mesh.vertices[boundary_vertex_index]
                        self.placable_points.append(list(point))
        return self.placable_points
    
    def _convert_placable_points_to_sklearn_X_y(self):
        for p in self.placable_points:
            self.X.append([p[0], p[1]])
            self.y.append(p[2])
        self.X = np.array(self.X)
        self.y = np.array(self.y)
        return self.X, self.y
    
    def _get_plane_point_groups(self):
        ransac = linear_model.RANSACRegressor()
        num_planes = 0
        while(num_planes < self.max_num_planes):
            plane_points = []
            plane_coefs = []
            num_popped_element = 0
            if(len(self.X) <= 2):
                break
            else:
                ransac.fit(self.X, self.y)
                inlier_mask = ransac.inlier_mask_ # a list of bool values indicating whether the element is inlier point
                plane_coefficients = ransac.estimator_.coef_
                plane_intercept = ransac.estimator_.intercept_
                plane_coefs = [plane_coefficients[0], plane_coefficients[1], plane_intercept]
                for i in range(len(self.X)):
                    if(len(self.X) > 0 and inlier_mask[i] == True):
                        X = self.X[i - num_popped_element]
                        y = self.y[i - num_popped_element]
                        plane_points.append([X[0], X[1], y])
                        self.X = np.delete(self.X, i-num_popped_element, 0)
                        self.y = np.delete(self.y, i-num_popped_element, 0)
                        num_popped_element+=1

                self.plane_point_groups.append(plane_points)
                self.plane_coef_groups.append(plane_coefs)
                num_planes += 1
        return self.plane_point_groups, self.plane_coef_groups
    
    def _get_convex_hulls(self):
        for plane_points in self.plane_point_groups:
            ch = CONVEX_HULL(plane_points)
            self.convex_hulls.append(ch)

class BOUNDING_BOX_2D:
    def __init__(self, origin, width, height, rot_z=0) -> None:
        self.origin = origin
        self.origin_x = origin[0]
        self.origin_y = origin[1]
        self.width = width
        self.height = height
        self.rot_z = rot_z

        self.setup()

    def setup(self):
        self._get_bb_vertices()
        self._apply_transform_to_bb_vertices()

    def _get_bb_vertices(self):
        self.bb_vertics = [[self.origin_x-self.width/2, self.origin_y-self.height/2],
                           [self.origin_x-self.width/2, self.origin_y+self.height/2],
                           [self.origin_x+self.width/2, self.origin_y-self.height/2],
                           [self.origin_x+self.width/2, self.origin_y+self.height/2],]
        return self.bb_vertics

    def _apply_transform_to_bb_vertices(self):
        origin = np.array([[1, 0, 0, self.origin_x],
                           [0, 1, 0, self.origin_y],
                           [0, 0, 1, 0],
                           [0, 0, 0, 1]])
        T = np.array([[np.cos(self.rot_z), -np.sin(self.rot_z), 0, 0],
                      [np.sin(self.rot_z), np.cos(self.rot_z),  0, 0],
                      [0,                  0,                   1, 0],
                      [0,                  0,                   0, 1]])
        T_left_up = np.array([[1, 0, 0, -self.width/2],
                              [0, 1, 0, self.height/2],
                              [0, 0, 1, 0],
                              [0, 0, 0, 1]])
        T_left_down = np.array([[1, 0, 0, -self.width/2],
                                [0, 1, 0, -self.height/2],
                                [0, 0, 1, 0],
                                [0, 0, 0, 1]])
        T_right_up = np.array([[1, 0, 0, self.width/2],
                               [0, 1, 0, self.height/2],
                               [0, 0, 1, 0],
                               [0, 0, 0, 1]])
        T_right_down = np.array([[1, 0, 0, self.width/2],
                               [0, 1, 0, -self.height/2],
                               [0, 0, 1, 0],
                               [0, 0, 0, 1]])
        transformed_origin = origin.dot(T)
        tv0 = transformed_origin.dot(T_left_down)
        tv1 = transformed_origin.dot(T_left_up)
        tv2 = transformed_origin.dot(T_right_down)
        tv3 = transformed_origin.dot(T_right_up)
        transformed_vertices = [tv0, tv1, tv2, tv3]
    
        for vertex_idx in range(len(self.bb_vertics)):
            self.bb_vertics[vertex_idx][0] = transformed_vertices[vertex_idx][0][3]
            self.bb_vertics[vertex_idx][1] = transformed_vertices[vertex_idx][1][3]