import open3d as o3d
import numpy as np
import torch
import cv2
from scipy.spatial import ConvexHull

def read_pts_file(pts_file):
    points = []
    with open(pts_file, 'r') as f:
        for line in f:
            values = line.split()
            if len(values) > 4:  # Check if there are at least 5 values in the line
                points.append([float(values[0]), float(values[1]), float(values[2])])

    # Convert to a NumPy array
    points = np.asarray(points)

    # Calculate the average Z value
    avg_z = np.mean(points[:, 2])

    # Filter out points with a Z value higher than the average
    filtered_points = points[points[:, 2] <= avg_z]

    return filtered_points
def apply_statistical_outlier_removal(pcd, nb_neighbors=20, std_ratio=2.0):
    """
    Apply a statistical outlier removal filter to the point cloud.

    Parameters:
    pcd (o3d.geometry.PointCloud): The input point cloud
    nb_neighbors (int): The number of nearest neighbors to analyze for each point
    std_ratio (float): The standard deviation ratio for the distance threshold

    Returns:
    o3d.geometry.PointCloud: The filtered point cloud
    """
    pcd_filtered, _ = pcd.remove_statistical_outlier(nb_neighbors=nb_neighbors, std_ratio=std_ratio)
    return pcd_filtered

def estimate_floor_plane(pcd, distance_threshold=0.01):
    """
    Estimate the floor plane using RANSAC.

    Parameters:
    pcd (o3d.geometry.PointCloud): The input point cloud
    distance_threshold (float): The maximum distance threshold for the RANSAC algorithm

    Returns:
    Tuple: The floor plane coefficients (a, b, c, d) for the equation ax + by + cz + d = 0
    """
    plane_model, inliers = pcd.segment_plane(distance_threshold=distance_threshold,
                                             ransac_n=3,
                                             num_iterations=1000)
    return plane_model

def generate_floor_points(plane_model, pcd, grid_resolution=0.1):
    """
    Generate a grid of points on the estimated floor plane.

    Parameters:
    plane_model (Tuple): The floor plane coefficients (a, b, c, d)
    pcd (o3d.geometry.PointCloud): The input point cloud
    grid_resolution (float): The resolution of the generated grid

    Returns:
    np.ndarray: The generated grid points on the floor plane
    """
    a, b, c, d = plane_model

    # Calculate the bounds of the point cloud
    points = np.asarray(pcd.points)
    min_x, min_y, _ = np.min(points, axis=0)
    max_x, max_y, _ = np.max(points, axis=0)

    # Generate a grid of points within the bounds on the floor plane
    x_coords = np.arange(min_x, max_x, grid_resolution)
    y_coords = np.arange(min_y, max_y, grid_resolution)
    grid_points = []

    for x in x_coords:
        for y in y_coords:
            z = (-d - a * x - b * y) / c
            grid_points.append([x, y, z])

    return np.asarray(grid_points)

def main():
    # Replace 'path_to_your_pts_file.pts' with the path to your .pts file
    pts_file = '111.pts'
    print("Reading point cloud data")
    point_cloud_data = read_pts_file(pts_file)
    print("Visualize the point cloud data")

    # Create Open3D point cloud object
    pcd = o3d.geometry.PointCloud()

    # Assign point cloud data to Open3D point cloud object
    pcd.points = o3d.utility.Vector3dVector(point_cloud_data)
    print("Apply statistical outlier removal filter")
    # Apply statistical outlier removal filter
    pcd_filtered = apply_statistical_outlier_removal(pcd)

    print("Estimate floor plane")
    # Estimate floor plane
    floor_plane_model = estimate_floor_plane(pcd_filtered)
    # Generate floor points
    floor_points = generate_floor_points(floor_plane_model, pcd_filtered)
    print("Visualize the floor points")
    # Create a new point cloud object for the floor points
    floor_pcd = o3d.geometry.PointCloud()
    floor_pcd.points = o3d.utility.Vector3dVector(floor_points)

    combined_pcd = pcd_filtered + floor_pcd

    # Create a visualization
    # Bird's eye view
    vis = o3d.visualization.Visualizer()
    vis.create_window()
    vis.add_geometry(combined_pcd)

    # Set up the bird's eye view
    view_control = vis.get_view_control()
    view_control.set_front([0, 1, 0])
    view_control.set_up([0, 0, 1])
    view_control.set_zoom(0.8)

    # Render the image
    # vis.poll_events()
    # vis.update_renderer()
    # image = vis.capture_screen_float_buffer()

    # # Save the image to a PNG file
    # output_file = "bird_eye_view.png"
    # if isinstance(image, o3d.cuda.pybind.geometry.Image):
    #     image_np = np.asarray(image)
    # else:
    #     image_np = image
    # image_np_uint8 = (image_np * 255).astype(np.uint8)
    # cv2.imwrite(output_file, cv2.cvtColor(image_np_uint8, cv2.COLOR_RGB2BGR))
    #
    # # Close the visualization window
    # # vis.destroy_window()

    # Visualize the point cloud
    o3d.visualization.draw_geometries([combined_pcd])

if __name__ == "__main__":
    main()