"""
PartNetE Datasets preprocessing

Author: Yujia Zhang (yujia.zhang.cs@gmail.com)
Please cite our work if the code is helpful to you.
"""

import os
import shutil
import argparse
import numpy as np
import trimesh
import open3d as o3d


def process_folder(target_dir):
    ply_path = os.path.join(target_dir, "pc.ply")
    label_path = os.path.join(target_dir, "label.npy")

    if os.path.exists(ply_path):
        mesh = trimesh.load(ply_path, process=False)
        coords = np.array(mesh.vertices, dtype=np.float32)
        np.save(os.path.join(target_dir, "coord.npy"), coords)

        if (
            hasattr(mesh, "vertex_normals")
            and mesh.vertex_normals is not None
            and len(mesh.vertex_normals) == len(coords)
        ):
            normals = np.array(mesh.vertex_normals, dtype=np.float32)
        else:
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(coords)
            search_param = o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30)
            pcd.estimate_normals(search_param=search_param)
            normals = np.asarray(pcd.normals).astype(np.float32)

        np.save(os.path.join(target_dir, "normal.npy"), normals)

        if (
            hasattr(mesh.visual, "vertex_colors")
            and mesh.visual.vertex_colors is not None
        ):
            colors = np.array(mesh.visual.vertex_colors[:, :3], dtype=np.uint8)
            np.save(os.path.join(target_dir, "color.npy"), colors)
        elif hasattr(mesh, "colors") and mesh.colors is not None:
            colors = np.array(mesh.colors[:, :3], dtype=np.uint8)
            np.save(os.path.join(target_dir, "color.npy"), colors)
        else:
            pass

        label_data = np.load(label_path, allow_pickle=True).item()
        segment = label_data["semantic_seg"]
        instance = label_data["instance_seg"]
        assert coords.shape[0] == segment.shape[0]
        segment_path = os.path.join(target_dir, "segment.npy")
        instance_path = os.path.join(target_dir, "instance.npy")
        np.save(segment_path, segment)
        np.save(instance_path, instance)
    else:
        print(f"Warning: pc.ply not found in {target_dir}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Preprocess PartNetE dataset by splitting PLY files and copying labels."
    )
    parser.add_argument(
        "--dataset_root",
        required=True,
        help="Path to the ScanNet dataset containing scene folders",
    )
    args = parser.parse_args()
    subfolders = ["few_shot", "test"]
    total_processed_count = 0

    for subfolder in subfolders:
        current_root = os.path.join(args.dataset_root, subfolder)
        processed_in_subfolder = 0
        for dirpath, dirnames, filenames in os.walk(current_root):
            if "pc.ply" in filenames and "label.npy" in filenames:
                process_folder(dirpath)
                processed_in_subfolder += 1