import numpy as np

def _sample_points_in_box3d(bbox_vertices, num_samples):
    """
    Sample points in a axis-aligned 3D bounding box\n
    - bbox_vertices: the vertices of the bounding box in the form: [[x0, y0, z0], [x1, y1, z1], ...]\n
    - num_samples: the number of samples to use\n

    Return:\n
    - points: the sampled points in the form: [[x0, y0, z0], [x1, y1, z1], ...]
    """

    # Compute the bounding box size
    bbox_size = np.max(bbox_vertices, axis=0) - np.min(bbox_vertices, axis=0)

    # Sample points in the bounding box
    points = np.random.rand(num_samples, 3) * bbox_size + np.min(bbox_vertices, axis=0)

    return points

def _apply_forward_transformations(points, transformations):
    """
    Apply forward transformations to the points\n
    - points: the points in the form: [[x0, y0, z0], [x1, y1, z1], ...]\n
    - transformations: list of transformations to apply\n

    Return:\n
    - points_transformed: the transformed points in the form: [[x0, y0, z0], [x1, y1, z1], ...]
    """

    # To homogeneous coordinates
    points_transformed = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1)

    # Apply the transformations one by one in order
    for transformation in transformations:
        if transformation["type"] == "translation":
            points_transformed = np.matmul(transformation["matrix"], points_transformed.T).T

        elif transformation["type"] == "rotation":
            axis_origin = np.append(transformation["rotation_axis_origin"], 0)
            points_recentered = points_transformed - axis_origin

            points_rotated = np.matmul(transformation["matrix"], points_recentered.T).T
            points_transformed = points_rotated + axis_origin

        elif transformation["type"] == "plucker":
            points_transformed = np.matmul(transformation["matrix"], points_transformed.T).T

        else:
            raise ValueError(f"Unknown transformation type: {transformation['type']}")

    return points_transformed[...,:3]

def _apply_backward_transformations(points, transformations):
    """
    Apply backward transformations to the points\n
    - points: the points in the form: [[x0, y0, z0], [x1, y1, z1], ...]\n
    - transformations: list of transformations to apply\n
        - The inverse of the transformations are applied in reverse order\n

    Return:\n
    - points_transformed: the transformed points in the form: [[x0, y0, z0], [x1, y1, z1], ...]

    Reference: https://mathematica.stackexchange.com/questions/106257/how-do-i-get-the-inverse-of-a-homogeneous-transformation-matrix
    """

    # To homogeneous coordinates
    points_transformed = np.concatenate([points, np.ones((points.shape[0], 1))], axis=1)

    # Apply the transformations one by one in reverse order
    for transformation in transformations[::-1]:
        inv_transformation = np.eye(4)
        inv_transformation[:3, :3] = transformation["matrix"][:3, :3].T
        inv_transformation[:3, 3] = -np.matmul(transformation["matrix"][:3, :3].T, transformation["matrix"][:3, 3])

        if transformation["type"] == "translation":
            points_transformed = np.matmul(inv_transformation, points_transformed.T).T

        elif transformation["type"] == "rotation":
            axis_origin = np.append(transformation["rotation_axis_origin"], 0)
            points_recentered = points_transformed - axis_origin

            points_rotated = np.matmul(inv_transformation, points_recentered.T).T
            points_transformed = points_rotated + axis_origin

        elif transformation["type"] == "plucker":
            points_transformed = np.matmul(inv_transformation, points_transformed.T).T

        else:
            raise ValueError(f"Unknown transformation type: {transformation['type']}")

    return points_transformed[...,:3]

def _count_points_in_box3d(points, bbox_vertices):
    """
    Count the number of points in a 3D bounding box\n
    - points: the points in the form: [[x0, y0, z0], [x1, y1, z1], ...]\n
    - bbox_vertices: the vertices of the bounding box in the form: [[x0, y0, z0], [x1, y1, z1], ...]\n
        - The bbox is assumed to be axis-aligned\n

    Return:\n
    - num_points_in_bbox: the number of points in the bounding box
    """

    # Count the number of points in the bounding box
    num_points_in_bbox = np.sum(np.all(points >= np.min(bbox_vertices, axis=0), axis=1) & np.all(points <= np.max(bbox_vertices, axis=0), axis=1))

    return num_points_in_bbox

def sampling_iou(bbox1_vertices, bbox2_vertices, bbox1_transformations, bbox2_transformations, num_samples=10000):
    """
    Compute the IoU between two bounding boxes\n
    - bbox1_vertices: the vertices of the first bounding box\n
    - bbox2_vertices: the vertices of the second bounding box\n
    - bbox1_transformations: list of transformations applied to the first bounding box\n
    - bbox2_transformations: list of transformations applied to the second bounding box\n
    - num_samples (optional): the number of samples to use per bounding box\n

    Return:\n
    - iou: the IoU between the two bounding boxes after applying the transformations
    """

    # Volume of the two bounding boxes
    bbox1_volume = np.prod(np.max(bbox1_vertices, axis=0) - np.min(bbox1_vertices, axis=0))
    bbox2_volume = np.prod(np.max(bbox2_vertices, axis=0) - np.min(bbox2_vertices, axis=0))

    # Sample points in the two bounding boxes
    bbox1_points = _sample_points_in_box3d(bbox1_vertices, num_samples)
    bbox2_points = _sample_points_in_box3d(bbox2_vertices, num_samples)

    # Transform the points
    forward_bbox1_points = _apply_forward_transformations(bbox1_points, bbox1_transformations)
    forward_bbox2_points = _apply_forward_transformations(bbox2_points, bbox2_transformations)    

    # Transform the forward points to the other box's rest pose frame
    forward_bbox1_points_in_rest_bbox2_frame = _apply_backward_transformations(forward_bbox1_points, bbox2_transformations)
    forward_bbox2_points_in_rest_bbox1_frame = _apply_backward_transformations(forward_bbox2_points, bbox1_transformations)

    # Count the number of points in the other bounding box
    num_bbox1_points_in_bbox2 = _count_points_in_box3d(forward_bbox1_points_in_rest_bbox2_frame, bbox2_vertices)
    num_bbox2_points_in_bbox1 = _count_points_in_box3d(forward_bbox2_points_in_rest_bbox1_frame, bbox1_vertices)
    
    # Compute the IoU
    intersect = (bbox1_volume * num_bbox1_points_in_bbox2 + bbox2_volume * num_bbox2_points_in_bbox1) / 2
    union = bbox1_volume * num_samples + bbox2_volume * num_samples - intersect
    iou = intersect / union
    
    return iou