import argparse
from collections import defaultdict
import glob
import json
import os

import numpy as np
import trimesh
import point_cloud_utils as pcu
from tqdm import tqdm


import numpy as np

def voxelize(points, voxel_size=0.1):
    """
    Converts a set of 3D points to a voxel grid representation.
    Points are quantized to the nearest voxel center.

    Parameters:
    - points: (N, 3) numpy array of 3D points
    - voxel_size: the size of each voxel (adjust this depending on your data)

    Returns:
    - voxels: Set of unique voxel coordinates
    """
    # Quantize the points to the nearest voxel
    quantized_points = np.floor(points / voxel_size).astype(int)
    # Use a set to get unique voxel coordinates
    voxels = set(map(tuple, quantized_points))
    return voxels

def calculate_iou(model_vox, target_vox, voxel_size=0.1):
    """
    Calculate the IoU (Intersection over Union) between two point clouds.

    Parameters:
    - model_vox: (N, 3) numpy array of the first point cloud
    - target_vox: (M, 3) numpy array of the second point cloud
    - voxel_size: Size of the voxels (default is 0.1)

    Returns:
    - iou: Intersection over Union (IoU) score
    """
    # Voxelize both point clouds
    model_voxels = voxelize(model_vox, voxel_size)
    target_voxels = voxelize(target_vox, voxel_size)

    # Calculate intersection and union
    intersection = len(model_voxels.intersection(target_voxels))
    union = len(model_voxels.union(target_voxels))

    # Compute IoU
    iou = intersection / union if union > 0 else 0.0
    return iou


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dir", type=str, default="./results/infer/PointClouds")
    parser.add_argument("--target_dir", type=str, default="./data/test_pc")
    parser.add_argument("--detail", action="store_true")
    args = parser.parse_args()

    if not os.path.exists(args.input_dir) or not os.path.exists(args.target_dir):
        print("Invalid input!")
        exit(1)

    model_prefix = os.path.join(args.input_dir, "*.ply")
    model_path_list = sorted(list(glob.glob(model_prefix)))

    distance_json = {}
    distance_list = []
    emu_distance_list = []
    hausdorff_distance_list = []
    voxel_iou_list = []
    for model_path in tqdm(model_path_list):
        model_name = os.path.basename(model_path)
        target_path = os.path.join(args.target_dir, model_name)
        if not os.path.exists(target_path):
            print(f"{target_path}: not found!")
            exit(1)

        model_pc = np.array(trimesh.load(model_path).vertices)
        target_pc = np.array(trimesh.load(target_path).vertices)

        distance = pcu.chamfer_distance(model_pc, target_pc)
        model_pc_downsampled = model_pc[np.random.choice(model_pc.shape[0], 1000, replace=False)]
        target_pc_downsampled = target_pc[np.random.choice(target_pc.shape[0], 1000, replace=False)]
        emu_distance, _ = pcu.earth_movers_distance(model_pc_downsampled, target_pc_downsampled)
        hausdorff_distance = pcu.hausdorff_distance(model_pc, target_pc)

        iou = calculate_iou(model_pc, target_pc, voxel_size=1/32.)

        distance_list.append(distance)
        emu_distance_list.append(emu_distance)
        hausdorff_distance_list.append(hausdorff_distance)
        voxel_iou_list.append(iou)
        model_id = os.path.splitext(model_name)[0]
        distance_json[model_id] = distance

        print(f"{model_id}: chamfer distance: {distance:.3f}, earth movers distance: {emu_distance:.3f}, hausdorff distance: {hausdorff_distance:.3f}, voxel IoU: {iou:.3f}")

    distance_json["mean"] = np.mean(distance_list)
    distance_json["mean_emu"] = np.mean(emu_distance_list)
    distance_json["mean_hausdorff"] = np.mean(hausdorff_distance_list)
    distance_json["mean_voxel_iou"] = np.mean(voxel_iou_list)
    print(f"mean chamfer distance: {np.mean(distance_list)}")
    print(f"mean earth movers distance: {np.mean(emu_distance_list)}")
    print(f"mean hausdorff distance: {np.mean(hausdorff_distance_list)}")
    print(f"mean voxel IoU: {np.mean(voxel_iou_list)}")
    with open(os.path.join(args.input_dir, "distance.json"), "w") as json_file:
        json.dump(distance_json, json_file, indent=4)