import numpy as np
import open3d as o3d
import os
import time
import json
import argparse

def save_point_cloud_by_class(coords_file, preds_file, classes_to_show, save_pcd_path):
    """
    加载点云坐标和预测类别，筛选指定类别，并保存为.ply文件。
    """
    # 1. 定义类别名称和颜色映射
    CLASS_NAMES = [
        'refrigerator', 'desk', 'curtain', 'sofa', 'bookshelf', 'bed',
        'table', 'window', 'cabinet', 'door', 'chair', 'floor', 'wall',
        'sink', 'toilet', 'bathtub', 'shower curtain', 'picture', 'counter'
    ]
    COLOR_MAP = np.array([
        [174, 199, 232], [255, 127, 14], [152, 223, 138], [214, 39, 40],
        [148, 103, 189], [255, 187, 120], [140, 86, 75], [152, 223, 138],
        [152, 223, 138], [152, 223, 138], [196, 156, 148], [127, 127, 127],
        [152, 223, 138], [188, 189, 34], [219, 219, 141], [227, 119, 194],
        [31, 119, 180], [255, 152, 150], [82, 84, 163]
    ])

    # 2. 加载数据
    try:
        print(f"正在加载坐标文件: {coords_file}")
        coords = np.load(coords_file)
        print(f"正在加载预测文件: {preds_file}")
        predictions = np.load(preds_file)
    except FileNotFoundError as e:
        print(f"错误: 找不到文件 {e.filename}。请确保文件路径正确。")
        return

    print(f"原始点云数量: {len(coords)}")
    if len(coords) != len(predictions):
        print("警告: 坐标点数和预测标签数不匹配！")
        return

    # 3. 根据 'classes_to_show' 筛选点云
    if isinstance(classes_to_show, str):
        target_classes = [classes_to_show]
    elif isinstance(classes_to_show, list):
        target_classes = classes_to_show
    else:
        print("错误: 'classes_to_show' 参数必须是字符串或列表。")
        return

    target_indices = [CLASS_NAMES.index(cn) for cn in target_classes if cn in CLASS_NAMES]
    if not target_indices:
        print(f"错误: 类别 '{classes_to_show}' 无效或未在 CLASS_NAMES 中找到。")
        return

    print(f"正在筛选类别: {[CLASS_NAMES[i] for i in target_indices]}")
    mask = np.isin(predictions, target_indices)
    coords = coords[mask]
    predictions = predictions[mask]

    if coords.shape[0] == 0:
        print(f"警告: 在场景中没有找到属于类别 {target_classes} 的点。")
        return
    print(f"筛选后剩余点云数量: {len(coords)}")

    # 4. 创建Open3D点云对象
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(coords)
    pcd.colors = o3d.utility.Vector3dVector(COLOR_MAP[predictions] / 255.0)

    # 5. 保存筛选后的点云
    if not save_pcd_path:
        print("警告: 未提供 'save_pcd_path'，不保存文件。")
        return

    print(f"正在保存筛选后的点云到: {save_pcd_path}")
    try:
        o3d.io.write_point_cloud(save_pcd_path, pcd, write_ascii=True)
        print("保存成功。")
    except Exception as e:
        print(f"错误: 保存点云文件失败: {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="输入单个场景文件夹用于结果点云保存",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        '-i',
        '--input_folder',
        type=str,
        required=True,
        help='指定输入场景的文件夹路径。'
    )
    args = parser.parse_args()
    scene_folder = args.input_folder
    scenece = os.path.basename(scene_folder)

    coords_file = os.path.join(scene_folder, 'scene/val/process_data/coord.npy')
    preds_file = os.path.join(scene_folder, 'output/pred.npy')

    output_dir = os.path.join(scene_folder, 'wall_ply')
    os.makedirs(output_dir, exist_ok=True)

    print(f"\n--- 正在处理场景: {scenece} ---")
    save_point_cloud_by_class(
        coords_file=coords_file,
        preds_file=preds_file,
        classes_to_show=['wall', 'window', 'door', 'cabinet'],
        save_pcd_path=f'{output_dir}/wall_instances.ply'
    )