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

from OCC.Core.STEPControl import STEPControl_Reader
from OCC.Core.IFSelect import IFSelect_RetDone
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.TopExp import TopExp_Explorer
from OCC.Core.TopAbs import TopAbs_FACE
from OCC.Core.BRep import BRep_Tool

try:
    from OCC.Core.TopoDS import topods
    USE_TOPODS = "topods"
except ImportError:
    from OCC.Core.TopoDS import TopoDS
    USE_TOPODS = "TopoDS"


# ============================================================
# STEP -> Open3D Mesh
# ============================================================

def read_step_file(step_path):
    reader = STEPControl_Reader()
    status = reader.ReadFile(step_path)

    if status != IFSelect_RetDone:
        raise RuntimeError(f"STEP 文件读取失败: {step_path}")

    reader.TransferRoots()
    shape = reader.OneShape()

    print("[OK] STEP loaded:", step_path)
    return shape


def occ_shape_to_numpy_mesh(shape):
    vertices = []
    triangles = []

    explorer = TopExp_Explorer(shape, TopAbs_FACE)

    face_count = 0
    tri_count = 0

    while explorer.More():
        current_shape = explorer.Current()

        if USE_TOPODS == "topods":
            face = topods.Face(current_shape)
        else:
            face = TopoDS.Face(current_shape)

        location = face.Location()
        triangulation = BRep_Tool.Triangulation(face, location)

        if triangulation is None:
            explorer.Next()
            continue

        trsf = location.Transformation()
        base_index = len(vertices)

        for i in range(1, triangulation.NbNodes() + 1):
            pnt = triangulation.Node(i)
            pnt.Transform(trsf)

            vertices.append([
                pnt.X(),
                pnt.Y(),
                pnt.Z()
            ])

        for i in range(1, triangulation.NbTriangles() + 1):
            tri = triangulation.Triangle(i)
            n1, n2, n3 = tri.Get()

            triangles.append([
                base_index + n1 - 1,
                base_index + n2 - 1,
                base_index + n3 - 1
            ])

            tri_count += 1

        face_count += 1
        explorer.Next()

    vertices = np.asarray(vertices, dtype=np.float64)
    triangles = np.asarray(triangles, dtype=np.int32)

    print("[OK] CAD mesh extracted")
    print("     CAD faces:", face_count)
    print("     mesh vertices:", vertices.shape)
    print("     mesh triangles:", triangles.shape)

    if vertices.shape[0] == 0 or triangles.shape[0] == 0:
        raise RuntimeError("CAD 三角网格提取失败")

    return vertices, triangles


def step_to_open3d_mesh(step_path, mesh_deflection=0.1):
    shape = read_step_file(step_path)

    print("[INFO] Meshing CAD...")
    mesh_algo = BRepMesh_IncrementalMesh(shape, mesh_deflection)
    mesh_algo.Perform()

    vertices, triangles = occ_shape_to_numpy_mesh(shape)

    mesh = o3d.geometry.TriangleMesh()
    mesh.vertices = o3d.utility.Vector3dVector(vertices)
    mesh.triangles = o3d.utility.Vector3iVector(triangles)

    mesh.remove_duplicated_vertices()
    mesh.remove_duplicated_triangles()
    mesh.remove_degenerate_triangles()
    mesh.compute_vertex_normals()

    return mesh


# ============================================================
# 点云处理
# ============================================================

def make_translation_matrix(dx, dy, dz):
    T = np.eye(4, dtype=np.float64)
    T[0, 3] = dx
    T[1, 3] = dy
    T[2, 3] = dz
    return T


def crop_front_tail_by_x(pcd, ratio=0.08):
    """
    按 X 轴裁剪车头和车尾。

    front: X 最小端
    tail:  X 最大端
    """

    points = np.asarray(pcd.points)

    min_bound = points.min(axis=0)
    max_bound = points.max(axis=0)

    x_min = min_bound[0]
    x_max = max_bound[0]
    x_len = x_max - x_min

    front_min = min_bound.copy()
    front_max = max_bound.copy()
    front_min[0] = x_min
    front_max[0] = x_min + ratio * x_len

    tail_min = min_bound.copy()
    tail_max = max_bound.copy()
    tail_min[0] = x_max - ratio * x_len
    tail_max[0] = x_max

    front_box = o3d.geometry.AxisAlignedBoundingBox(
        min_bound=front_min,
        max_bound=front_max
    )

    tail_box = o3d.geometry.AxisAlignedBoundingBox(
        min_bound=tail_min,
        max_bound=tail_max
    )

    front_pcd = pcd.crop(front_box)
    tail_pcd = pcd.crop(tail_box)

    print("[OK] Crop front/tail by X")
    print("     X range:", x_min, x_max)
    print("     front points:", np.asarray(front_pcd.points).shape)
    print("     tail points:", np.asarray(tail_pcd.points).shape)

    return front_pcd, tail_pcd


def add_gaussian_noise(pcd, noise_std=0.0):
    """
    给点云添加高斯噪声。
    noise_std=0 表示不加噪声。
    """

    if noise_std <= 0:
        return pcd

    points = np.asarray(pcd.points)
    noise = np.random.normal(
        loc=0.0,
        scale=noise_std,
        size=points.shape
    )

    noisy_points = points + noise

    noisy_pcd = o3d.geometry.PointCloud()
    noisy_pcd.points = o3d.utility.Vector3dVector(noisy_points)

    if pcd.has_colors():
        noisy_pcd.colors = pcd.colors

    if pcd.has_normals():
        noisy_pcd.normals = pcd.normals

    return noisy_pcd


def visualize_pair(cad_pcd, shifted_pcd):
    cad_vis = o3d.geometry.PointCloud(cad_pcd)
    shifted_vis = o3d.geometry.PointCloud(shifted_pcd)

    cad_vis.paint_uniform_color([0.0, 0.5, 1.0])
    shifted_vis.paint_uniform_color([1.0, 0.0, 0.0])

    o3d.visualization.draw_geometries(
        [cad_vis, shifted_vis],
        window_name="Blue = CAD sampled point cloud, Red = shifted test point cloud"
    )


# ============================================================
# 主函数
# ============================================================

def generate_shifted_train_pointcloud(
    cad_path,
    output_dir,
    sample_points=1000000,
    mesh_deflection=0.1,
    crop_ratio=0.08,
    dx=100.0,
    dy=20.0,
    dz=-10.0,
    noise_std=0.0,
    visualize=True
):
    os.makedirs(output_dir, exist_ok=True)

    # 1. CAD -> mesh
    cad_mesh = step_to_open3d_mesh(
        cad_path,
        mesh_deflection=mesh_deflection
    )

    # 2. mesh -> 点云
    print(f"[INFO] Sampling CAD surface: {sample_points} points")
    cad_pcd = cad_mesh.sample_points_uniformly(
        number_of_points=sample_points
    )

    cad_pcd_path = os.path.join(output_dir, "cad_sampled_whole.ply")
    o3d.io.write_point_cloud(cad_pcd_path, cad_pcd)
    print("[OK] Saved:", cad_pcd_path)

    # 3. 裁剪原始 CAD 点云的车头 / 车尾
    cad_front, cad_tail = crop_front_tail_by_x(
        cad_pcd,
        ratio=crop_ratio
    )

    o3d.io.write_point_cloud(
        os.path.join(output_dir, "cad_front.ply"),
        cad_front
    )

    o3d.io.write_point_cloud(
        os.path.join(output_dir, "cad_tail.ply"),
        cad_tail
    )

    # 4. 生成固定平移矩阵
    T = make_translation_matrix(dx, dy, dz)

    print("\n========== Fixed Translation ==========")
    print("dx:", dx)
    print("dy:", dy)
    print("dz:", dz)
    print("T:")
    print(T)

    # 5. 对整车点云施加固定平移
    shifted_pcd = o3d.geometry.PointCloud(cad_pcd)
    shifted_pcd.transform(T)

    # 6. 可选：加噪声，模拟真实扫描点云
    shifted_pcd = add_gaussian_noise(
        shifted_pcd,
        noise_std=noise_std
    )

    shifted_path = os.path.join(output_dir, "shifted_whole_pointcloud.ply")
    o3d.io.write_point_cloud(shifted_path, shifted_pcd)
    print("[OK] Saved:", shifted_path)

    # 7. 裁剪平移后的车头 / 车尾
    shifted_front, shifted_tail = crop_front_tail_by_x(
        shifted_pcd,
        ratio=crop_ratio
    )

    shifted_front_path = os.path.join(output_dir, "shifted_front.ply")
    shifted_tail_path = os.path.join(output_dir, "shifted_tail.ply")

    o3d.io.write_point_cloud(shifted_front_path, shifted_front)
    o3d.io.write_point_cloud(shifted_tail_path, shifted_tail)

    print("[OK] Saved:", shifted_front_path)
    print("[OK] Saved:", shifted_tail_path)

    # 8. 保存真实位移矩阵
    matrix_path = os.path.join(output_dir, "GT_translation_matrix.txt")
    np.savetxt(matrix_path, T, fmt="%.10f")
    print("[OK] Saved:", matrix_path)

    # 9. 保存真实位移向量
    translation_path = os.path.join(output_dir, "GT_translation_vector.txt")
    np.savetxt(
        translation_path,
        np.array([dx, dy, dz], dtype=np.float64),
        fmt="%.10f"
    )
    print("[OK] Saved:", translation_path)

    # 10. 可视化
    if visualize:
        # 百万级点云显示会卡，所以这里可视化时降采样
        cad_vis = cad_pcd.voxel_down_sample(voxel_size=20.0)
        shifted_vis = shifted_pcd.voxel_down_sample(voxel_size=20.0)

        visualize_pair(cad_vis, shifted_vis)

    return {
        "cad_sampled_path": cad_pcd_path,
        "shifted_pointcloud_path": shifted_path,
        "shifted_front_path": shifted_front_path,
        "shifted_tail_path": shifted_tail_path,
        "translation_matrix_path": matrix_path,
        "translation_vector": np.array([dx, dy, dz], dtype=np.float64)
    }


# ============================================================
# 命令行入口
# ============================================================

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="从整车 CAD 生成测试点云，并施加固定前后/整体位移"
    )

    parser.add_argument(
        "--cad",
        default=r"datasets/CAD/M02车.stp",
        help="整车 CAD STEP/STP 文件路径"
    )

    parser.add_argument(
        "--output_dir",
        default=r"datasets\shift_test_output",
        help="输出目录"
    )

    parser.add_argument(
        "--sample_points",
        type=int,
        default=1000000,
        help="从 CAD 表面采样的点数，默认 1000000"
    )

    parser.add_argument(
        "--mesh_deflection",
        type=float,
        default=0.1,
        help="CAD 三角化精度，越小越精细，默认 0.1"
    )

    parser.add_argument(
        "--crop_ratio",
        type=float,
        default=0.08,
        help="车头/车尾裁剪比例，默认 0.08"
    )

    parser.add_argument(
        "--dx",
        type=float,
        default=100.0,
        help="X 方向固定平移量"
    )

    parser.add_argument(
        "--dy",
        type=float,
        default=20.0,
        help="Y 方向固定平移量"
    )

    parser.add_argument(
        "--dz",
        type=float,
        default=-10.0,
        help="Z 方向固定平移量"
    )

    parser.add_argument(
        "--noise_std",
        type=float,
        default=0.0,
        help="高斯噪声标准差，默认 0 表示不加噪声"
    )

    parser.add_argument(
        "--no_visualize",
        action="store_true",
        help="不显示可视化窗口"
    )

    args = parser.parse_args()

    result = generate_shifted_train_pointcloud(
        cad_path=args.cad,
        output_dir=args.output_dir,
        sample_points=args.sample_points,
        mesh_deflection=args.mesh_deflection,
        crop_ratio=args.crop_ratio,
        dx=args.dx,
        dy=args.dy,
        dz=args.dz,
        noise_std=args.noise_std,
        visualize=not args.no_visualize
    )

    print("\n========== Result ==========")
    print("CAD sampled point cloud:", result["cad_sampled_path"])
    print("Shifted whole point cloud:", result["shifted_pointcloud_path"])
    print("Shifted front point cloud:", result["shifted_front_path"])
    print("Shifted tail point cloud:", result["shifted_tail_path"])
    print("GT translation vector:", result["translation_vector"])
    print("GT matrix:", result["translation_matrix_path"])