#!/usr/bin/env python3
import argparse
import os
from pathlib import Path
import numpy as np
import open3d as o3d

def parse_mapping(mapping_str: str):
    """
    解析映射字符串，例如 'z,y,-x'
    返回一个变换矩阵或映射逻辑
    """
    mapping_str = mapping_str.lower().replace(" ", "")
    parts = mapping_str.split(",")
    if len(parts) != 3:
        raise ValueError("映射关系必须包含三个部分，用逗号分隔，例如 'z,y,-x'")
    
    # 构造变换矩阵
    # [new_x, new_y, new_z, 1]^T = T * [x, y, z, 1]^T
    transform = np.zeros((4, 4))
    transform[3, 3] = 1.0
    
    axis_map = {'x': 0, 'y': 1, 'z': 2}
    
    for i, part in enumerate(parts):
        sign = 1.0
        if part.startswith('-'):
            sign = -1.0
            axis = part[1:]
        elif part.startswith('+'):
            axis = part[1:]
        else:
            axis = part
            
        if axis not in axis_map:
            raise ValueError(f"无效的轴名称: {axis}. 必须是 x, y, 或 z")
            
        transform[i, axis_map[axis]] = sign
        
    return transform

def process_pcd(file_path, transform, output_path):
    print(f"正在处理: {file_path.name} ...", end="\r")
    pcd = o3d.io.read_point_cloud(str(file_path))
    if pcd.is_empty():
        print(f"\n[跳过] {file_path.name} 为空或无法读取")
        return False
    
    # 应用变换
    pcd.transform(transform)
    
    # 保存
    success = o3d.io.write_point_cloud(str(output_path), pcd)
    return success

def main():
    parser = argparse.ArgumentParser(description="PCD 坐标修复工具 - 批量转换坐标轴")
    parser.add_argument("input_dir", help="输入 PCD 文件夹路径")
    parser.add_argument("--map", required=True, help="坐标映射关系，例如 'z,y,-x' (新X=原Z, 新Y=原Y, 新Z=原-X)")
    parser.add_argument("--out_dir", help="输出文件夹路径（如果不指定，将直接覆盖原文件）")
    parser.add_argument("--ext", default=".pcd", help="要处理的文件后缀 (默认: .pcd)")

    args = parser.parse_args()
    
    input_path = Path(args.input_dir)
    if not input_path.is_dir():
        print(f"错误: {args.input_dir} 不是一个有效的目录")
        return

    try:
        transform_matrix = parse_mapping(args.map)
        print("解析到的变换矩阵:")
        print(transform_matrix[:3, :3])
    except ValueError as e:
        print(f"映射解析错误: {e}")
        return

    output_dir = Path(args.out_dir) if args.out_dir else input_path
    output_dir.mkdir(parents=True, exist_ok=True)

    pcd_files = list(input_path.glob(f"*{args.ext}"))
    if not pcd_files:
        print(f"在 {args.input_dir} 中没有找到 {args.ext} 文件")
        return

    print(f"找到 {len(pcd_files)} 个文件，准备开始处理...")
    
    count = 0
    for pcd_file in pcd_files:
        out_file = output_dir / pcd_file.name
        if process_pcd(pcd_file, transform_matrix, out_file):
            count += 1
            
    print(f"\n完成！成功处理并保存了 {count} 个文件。")

if __name__ == "__main__":
    main()
