import open3d as o3d
import numpy as np
import os


def reconstruct_and_save_ply(npy_dir, output_ply_path):
    """
    Reads coord.npy, color.npy, and normal.npy from a directory,
    reconstructs a point cloud, and saves it as a .ply file.
    """
    print(f"Processing directory: {npy_dir}")

    # Define file paths
    coord_file = os.path.join(npy_dir, 'coord.npy')
    color_file = os.path.join(npy_dir, 'color.npy')
    normal_file = os.path.join(npy_dir, 'normal.npy')

    # Check if coordinate file exists (it's required)
    if not os.path.exists(coord_file):
        print(f"  - Error: coord.npy not found in {npy_dir}. Skipping.")
        return

    # Load coordinates
    try:
        coords = np.load(coord_file)
        print(f"  - Loaded {len(coords)} points from coord.npy.")
    except Exception as e:
        print(f"  - Error loading coord.npy: {e}. Skipping.")
        return

    # Create point cloud object
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(coords)

    # Load colors if available
    if os.path.exists(color_file):
        try:
            colors = np.load(color_file)
            if len(colors) == len(coords):
                pcd.colors = o3d.utility.Vector3dVector(colors)
                print("  - Loaded colors.")
            else:
                print("  - Warning: Mismatch between point count and color count. Skipping colors.")
        except Exception as e:
            print(f"  - Error loading color.npy: {e}.")
    else:
        print("  - info: color.npy not found.")

    # Load normals if available
    if os.path.exists(normal_file):
        try:
            normals = np.load(normal_file)
            if len(normals) == len(coords):
                pcd.normals = o3d.utility.Vector3dVector(normals)
                print("  - Loaded normals.")
            else:
                print("  - Warning: Mismatch between point count and normal count. Skipping normals.")
        except Exception as e:
            print(f"  - Error loading normal.npy: {e}.")
    else:
        print("  - info: normal.npy not found.")

    # Save the reconstructed point cloud
    try:
        o3d.io.write_point_cloud(output_ply_path, pcd, write_ascii=True)
        print(f"  - Successfully saved reconstructed point cloud to: {output_ply_path}")
    except Exception as e:
        print(f"  - Error saving .ply file: {e}")


if __name__ == '__main__':
    scenes_folder = "/media/gu/d54b9541-2b55-4c75-b059-3006d51983d53/lqc/Downloads/Scenes"
    scenes = os.listdir(scenes_folder)
    for scene_name in scenes:
        print(f"\nVerifying scene: {scene_name}...")
        scene_folder = os.path.join(scenes_folder, scene_name)
        npy_folder = os.path.join(scene_folder, "scene/val/process_data")

        # Directory where the .npy files are located
        process_folder = os.path.join(scene_folder, 'denoise_ply')
        if not os.path.exists(process_folder):
            os.makedirs(process_folder)

        # Path for the output .ply file
        output_ply = os.path.join(scene_folder, 'verified_cut_result.ply')

        if os.path.isdir(process_folder):
            reconstruct_and_save_ply(npy_folder, output_ply)
        else:
            print(f"  - Process data folder not found: {process_folder}. Skipping scene.")