import gradio as gr
import open3d as o3d
import numpy as np
import pandas as pd
from PIL import Image
from gradio.components import File, Slider, Image, Dataframe, Model3D
import tempfile

def point_cloud_to_image(point_cloud):
    # Create Open3D visualization object
    vis = o3d.visualization.Visualizer()
    vis.create_window()
    vis.add_geometry(point_cloud)

    # Render the visualization and get the image
    vis.poll_events()
    vis.update_renderer()
    img = vis.capture_screen_float_buffer()

    # Convert image to PIL Image format
    pil_img = Image.fromarray(np.uint8(np.asarray(img)*255))

    return pil_img

def create_file_mesh_from_pcd(pcd, alpha): 
    
    temp_file = tempfile.NamedTemporaryFile(suffix=".obj", delete=False)
    file_path = temp_file.name
    temp_file.close()
    
    # print(f"alpha={alpha:.3f}")
    mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_alpha_shape(pcd, alpha)
    mesh.compute_vertex_normals()
    o3d.io.write_triangle_mesh(file_path, mesh, write_triangle_uvs=True)
    return file_path

def plane_detection(pointcloud_path, voxel_size=0.05, distance_threshold=0.01, num_iterations=1000, alpha=2):
    # Load point cloud from file
    pcd = o3d.io.read_point_cloud(pointcloud_path.name)

    # Downsample the point cloud to reduce computational load
    pcd = pcd.voxel_down_sample(voxel_size=voxel_size)

    # Find plane using RANSAC algorithm
    plane_model, inliers = pcd.segment_plane(distance_threshold=distance_threshold,
                                             ransac_n=3,
                                             num_iterations=num_iterations)

    # Get inlier and outlier point cloud
    inlier_cloud = pcd.select_by_index(inliers)
    outlier_cloud = pcd.select_by_index(inliers, invert=True)
    # extract the coefficients of the plane
    a, b, c, d = plane_model
    plane_model = np.array([[a], [b], [c], [d]])
    df = pd.DataFrame(plane_model.reshape(1, 4), columns=["a", "b", "c", "d"])
    input_path = create_file_mesh_from_pcd(pcd, alpha)
    inlier_path = create_file_mesh_from_pcd(inlier_cloud, alpha)
    outlier_path = create_file_mesh_from_pcd(outlier_cloud, alpha)
    # Return inlier point cloud, outlier point cloud, and plane model
    # return point_cloud_to_image(inlier_cloud), point_cloud_to_image(outlier_cloud), df
    return input_path, inlier_path, outlier_path, df

outputs = [
    # show pcd inlier point cloud
    Model3D(label="Input Cloud", clear_color=[1.0, 1.0, 1.0, 1.0]),
    # show pcd inlier point cloud
    Model3D(label="Inlier Cloud", clear_color=[1.0, 1.0, 1.0, 1.0]),
    # show pcd outlier point cloud
    Model3D(label="Outlier Cloud", clear_color=[1.0, 1.0, 1.0, 1.0]),
    # show the centroids and counts, which is a numpy array
    Dataframe(label="Coefficients of the plane model", type="pandas")
]

# Create Gradio interface
iface = gr.Interface(plane_detection, 
                     inputs=[
                         File(label="Point cloud file (.ply or .pcd format)"),
                         Slider(label="Voxel size", minimum=0.001, maximum=50, step=1, value=2),
                         Slider(label="Distance threshold", minimum=0.001, maximum=50, step=0.01, value=5),
                         Slider(label="Number of iterations", minimum=1, maximum=10000, step=1, value=100),
                         Slider(label="Alpha for surface reconstruction", minimum=0.02, maximum=100, step=0.01, value=2),
                     ],
                     outputs=outputs,
                     title="Plane Detection using RANSAC",
                     description="This app takes as input a point cloud file (.ply or .pcd format), voxel size, distance threshold, and number of iterations, finds a plane using RANSAC algorithm, displays the inlier and outlier point clouds, and returns the inlier point cloud, outlier point cloud, and the plane model.",
                     allow_flagging="never",
                     examples=[
                        # ["Pointclouds/p1.ply", 2, 5, 100, 2]
                        ["Pointclouds/cloud_bin_0.ply", 0.001, 0.03, 100, 0.05],
                        ["Pointclouds/cloud_bin_20.ply", 0.001, 0.03, 100, 0.05],
                        ["Pointclouds/cloud_bin_30.ply", 0.001, 0.03, 100, 0.05],
                        ["Pointclouds/cloud_bin_40.ply", 0.001, 0.03, 100, 0.05],
                        ["Pointclouds/cloud_bin_50.ply", 0.001, 0.03, 100, 0.05],
                        ],
                     )


# Launch the interface
iface.launch()