import argparse
import posixpath
import time
from pathlib import Path
from zipfile import ZipFile, ZIP_DEFLATED
from xml.sax.saxutils import escape
from xml.etree import ElementTree as ET
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
]

DEFAULT_CANDIDATE_TRIANGLE_COUNT = 32
DEFAULT_RESULT_TEMPLATE = Path("excel/测量结果.xlsx")

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 read_mesh_file(mesh_path):
    mesh_path = Path(mesh_path)

    if o3d is not None:
        mesh = o3d.io.read_triangle_mesh(str(mesh_path))
        if mesh.has_vertices() and mesh.has_triangles():
            vertices_np = np.asarray(mesh.vertices, dtype=np.float64)
            faces_np = np.asarray(mesh.triangles, dtype=np.int64)
            print("Mesh 文件读取完成")
            print(f"Mesh 顶点: {vertices_np.shape}")
            print(f"Mesh 三角面: {faces_np.shape}")
            return vertices_np, faces_np

    mesh = trimesh.load(str(mesh_path), force="mesh")
    vertices_np = np.asarray(mesh.vertices, dtype=np.float64)
    faces_np = np.asarray(mesh.faces, dtype=np.int64)

    if vertices_np.shape[0] == 0 or faces_np.shape[0] == 0:
        raise RuntimeError(f"Mesh 文件读取失败或为空: {mesh_path}")

    print("Mesh 文件读取完成")
    print(f"Mesh 顶点: {vertices_np.shape}")
    print(f"Mesh 三角面: {faces_np.shape}")

    return vertices_np, faces_np


def read_model_mesh(model_path, mesh_deflection):
    model_path = Path(model_path)
    suffix = model_path.suffix.lower()

    if suffix in {".step", ".stp"}:
        shape = read_step_file(model_path)

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

        return occ_shape_to_numpy_mesh(shape)

    if suffix in {".stl", ".obj", ".ply"}:
        return read_mesh_file(model_path)

    raise ValueError(f"不支持的模型文件格式: {model_path}")


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

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

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

    p = points[:, None, :]

    if triangles.dim() == 3:
        a = triangles[None, :, 0, :]
        b = triangles[None, :, 1, :]
        c = triangles[None, :, 2, :]
    elif triangles.dim() == 4:
        a = triangles[:, :, 0, :]
        b = triangles[:, :, 1, :]
        c = triangles[:, :, 2, :]
    else:
        raise ValueError(
            f"triangles must have shape [T, 3, 3] or [P, K, 3, 3], got {tuple(triangles.shape)}"
        )

    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)

    return torch.where(torch.isinf(dist2), face_dist2, dist2)


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

def nearest_aabb_candidate_indices(points, triangle_mins, triangle_maxs, candidate_count):
    """
    参考 test8.py，用点到三角形 AABB 的距离先筛出每个点附近的候选三角形。

    points:
        [P, 3]
    triangle_mins / triangle_maxs:
        [M, 3]

    返回:
        candidate_indices: [P, K]
    """

    below = (triangle_mins[None, :, :] - points[:, None, :]).clamp_min(0.0)
    above = (points[:, None, :] - triangle_maxs[None, :, :]).clamp_min(0.0)
    aabb_dist2 = ((below + above) ** 2).sum(dim=-1)
    k = min(candidate_count, triangle_mins.shape[0])
    return torch.topk(aabb_dist2, k=k, dim=1, largest=False).indices


def nearest_mesh_distances(
    points,
    triangles,
    point_chunk_size=2048,
    triangle_chunk_size=4096,
    candidate_triangle_count=DEFAULT_CANDIDATE_TRIANGLE_COUNT
):
    """
    分块计算点云到三角网格的最近距离。

    默认使用 test8.py 的 AABB 候选三角形筛选方法：
        1. 先按点到三角形包围盒的距离选出 K 个候选三角形；
        2. 再对候选三角形计算精确点到三角形距离。

    如果 candidate_triangle_count <= 0，则退回旧的全量三角面分块精确搜索。

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

    返回:
        distances: [N]
    """

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

    if candidate_triangle_count > 0:
        triangle_mins = triangles.min(dim=1).values
        triangle_maxs = triangles.max(dim=1).values

        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]

            candidate_idx = nearest_aabb_candidate_indices(
                point_chunk,
                triangle_mins,
                triangle_maxs,
                candidate_triangle_count
            )
            candidate_triangles = triangles[candidate_idx]
            dist2 = point_to_triangle_distance_squared(
                point_chunk,
                candidate_triangles
            )
            all_distances.append(torch.sqrt(dist2.min(dim=1).values))

        return torch.cat(all_distances, dim=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))

    return torch.cat(all_distances, dim=0)


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

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

    pcd_show = build_heatmap_point_cloud(points_np, distances_np)
    print("色阶规则：深蓝/紫色 = 最近，红色 = 最远")

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


def build_heatmap_point_cloud(points_np, distances_np):
    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)

    return pcd_show


def count_task_result_points(task_results):
    total = 0
    for task_id in sorted(task_results, key=task_sort_key):
        result = task_results[task_id]
        for sample_result in result.get("sample_results", []):
            total += int(len(sample_result["points_np"]))
    return total


def read_binary_xyz_ply_header(ply_path):
    with open(ply_path, "rb") as file:
        header_lines = []
        while True:
            line = file.readline()
            if not line:
                raise ValueError(f"Invalid PLY header: {ply_path}")
            header_lines.append(line)
            if line.strip() == b"end_header":
                break

        header_size = file.tell()

    text_lines = [
        line.decode("ascii", errors="replace").strip()
        for line in header_lines
    ]

    if len(text_lines) < 3 or text_lines[0] != "ply":
        raise ValueError(f"Not a PLY file: {ply_path}")

    if text_lines[1] != "format binary_little_endian 1.0":
        raise ValueError(
            "Only binary_little_endian PLY context point clouds can be streamed."
        )

    vertex_count = None
    vertex_properties = []
    in_vertex = False
    for line in text_lines[2:]:
        parts = line.split()
        if not parts:
            continue
        if parts[0] == "element":
            in_vertex = parts[1] == "vertex"
            if in_vertex:
                vertex_count = int(parts[2])
            continue
        if in_vertex and parts[0] == "property":
            if len(parts) >= 2 and parts[1] == "list":
                raise ValueError("List properties are not supported in streamed PLY.")
            if len(parts) >= 3:
                vertex_properties.append((parts[1], parts[2]))

    expected = [("float", "x"), ("float", "y"), ("float", "z")]
    if vertex_count is None or vertex_properties[:3] != expected:
        raise ValueError(
            "Context PLY must start with vertex properties: float x, float y, float z."
        )
    if len(vertex_properties) != 3:
        raise ValueError(
            "Context PLY must contain only float x/y/z vertices for streamed merge."
        )

    return vertex_count, header_size


def write_xyz_rgb_chunk(file, points, colors):
    points = np.asarray(points)
    colors = np.asarray(colors)
    count = len(points)
    if count == 0:
        return

    vertex_dtype = np.dtype([
        ("x", "<f4"),
        ("y", "<f4"),
        ("z", "<f4"),
        ("red", "u1"),
        ("green", "u1"),
        ("blue", "u1"),
    ])
    vertices = np.empty(count, dtype=vertex_dtype)
    vertices["x"] = points[:, 0].astype(np.float32, copy=False)
    vertices["y"] = points[:, 1].astype(np.float32, copy=False)
    vertices["z"] = points[:, 2].astype(np.float32, copy=False)

    if colors.ndim == 1:
        vertices["red"] = int(colors[0])
        vertices["green"] = int(colors[1])
        vertices["blue"] = int(colors[2])
    else:
        colors = colors.astype(np.uint8, copy=False)
        vertices["red"] = colors[:, 0]
        vertices["green"] = colors[:, 1]
        vertices["blue"] = colors[:, 2]

    file.write(vertices.tobytes())


def write_context_xyz_ply_stream(input_path, output_file, header_size, point_count):
    record_size = 12
    chunk_points = 1_000_000
    gray = np.array([158, 158, 158], dtype=np.uint8)

    with open(input_path, "rb") as input_file:
        input_file.seek(header_size)
        remaining = point_count
        while remaining > 0:
            current = min(chunk_points, remaining)
            data = input_file.read(current * record_size)
            if len(data) != current * record_size:
                raise IOError(f"Unexpected end of context PLY: {input_path}")
            points = np.frombuffer(data, dtype="<f4").reshape(-1, 3)
            write_xyz_rgb_chunk(output_file, points, gray)
            remaining -= current


def write_task_results_stream(task_results, output_file):
    chunk_points = 1_000_000
    for task_id in sorted(task_results, key=task_sort_key):
        result = task_results[task_id]
        for sample_result in result.get("sample_results", []):
            points_np = np.asarray(sample_result["points_np"])
            distances_np = np.asarray(sample_result["distances_np"])
            if len(points_np) == 0:
                continue

            d_min = float(distances_np.min())
            d_max = float(distances_np.max())
            denom = d_max - d_min + 1e-12

            for start in range(0, len(points_np), chunk_points):
                end = min(start + chunk_points, len(points_np))
                dist_norm = (distances_np[start:end] - d_min) / denom
                colors = (plt.cm.turbo(dist_norm)[:, :3] * 255.0 + 0.5).astype(np.uint8)
                write_xyz_rgb_chunk(output_file, points_np[start:end], colors)


def save_combined_successful_results_stream(
    task_results,
    output_path,
    context_pointcloud_path=None
):
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    context_point_count = 0
    context_header_size = None
    context_path = Path(context_pointcloud_path) if context_pointcloud_path else None
    if context_path:
        if context_path.exists():
            context_point_count, context_header_size = read_binary_xyz_ply_header(context_path)
            print(f"完整点云将以流式方式加入合并结果: {context_path}")
        else:
            print(f"完整点云不存在，跳过: {context_path}")
            context_path = None

    task_point_count = count_task_result_points(task_results)
    total_point_count = context_point_count + task_point_count
    if total_point_count == 0:
        print("没有成功执行的可视化结果，跳过合并点云保存。")
        return None

    header = (
        "ply\n"
        "format binary_little_endian 1.0\n"
        f"element vertex {total_point_count}\n"
        "property float x\n"
        "property float y\n"
        "property float z\n"
        "property uchar red\n"
        "property uchar green\n"
        "property uchar blue\n"
        "end_header\n"
    )

    with open(output_path, "wb") as output_file:
        output_file.write(header.encode("ascii"))
        if context_path:
            write_context_xyz_ply_stream(
                context_path,
                output_file,
                context_header_size,
                context_point_count
            )
        write_task_results_stream(task_results, output_file)

    print(f"合并点云已保存: {output_path}")
    return output_path


def build_combined_successful_point_cloud(
    task_results,
    context_pointcloud_path="datasets/train/train_2_10M.ply",
    arrange_grid=False
):
    if o3d is None:
        return None

    task_geometries = []
    for task_id in sorted(task_results, key=task_sort_key):
        result = task_results[task_id]
        for sample_result in result.get("sample_results", []):
            pcd_show = build_heatmap_point_cloud(
                sample_result["points_np"],
                sample_result["distances_np"]
            )
            task_geometries.append(pcd_show)

    if not task_geometries:
        return None

    geometries = []
    if context_pointcloud_path:
        context_path = Path(context_pointcloud_path)
        if context_path.exists():
            context = o3d.io.read_point_cloud(str(context_path))
            if len(context.points) > 0:
                if not context.has_colors():
                    context.paint_uniform_color([0.62, 0.62, 0.62])
                geometries.append(context)
                print(f"完整点云已加入合并结果: {context_path}")
        else:
            print(f"完整点云不存在，跳过: {context_path}")

    geometries.extend(task_geometries)

    if arrange_grid and len(geometries) > 1:
        extents = [
            np.asarray(geometry.get_max_bound()) - np.asarray(geometry.get_min_bound())
            for geometry in geometries
        ]
        max_extent = max(float(np.max(extent)) for extent in extents)
        spacing = max(max_extent * 1.35, 1.0)
        cols = int(np.ceil(np.sqrt(len(geometries))))

        for index, geometry in enumerate(geometries):
            row, col = divmod(index, cols)
            center = (
                np.asarray(geometry.get_min_bound())
                + np.asarray(geometry.get_max_bound())
            ) / 2.0
            target_center = np.array([col * spacing, -row * spacing, 0.0])
            geometry.translate(target_center - center)

    combined = o3d.geometry.PointCloud()
    for geometry in geometries:
        if not geometry.has_colors():
            geometry.paint_uniform_color([0.62, 0.62, 0.62])
        combined += geometry
    return combined


def save_all_successful_results(
    task_results,
    output_path="vis/combined_contour.ply",
    context_pointcloud_path="datasets/train/train_2_10M.ply"
):
    try:
        return save_combined_successful_results_stream(
            task_results,
            output_path=output_path,
            context_pointcloud_path=context_pointcloud_path
        )
    except Exception as exc:
        context_path = Path(context_pointcloud_path) if context_pointcloud_path else None
        if context_path and context_path.exists() and context_path.stat().st_size > 1024 ** 3:
            raise RuntimeError(
                f"大文件流式合并失败，无法回退到 Open3D 整体加载: {exc}"
            ) from exc
        print(f"流式合并失败，回退到 Open3D 合并: {exc}")

    if o3d is None:
        print("未安装 open3d，跳过合并点云保存。")
        return None

    combined = build_combined_successful_point_cloud(
        task_results,
        context_pointcloud_path=context_pointcloud_path
    )
    if combined is None or len(combined.points) == 0:
        print("没有成功执行的可视化结果，跳过合并点云保存。")
        return None

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    ok = o3d.io.write_point_cloud(str(output_path), combined)
    if not ok:
        raise IOError(f"合并点云保存失败: {output_path}")

    print(f"合并点云已保存: {output_path}")
    return output_path


def visualize_all_successful_results(
    task_results,
    context_pointcloud_path="datasets/train/train_2_10M.ply"
):
    if o3d is None:
        print("未安装 open3d，跳过可视化。")
        return

    combined = build_combined_successful_point_cloud(
        task_results,
        context_pointcloud_path=context_pointcloud_path
    )
    if combined is None or len(combined.points) == 0:
        print("没有成功执行的可视化结果。")
        return

    print("色阶规则：深蓝/紫色 = 最近，红色 = 最远")
    o3d.visualization.draw_geometries(
        [combined],
        window_name="批量点云到 CAD 距离热力图",
        width=1280,
        height=900
    )


def visualize_task_result(task_results, task_id, sample_id=None):
    """
    task_results: run_batch_tasks 返回的字典，key 为任务号。
    task_id: 需要可视化的任务号。
    sample_id: 可选样本名，例如 13_2。不传时显示该任务第一个成功样本。
    """

    result = task_results.get(str(task_id)) or task_results.get(int(task_id))

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

    if "sample_results" in result:
        sample_results = result["sample_results"]
        if not sample_results:
            raise KeyError(f"任务 {task_id} 没有可视化数据。")

        if sample_id is None:
            sample_result = sample_results[0]
        else:
            sample_result = next(
                (
                    item for item in sample_results
                    if item.get("sample_id") == str(sample_id)
                ),
                None
            )
            if sample_result is None:
                raise KeyError(f"任务 {task_id} 中未找到样本 {sample_id}。")

        visualize_point_cloud_heatmap(
            sample_result["points_np"],
            sample_result["distances_np"]
        )
        return

    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 _xlsx_col_to_index(cell_ref):
    letters = "".join(ch for ch in cell_ref if ch.isalpha())
    index = 0
    for letter in letters:
        index = index * 26 + ord(letter.upper()) - ord("A") + 1
    return index


def _xlsx_rel_target_to_path(target):
    target = target.replace("\\", "/").lstrip("/")
    if target.startswith("xl/"):
        return target
    return posixpath.normpath(posixpath.join("xl", target))


def _xlsx_cell_value(cell, shared_strings, ns):
    value_node = cell.find("a:v", ns)
    inline_node = cell.find("a:is/a:t", ns)
    cell_type = cell.attrib.get("t")

    if cell_type == "s" and value_node is not None:
        return shared_strings[int(value_node.text)]
    if cell_type == "inlineStr" and inline_node is not None:
        return inline_node.text or ""
    if value_node is not None:
        return value_node.text or ""
    return ""


def read_measurement_template(template_path):
    template_path = Path(template_path)
    ns = {
        "a": "http://schemas.openxmlformats.org/spreadsheetml/2006/main",
        "r": "http://schemas.openxmlformats.org/officeDocument/2006/relationships",
    }

    with ZipFile(template_path) as xlsx:
        names = xlsx.namelist()

        shared_strings = []
        if "xl/sharedStrings.xml" in names:
            shared_root = ET.fromstring(xlsx.read("xl/sharedStrings.xml"))
            for item in shared_root.findall("a:si", ns):
                shared_strings.append(
                    "".join((text.text or "") for text in item.findall(".//a:t", ns))
                )

        workbook_root = ET.fromstring(xlsx.read("xl/workbook.xml"))
        first_sheet = workbook_root.find("a:sheets/a:sheet", ns)
        if first_sheet is None:
            raise RuntimeError(f"模板中没有工作表: {template_path}")

        sheet_name = first_sheet.attrib.get("name", "Sheet1")
        relation_id = first_sheet.attrib[
            "{http://schemas.openxmlformats.org/officeDocument/2006/relationships}id"
        ]

        rels_root = ET.fromstring(xlsx.read("xl/_rels/workbook.xml.rels"))
        relation_targets = {
            rel.attrib["Id"]: rel.attrib["Target"]
            for rel in rels_root
        }
        sheet_path = _xlsx_rel_target_to_path(relation_targets[relation_id])

        sheet_root = ET.fromstring(xlsx.read(sheet_path))
        parsed_rows = []
        max_col_count = 0

        for row in sheet_root.findall("a:sheetData/a:row", ns):
            values_by_col = {}
            for cell in row.findall("a:c", ns):
                col_index = _xlsx_col_to_index(cell.attrib["r"])
                values_by_col[col_index] = _xlsx_cell_value(cell, shared_strings, ns)
                max_col_count = max(max_col_count, col_index)

            if values_by_col:
                parsed_rows.append(values_by_col)

        table = []
        for values_by_col in parsed_rows:
            table.append([
                values_by_col.get(col_index, "")
                for col_index in range(1, max_col_count + 1)
            ])

    if not table:
        raise RuntimeError(f"模板为空: {template_path}")

    return sheet_name, table[0], table[1:]


def _normalize_template_task_id(value):
    text = str(value).strip()
    if text.endswith(".0") and text[:-2].isdigit():
        return text[:-2]
    return text


def _find_header_index(headers, candidates):
    normalized = [str(header).strip() for header in headers]
    for candidate in candidates:
        if candidate in normalized:
            return normalized.index(candidate)
    raise ValueError(f"模板缺少列: {'/'.join(candidates)}")


def _row_measurement_conclusion(row):
    mean_dist = row.get("mean_dist")
    if mean_dist is None:
        return "失败"
    try:
        if not np.isfinite(float(mean_dist)):
            return "失败"
    except (TypeError, ValueError):
        return "失败"
    return "成功"


def _build_measurement_table_from_template(rows, template_path):
    template_sheet_name, headers, template_rows = read_measurement_template(template_path)
    task_col = _find_header_index(headers, ["任务号", "任务编号", "task_id"])
    result_col = _find_header_index(headers, ["测量结果"])
    conclusion_col = _find_header_index(headers, ["测量结论", "结论", "状态"])

    result_by_task = {
        _normalize_template_task_id(row.get("task_id")): row
        for row in rows
    }
    used_task_ids = set()
    table = [headers]

    for template_row in template_rows:
        output_row = list(template_row)
        if len(output_row) < len(headers):
            output_row.extend([""] * (len(headers) - len(output_row)))

        task_id = _normalize_template_task_id(output_row[task_col])
        result_row = result_by_task.get(task_id)
        if result_row is None:
            continue

        output_row[result_col] = result_row.get("mean_dist")
        output_row[conclusion_col] = _row_measurement_conclusion(result_row)
        table.append(output_row)
        used_task_ids.add(task_id)

    for task_id, result_row in result_by_task.items():
        if task_id in used_task_ids:
            continue

        output_row = [""] * len(headers)
        output_row[task_col] = task_id
        output_row[result_col] = result_row.get("mean_dist")
        output_row[conclusion_col] = _row_measurement_conclusion(result_row)
        table.append(output_row)

    return template_sheet_name, table


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

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

    if result_template_path and Path(result_template_path).exists():
        sheet_name, table = _build_measurement_table_from_template(
            rows,
            result_template_path
        )
    else:
        headers = ["任务号", "样本数", "平均距离(mm)", "最大距离(mm)", "状态", "说明"]
        table = [headers]

        for row in rows:
            table.append([
                row.get("task_id"),
                row.get("sample_count"),
                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,
    candidate_triangle_count=DEFAULT_CANDIDATE_TRIANGLE_COUNT,
    visualize=False
):
    print("计算设备:", device)
    print("数据精度:", dtype)

    # 读取 STEP CAD 或 STL/OBJ/PLY Mesh，并统一转换为 numpy 三角网格。
    vertices_np, faces_np = read_model_mesh(CAD_path, mesh_deflection)

    # 读取点云
    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("开始计算点云到模型表面的最近距离...")
    print("点云张量:", points.shape)
    print("三角面张量:", triangles.shape)
    print("候选三角面数量:", candidate_triangle_count)

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

    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 task_id_from_sample_name(sample_name):
    return str(sample_name).split("_", 1)[0]


def normalize_task_id(task_id):
    return str(task_id).strip()


def task_sort_key(value):
    base = task_id_from_sample_name(value)
    if str(base).isdigit():
        return (0, int(base), str(value))
    return (1, str(base), str(value))


def pointcloud_path_for_sample(data_root, sample_id, pointcloud_kind):
    data_root = Path(data_root)

    if pointcloud_kind == "clean":
        return data_root / "PointCloud" / f"{sample_id}.ply"
    if pointcloud_kind == "noise":
        return data_root / "Noise_PointCloud" / f"{sample_id}_noisy.ply"

    raise ValueError("pointcloud_kind must be 'noise' or 'clean'.")


def mesh_sample_paths_for_task(data_root, requested_task):
    data_root = Path(data_root)
    mesh_dir = data_root / "Mesh"
    requested_task = normalize_task_id(requested_task)

    if not mesh_dir.exists():
        return []

    mesh_paths = []
    for pattern in ("*.stl", "*.obj", "*.ply"):
        mesh_paths.extend(mesh_dir.glob(pattern))

    if "_" in requested_task:
        matched = [
            path for path in mesh_paths
            if path.stem == requested_task
        ]
    else:
        matched = [
            path for path in mesh_paths
            if task_id_from_sample_name(path.stem) == requested_task
        ]

    return sorted(
        matched,
        key=lambda path: task_sort_key(path.stem)
    )


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

    mesh_paths = mesh_sample_paths_for_task(data_root, task_id)
    if mesh_paths:
        return [
            {
                "task_id": task_id_from_sample_name(mesh_path.stem),
                "sample_id": mesh_path.stem,
                "model_path": mesh_path,
                "pointcloud_path": pointcloud_path_for_sample(
                    data_root,
                    mesh_path.stem,
                    pointcloud_kind
                ),
            }
            for mesh_path in mesh_paths
        ]

    sample_id = task_id
    return [
        {
            "task_id": task_id_from_sample_name(sample_id),
            "sample_id": sample_id,
            "model_path": data_root / "CAD" / f"{sample_id}.step",
            "pointcloud_path": pointcloud_path_for_sample(
                data_root,
                sample_id,
                pointcloud_kind
            ),
        }
    ]

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,
    candidate_triangle_count=DEFAULT_CANDIDATE_TRIANGLE_COUNT,
    result_template_path=DEFAULT_RESULT_TEMPLATE
):
    global task_distance_data

    task_summaries = {}
    task_distance_data = {}

    for task_id in task_ids:
        task_inputs = get_task_paths(
            task_id,
            data_root=data_root,
            pointcloud_kind=pointcloud_kind
        )

        for task_input in task_inputs:
            base_task_id = task_input["task_id"]
            sample_id = task_input["sample_id"]
            model_path = task_input["model_path"]
            pointcloud_path = task_input["pointcloud_path"]

            summary = task_summaries.setdefault(
                base_task_id,
                {
                    "task_id": base_task_id,
                    "sample_ids": [],
                    "point_count": 0,
                    "distance_sum": 0.0,
                    "max_dist": None,
                    "failures": [],
                    "sample_results": [],
                }
            )
            summary["sample_ids"].append(sample_id)

            print(f"===== 任务 {base_task_id} / 样本 {sample_id} =====")

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

            if missing_paths:
                message = f"{sample_id} 缺少文件: " + "; ".join(missing_paths)
                print(message)
                summary["failures"].append(message)
                continue

            try:
                result = compute_point_CAD_distance(
                    CAD_path=str(model_path),
                    pointcloud_path=str(pointcloud_path),
                    mesh_deflection=mesh_deflection,
                    point_chunk_size=point_chunk_size,
                    triangle_chunk_size=triangle_chunk_size,
                    candidate_triangle_count=candidate_triangle_count,
                    visualize=False
                )
            except Exception as exc:
                message = f"{sample_id} 失败: {exc}"
                print(message)
                summary["failures"].append(message)
                continue

            distances_np = result["distances_np"]
            point_count = len(distances_np)
            summary["point_count"] += point_count
            summary["distance_sum"] += float(distances_np.sum())
            summary["max_dist"] = (
                result["max_dist"]
                if summary["max_dist"] is None
                else max(summary["max_dist"], result["max_dist"])
            )
            summary["sample_results"].append({
                "sample_id": sample_id,
                "model_path": str(model_path),
                "pointcloud_path": str(pointcloud_path),
                "points_np": result["points_np"],
                "distances_np": distances_np,
                "mean_dist": result["mean_dist"],
                "max_dist": result["max_dist"],
                "min_dist": result["min_dist"],
            })

    result_rows = []
    for task_id in sorted(
        task_summaries,
        key=task_sort_key
    ):
        summary = task_summaries[task_id]
        success_count = len(summary["sample_results"])
        failure_count = len(summary["failures"])

        if success_count:
            mean_dist = summary["distance_sum"] / summary["point_count"]
            max_dist = summary["max_dist"]
            status = "完成" if failure_count == 0 else "部分完成"
        else:
            mean_dist = None
            max_dist = None
            status = "失败"

        message_parts = []
        if summary["sample_ids"]:
            message_parts.append("样本: " + ",".join(summary["sample_ids"]))
        if summary["failures"]:
            message_parts.append("失败: " + " | ".join(summary["failures"]))

        result_rows.append({
            "task_id": task_id,
            "sample_count": success_count,
            "mean_dist": mean_dist,
            "max_dist": max_dist,
            "status": status,
            "message": "; ".join(message_parts)
        })

        task_distance_data[task_id] = {
            "sample_results": summary["sample_results"],
            "sample_ids": [item["sample_id"] for item in summary["sample_results"]],
        }

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

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

    return result_rows, task_distance_data


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

if __name__ == "__main__":
    start_time = time.perf_counter()

    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="批量任务号或样本名，多个值用英文逗号分隔。例：13 会匹配 13_1/13_2；13_1 只匹配该样本。"
    )

    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(
        "--result_template",
        default=str(DEFAULT_RESULT_TEMPLATE),
        help="测量结果模板 Excel 路径。新表会按任务号复制模板对应行，并写入测量结果和测量结论。"
    )

    parser.add_argument(
        "--combined_output",
        default="vis/combined_contour.ply",
        help="批量模式下所有成功热力图合并后的 PLY 输出路径。"
    )

    parser.add_argument(
        "--combined_context_pointcloud",
        default="datasets/train/train_2_10M.ply",
        help="批量合并热力图时一起加入的完整点云路径，可切换为 train_2_1T.ply 或 train_2_100M.ply。"
    )

    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="全量精确搜索模式下的三角面分块大小；candidate_triangle_count <= 0 时使用。"
    )

    parser.add_argument(
        "--candidate_triangle_count",
        default=DEFAULT_CANDIDATE_TRIANGLE_COUNT,
        type=int,
        help="每个点按 AABB 预筛的候选三角面数量；设为 0 或负数则使用全量精确搜索。"
    )

    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,
            candidate_triangle_count=args.candidate_triangle_count,
            visualize=args.visualize
        )

        print("===== 计算结果 =====")
        print(f"平均距离: {result['mean_dist']:.8f} mm")
        print(f"最大距离: {result['max_dist']:.8f} mm")
        elapsed_time = time.perf_counter() - start_time
        print(f"运行时间: {elapsed_time:.2f} 秒")
    else:
        task_ids = [
            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,
            candidate_triangle_count=args.candidate_triangle_count,
            result_template_path=args.result_template
        )

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

        elapsed_time = time.perf_counter() - start_time
        print(f"运行时间: {elapsed_time:.2f} 秒")

        print("task_distance_data 已按任务号保存 sample_results。")
        save_all_successful_results(
            task_distance_data,
            output_path=args.combined_output,
            context_pointcloud_path=args.combined_context_pointcloud
        )
        if args.visualize:
            visualize_all_successful_results(
                task_distance_data,
                context_pointcloud_path=args.combined_context_pointcloud
            )