import argparse
from pathlib import Path
from zipfile import ZipFile, ZIP_DEFLATED
from xml.sax.saxutils import escape
import numpy as np
import torch
import trimesh
import matplotlib.pyplot as plt

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

# 兼容不同 pythonocc 版本
try:
    from OCC.Core.TopoDS import topods
    USE_TOPODS = "topods"
except ImportError:
    from OCC.Core.TopoDS import TopoDS
    USE_TOPODS = "TopoDS"

try:
    import open3d as o3d
except ModuleNotFoundError:
    o3d = None


# ============================================================
# 设备与精度
# ============================================================

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float64

DEFAULT_TASK_IDS = [
    11, 13, 17, 20, 25, 26, 27, 31, 34,
    35, 36, 37, 38, 39, 40, 41, 42, 45
]

task_distance_data = {}


# ============================================================
# STEP 读取
# ============================================================

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("STEP 文件读取完成")

    return shape


# ============================================================
# OCC Shape -> numpy mesh
# ============================================================

def occ_shape_to_numpy_mesh(shape):
    """
    从 OCC Shape 中提取三角网格。

    返回:
        vertices_np: [N, 3]
        faces_np: [M, 3]
    """

    vertices = []
    faces = []

    explorer = TopExp_Explorer(shape, TopAbs_FACE)

    face_count = 0
    triangle_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)

        # 读取当前 Face 的顶点
        for i in range(1, triangulation.NbNodes() + 1):
            pnt = triangulation.Node(i)

            # 应用局部变换
            pnt.Transform(trsf)

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

        # 读取当前 Face 的三角形
        for i in range(1, triangulation.NbTriangles() + 1):
            tri = triangulation.Triangle(i)

            n1, n2, n3 = tri.Get()

            # OCC 索引从 1 开始，numpy / torch 索引从 0 开始
            faces.append([
                base_index + n1 - 1,
                base_index + n2 - 1,
                base_index + n3 - 1
            ])

            triangle_count += 1

        face_count += 1
        explorer.Next()

    vertices_np = np.asarray(vertices, dtype=np.float64)
    faces_np = np.asarray(faces, dtype=np.int64)

    print("CAD 网格提取完成")
    print(f"CAD 面数量: {face_count}")
    print(f"网格顶点: {vertices_np.shape}")
    print(f"网格三角面: {faces_np.shape}")

    if vertices_np.shape[0] == 0 or faces_np.shape[0] == 0:
        raise RuntimeError(
            "CAD 三角网格提取失败。可以尝试减小 MESH_DEFLECTION，或者检查 STEP 文件是否正常。"
        )

    return vertices_np, faces_np


# ============================================================
# 读取点云
# ============================================================

def read_point_cloud(pointcloud_path):
    if o3d is not None:
        pcd = o3d.io.read_point_cloud(pointcloud_path)
        points_np = np.asarray(pcd.points, dtype=np.float64)
    else:
        pcd = trimesh.load(pointcloud_path)
        points_np = np.asarray(pcd.vertices, dtype=np.float64)

    if points_np.shape[0] == 0:
        raise RuntimeError(f"点云读取失败或点数为 0: {pointcloud_path}")

    print("点云读取完成")
    print(f"点云点数: {points_np.shape}")

    return points_np


# ============================================================
# 点到三角形距离
# ============================================================

def point_to_triangle_distance_squared(points, triangles):
    """
    points:
        [P, 3]
    triangles:
        [T, 3, 3]

    返回:
        dist2: [P, T]
    """

    p = points[:, None, :]

    a = triangles[None, :, 0, :]
    b = triangles[None, :, 1, :]
    c = triangles[None, :, 2, :]

    ab = b - a
    ac = c - a
    ap = p - a

    d1 = (ab * ap).sum(dim=-1)
    d2 = (ac * ap).sum(dim=-1)

    inf = torch.full_like(d1, torch.inf)

    # 区域 A
    mask = (d1 <= 0.0) & (d2 <= 0.0)
    dist2 = torch.where(mask, (ap * ap).sum(dim=-1), inf)

    # 区域 B
    bp = p - b
    d3 = (ab * bp).sum(dim=-1)
    d4 = (ac * bp).sum(dim=-1)

    mask = (d3 >= 0.0) & (d4 <= d3)
    dist2 = torch.minimum(
        dist2,
        torch.where(mask, (bp * bp).sum(dim=-1), inf)
    )

    # 边 AB
    vc = d1 * d4 - d3 * d2
    mask = (vc <= 0.0) & (d1 >= 0.0) & (d3 <= 0.0)

    v = d1 / (d1 - d3).clamp_min(torch.finfo(points.dtype).eps)
    closest = a + v[..., None] * ab
    edge_dist2 = ((p - closest) ** 2).sum(dim=-1)

    dist2 = torch.minimum(
        dist2,
        torch.where(mask, edge_dist2, inf)
    )

    # 区域 C
    cp = p - c
    d5 = (ab * cp).sum(dim=-1)
    d6 = (ac * cp).sum(dim=-1)

    mask = (d6 >= 0.0) & (d5 <= d6)
    dist2 = torch.minimum(
        dist2,
        torch.where(mask, (cp * cp).sum(dim=-1), inf)
    )

    # 边 AC
    vb = d5 * d2 - d1 * d6
    mask = (vb <= 0.0) & (d2 >= 0.0) & (d6 <= 0.0)

    w = d2 / (d2 - d6).clamp_min(torch.finfo(points.dtype).eps)
    closest = a + w[..., None] * ac
    edge_dist2 = ((p - closest) ** 2).sum(dim=-1)

    dist2 = torch.minimum(
        dist2,
        torch.where(mask, edge_dist2, inf)
    )

    # 边 BC
    va = d3 * d6 - d5 * d4
    mask = (va <= 0.0) & ((d4 - d3) >= 0.0) & ((d5 - d6) >= 0.0)

    w = (d4 - d3) / ((d4 - d3) + (d5 - d6)).clamp_min(
        torch.finfo(points.dtype).eps
    )
    closest = b + w[..., None] * (c - b)
    edge_dist2 = ((p - closest) ** 2).sum(dim=-1)

    dist2 = torch.minimum(
        dist2,
        torch.where(mask, edge_dist2, inf)
    )

    # 三角形内部区域
    denom = (va + vb + vc).clamp_min(torch.finfo(points.dtype).eps)

    v = vb / denom
    w = vc / denom

    closest = a + ab * v[..., None] + ac * w[..., None]
    face_dist2 = ((p - closest) ** 2).sum(dim=-1)

    dist2 = torch.minimum(dist2, face_dist2)

    return dist2


# ============================================================
# 分块计算最近距离
# ============================================================

def nearest_mesh_distances(
    points,
    triangles,
    point_chunk_size=2048,
    triangle_chunk_size=4096
):
    """
    为了避免显存爆炸，点和三角面都分块计算。

    points:
        [N, 3]
    triangles:
        [M, 3, 3]

    返回:
        distances: [N]
    """

    all_distances = []

    num_points = points.shape[0]
    num_triangles = triangles.shape[0]

    for p_start in range(0, num_points, point_chunk_size):
        p_end = min(p_start + point_chunk_size, num_points)

        point_chunk = points[p_start:p_end]

        min_dist2 = torch.full(
            (point_chunk.shape[0],),
            torch.inf,
            dtype=points.dtype,
            device=points.device
        )

        for t_start in range(0, num_triangles, triangle_chunk_size):
            t_end = min(t_start + triangle_chunk_size, num_triangles)

            triangle_chunk = triangles[t_start:t_end]

            dist2 = point_to_triangle_distance_squared(
                point_chunk,
                triangle_chunk
            )

            min_dist2 = torch.minimum(
                min_dist2,
                dist2.min(dim=1).values
            )

        all_distances.append(torch.sqrt(min_dist2))

        # print(f"processed points: {p_end}/{num_points}")

    distances = torch.cat(all_distances, dim=0)

    return distances


# ============================================================
# 热力图显示
# ============================================================

def visualize_point_cloud_heatmap(points_np, distances_np):
    if o3d is None:
        print("未安装 open3d，跳过可视化。")
        return

    d_min = distances_np.min()
    d_max = distances_np.max()

    dist_norm = (distances_np - d_min) / (d_max - d_min + 1e-12)

    colors = plt.cm.turbo(dist_norm)[:, :3]

    pcd_show = o3d.geometry.PointCloud()
    pcd_show.points = o3d.utility.Vector3dVector(points_np)
    pcd_show.colors = o3d.utility.Vector3dVector(colors)

    print("色阶规则：深蓝/紫色 = 最近，红色 = 最远")

    o3d.visualization.draw_geometries(
        [pcd_show],
        window_name="点云到 CAD 距离热力图"
    )


def visualize_task_result(task_results, task_id):
    """
    task_results: run_batch_tasks 返回的字典，key 为任务号。
    task_id: 需要可视化的任务号。
    """

    result = task_results.get(int(task_id))

    if result is None:
        raise KeyError(f"未找到任务 {task_id} 的计算结果。")

    visualize_point_cloud_heatmap(
        result["points_np"],
        result["distances_np"]
    )


def _xlsx_cell_ref(row_index, col_index):
    letters = ""
    col = col_index

    while col:
        col, remainder = divmod(col - 1, 26)
        letters = chr(65 + remainder) + letters

    return f"{letters}{row_index}"


def _xlsx_cell_xml(row_index, col_index, value):
    cell_ref = _xlsx_cell_ref(row_index, col_index)

    if isinstance(value, (int, float, np.integer, np.floating)) and not isinstance(value, bool):
        return f'<c r="{cell_ref}"><v>{value}</v></c>'

    text = escape("" if value is None else str(value))
    return f'<c r="{cell_ref}" t="inlineStr"><is><t>{text}</t></is></c>'


def write_flatness_excel(rows, output_path, sheet_name="平整度"):
    """
    使用 Python 标准库生成 .xlsx，避免额外依赖。
    """

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    headers = ["任务号", "平均距离(mm)", "最大距离(mm)", "状态", "说明"]
    table = [headers]

    for row in rows:
        table.append([
            row.get("task_id"),
            row.get("mean_dist"),
            row.get("max_dist"),
            row.get("status", ""),
            row.get("message", "")
        ])

    sheet_rows = []
    for row_index, row_values in enumerate(table, start=1):
        cells = [
            _xlsx_cell_xml(row_index, col_index, value)
            for col_index, value in enumerate(row_values, start=1)
        ]
        sheet_rows.append(f'<row r="{row_index}">{"".join(cells)}</row>')

    safe_sheet_name = escape(sheet_name[:31])

    workbook_xml = (
        '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
        '<workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" '
        'xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">'
        f'<sheets><sheet name="{safe_sheet_name}" sheetId="1" r:id="rId1"/></sheets>'
        '</workbook>'
    )

    worksheet_xml = (
        '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
        '<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
        '<sheetData>'
        f'{"".join(sheet_rows)}'
        '</sheetData>'
        '</worksheet>'
    )

    content_types_xml = (
        '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
        '<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">'
        '<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>'
        '<Default Extension="xml" ContentType="application/xml"/>'
        '<Override PartName="/xl/workbook.xml" '
        'ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>'
        '<Override PartName="/xl/worksheets/sheet1.xml" '
        'ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>'
        '</Types>'
    )

    root_rels_xml = (
        '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
        '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
        '<Relationship Id="rId1" '
        'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" '
        'Target="xl/workbook.xml"/>'
        '</Relationships>'
    )

    workbook_rels_xml = (
        '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
        '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
        '<Relationship Id="rId1" '
        'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" '
        'Target="worksheets/sheet1.xml"/>'
        '</Relationships>'
    )

    with ZipFile(output_path, "w", ZIP_DEFLATED) as xlsx:
        xlsx.writestr("[Content_Types].xml", content_types_xml)
        xlsx.writestr("_rels/.rels", root_rels_xml)
        xlsx.writestr("xl/workbook.xml", workbook_xml)
        xlsx.writestr("xl/_rels/workbook.xml.rels", workbook_rels_xml)
        xlsx.writestr("xl/worksheets/sheet1.xml", worksheet_xml)

    return output_path


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

def compute_point_CAD_distance(
    CAD_path,
    pointcloud_path,
    mesh_deflection=0.1,
    point_chunk_size=2048,
    triangle_chunk_size=4096,
    visualize=False
):
    print("计算设备:", device)
    print("数据精度:", dtype)

    # STEP -> OCC Shape
    shape = read_step_file(CAD_path)

    # OCC Shape 三角化
    print("正在对 CAD 模型进行三角化...")
    mesh = BRepMesh_IncrementalMesh(shape, mesh_deflection)
    mesh.Perform()

    # 提取 numpy 网格
    vertices_np, faces_np = occ_shape_to_numpy_mesh(shape)

    # 读取点云
    points_np = read_point_cloud(pointcloud_path)

    # numpy -> torch
    points = torch.as_tensor(points_np, dtype=dtype, device=device)
    vertices = torch.as_tensor(vertices_np, dtype=dtype, device=device)
    faces = torch.as_tensor(faces_np, dtype=torch.long, device=device)

    triangles = vertices[faces]

    print("开始计算点云到 CAD 的最近距离...")
    print("点云张量:", points.shape)
    print("三角面张量:", triangles.shape)

    with torch.no_grad():
        distances_torch = nearest_mesh_distances(
            points,
            triangles,
            point_chunk_size=point_chunk_size,
            triangle_chunk_size=triangle_chunk_size
        )

    distances_np = distances_torch.detach().cpu().numpy()

    mean_dist = float(distances_np.mean())
    max_dist = float(distances_np.max())
    min_dist = float(distances_np.min())

    # print("Distance computation finished")
    print(f"最小距离: {min_dist:.8f}")
    print(f"平均距离: {mean_dist:.8f}")
    print(f"最大距离: {max_dist:.8f}")

    if visualize:
        visualize_point_cloud_heatmap(points_np, distances_np)

    return {
        "mean_dist": mean_dist,
        "max_dist": max_dist,
        "min_dist": min_dist,
        "points_np": points_np,
        "distances_np": distances_np
    }


def get_task_paths(task_id, data_root="High_Speed_Data", pointcloud_kind="noise"):
    data_root = Path(data_root)
    task_id = int(task_id)

    cad_path = data_root / "CAD" / f"{task_id}.step"

    if pointcloud_kind == "clean":
        pointcloud_path = data_root / "PointCloud" / f"{task_id}.ply"
    else:
        pointcloud_path = data_root / "Noise_PointCloud" / f"{task_id}_noisy.ply"

    return cad_path, pointcloud_path


def run_batch_tasks(
    task_ids=DEFAULT_TASK_IDS,
    data_root="High_Speed_Data",
    pointcloud_kind="noise",
    excel_output="平整度.xlsx",
    mesh_deflection=0.1,
    point_chunk_size=2048,
    triangle_chunk_size=4096
):
    global task_distance_data

    result_rows = []
    task_distance_data = {}

    for task_id in task_ids:
        task_id = int(task_id)
        cad_path, pointcloud_path = get_task_paths(
            task_id,
            data_root=data_root,
            pointcloud_kind=pointcloud_kind
        )

        print(f"===== 任务 {task_id} =====")

        missing_paths = [
            str(path)
            for path in [cad_path, pointcloud_path]
            if not path.exists()
        ]

        if missing_paths:
            message = "缺少文件: " + "; ".join(missing_paths)
            print(message)
            result_rows.append({
                "task_id": task_id,
                "mean_dist": None,
                "max_dist": None,
                "status": "失败",
                "message": message
            })
            continue

        try:
            result = compute_point_CAD_distance(
                CAD_path=str(cad_path),
                pointcloud_path=str(pointcloud_path),
                mesh_deflection=mesh_deflection,
                point_chunk_size=point_chunk_size,
                triangle_chunk_size=triangle_chunk_size,
                visualize=False
            )
        except Exception as exc:
            message = str(exc)
            print(f"任务 {task_id} 失败: {message}")
            result_rows.append({
                "task_id": task_id,
                "mean_dist": None,
                "max_dist": None,
                "status": "失败",
                "message": message
            })
            continue

        task_distance_data[task_id] = {
            "points_np": result["points_np"],
            "distances_np": result["distances_np"]
        }

        result_rows.append({
            "task_id": task_id,
            "mean_dist": result["mean_dist"],
            "max_dist": result["max_dist"],
            "status": "完成",
            "message": ""
        })

    excel_path = write_flatness_excel(
        result_rows,
        excel_output,
        sheet_name="平整度"
    )

    print(f"Excel 已保存: {excel_path}")

    return result_rows, task_distance_data


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

if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="计算点云到 CAD STEP 模型的最近距离，支持单任务或批量任务。"
    )

    parser.add_argument(
        "--CAD",
        "-c",
        default="High_Speed_Data/CAD/11.step",
        help="CAD STEP 文件路径，仅在 --single 单任务模式下使用。"
    )

    parser.add_argument(
        "--pointcloud",
        "-p",
        default="High_Speed_Data/Noise_PointCloud/11_noisy.ply",
        help="点云 PLY 文件路径，仅在 --single 单任务模式下使用。"
    )

    parser.add_argument(
        "--tasks",
        default=",".join(str(task_id) for task_id in DEFAULT_TASK_IDS),
        help="批量任务号，多个任务用英文逗号分隔。"
    )

    parser.add_argument(
        "--data_root",
        default="High_Speed_Data",
        help="批量任务的数据根目录。"
    )

    parser.add_argument(
        "--pointcloud_kind",
        choices=["noise", "clean"],
        default="noise",
        help="批量任务点云来源：noise=Noise_PointCloud，clean=PointCloud。"
    )

    parser.add_argument(
        "--excel_output",
        default="平整度.xlsx",
        help="批量结果 Excel 输出路径。"
    )

    parser.add_argument(
        "--single",
        action="store_true",
        help="使用 --CAD 和 --pointcloud 执行单任务，而不是默认批量任务。"
    )

    parser.add_argument(
        "--visualize",
        action="store_true",
        help="显示单任务热力图。批量模式不会自动可视化。"
    )

    parser.add_argument(
        "--mesh_deflection",
        default=0.1,
        type=float,
        help="CAD 三角化精度，越小越精细，但计算越慢。"
    )

    parser.add_argument(
        "--point_chunk_size",
        default=2048,
        type=int,
        help="点云分块大小，显存不够时可调小。"
    )

    parser.add_argument(
        "--triangle_chunk_size",
        default=4096,
        type=int,
        help="三角面分块大小，显存不够时可调小。"
    )

    args = parser.parse_args()

    if args.single:
        result = compute_point_CAD_distance(
            CAD_path=args.CAD,
            pointcloud_path=args.pointcloud,
            mesh_deflection=args.mesh_deflection,
            point_chunk_size=args.point_chunk_size,
            triangle_chunk_size=args.triangle_chunk_size,
            visualize=args.visualize
        )

        print("===== 计算结果 =====")
        print(f"平均距离: {result['mean_dist']:.8f} mm")
        print(f"最大距离: {result['max_dist']:.8f} mm")
    else:
        task_ids = [
            int(task_id.strip())
            for task_id in args.tasks.split(",")
            if task_id.strip()
        ]

        result_rows, task_distance_data = run_batch_tasks(
            task_ids=task_ids,
            data_root=args.data_root,
            pointcloud_kind=args.pointcloud_kind,
            excel_output=args.excel_output,
            mesh_deflection=args.mesh_deflection,
            point_chunk_size=args.point_chunk_size,
            triangle_chunk_size=args.triangle_chunk_size
        )

        print("===== 批量计算结果 =====")
        for row in result_rows:
            if row["status"] == "完成":
                print(
                    f"{row['task_id']}: "
                    f"平均距离={row['mean_dist']:.8f} mm, "
                    f"最大距离={row['max_dist']:.8f} mm"
                )
            else:
                print(f"{row['task_id']}: {row['status']} - {row['message']}")

        print("task_distance_data 已按任务号保存 points_np 和 distances_np。")