import gc
import time
import zipfile
from pathlib import Path
from xml.sax.saxutils import escape as xml_escape

import numpy as np
import open3d as o3d
import torch


MESH_DIR = Path("Mesh")
POINT_CLOUD_DIR = Path("CloudPoint_3000000")
RESULT_DIR = Path("result")
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.float64
QUERY_CHUNK_SIZE = 8192 if DEVICE.type == "cuda" else 4096
CANDIDATE_TRIANGLE_COUNT = 32
SUMMARY_XLSX_PATH = RESULT_DIR / "task_distance_summary.xlsx"
SUMMARY_CSV_PATH = RESULT_DIR / "task_distance_summary.csv"
OPEN_VISUALIZATION_WINDOW = False
COLOR_MAX_PERCENTILE = 99.0
POINT_SIZE = 2.0


def read_triangle_mesh(stl_path):
    mesh = o3d.io.read_triangle_mesh(stl_path)
    if mesh.is_empty():
        raise ValueError(f"Failed to read mesh: {stl_path}")
    if not mesh.has_triangles():
        raise ValueError(f"Mesh has no triangles: {stl_path}")
    return mesh


def read_point_cloud_points(ply_path):
    pcd = o3d.io.read_point_cloud(ply_path)
    if pcd.is_empty():
        raise ValueError(f"Failed to read point cloud: {ply_path}")
    return np.asarray(pcd.points, dtype=np.float64).copy()


def point_to_triangle_distance_squared(points, triangles):
    p = points[:, None, :]
    a = triangles[:, :, 0, :]
    b = triangles[:, :, 1, :]
    c = triangles[:, :, 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)

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

    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))

    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))

    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))

    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))

    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):
    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 compute_mesh_distances(mesh, points_np, chunk_size, candidate_count):
    vertices = np.asarray(mesh.vertices, dtype=np.float64)
    triangles_idx = np.asarray(mesh.triangles, dtype=np.int64)
    points = torch.as_tensor(points_np, dtype=DTYPE, device=DEVICE)
    vertices = torch.as_tensor(vertices, dtype=DTYPE, device=DEVICE)
    triangles_idx = torch.as_tensor(triangles_idx, dtype=torch.long, device=DEVICE)
    triangles = vertices[triangles_idx]
    triangle_mins = triangles.min(dim=1).values
    triangle_maxs = triangles.max(dim=1).values

    distances = []
    with torch.no_grad():
        for start in range(0, points.shape[0], chunk_size):
            chunk = points[start : start + chunk_size]
            candidate_idx = nearest_aabb_candidate_indices(
                chunk,
                triangle_mins,
                triangle_maxs,
                candidate_count,
            )
            candidate_triangles = triangles[candidate_idx]
            dist2 = point_to_triangle_distance_squared(chunk, candidate_triangles)
            chunk_distances = torch.sqrt(dist2.min(dim=1).values)
            distances.append(chunk_distances.cpu())

    if DEVICE.type == "cuda":
        torch.cuda.synchronize()

    return torch.cat(distances, dim=0).numpy()


def task_id_from_mesh_path(mesh_path):
    return mesh_path.stem.split("_", 1)[0]


def discover_mesh_tasks(mesh_dir, point_cloud_dir):
    mesh_paths = sorted(
        mesh_dir.glob("*.stl"),
        key=lambda path: (int(task_id_from_mesh_path(path)), path.stem),
    )
    if not mesh_paths:
        raise FileNotFoundError(f"No STL files found in {mesh_dir}")

    tasks = {}
    for mesh_path in mesh_paths:
        point_cloud_path = point_cloud_dir / f"{mesh_path.stem}.ply"
        if not point_cloud_path.exists():
            raise FileNotFoundError(
                f"Missing point cloud for {mesh_path.name}: {point_cloud_path}"
            )
        task_id = task_id_from_mesh_path(mesh_path)
        tasks.setdefault(task_id, []).append((mesh_path, point_cloud_path))
    return tasks


def distance_colors(distances):
    finite_mask = np.isfinite(distances)
    if not finite_mask.any():
        raise ValueError("No finite distances to visualize.")

    finite_distances = distances[finite_mask]
    color_min = float(finite_distances.min())
    if COLOR_MAX_PERCENTILE is None:
        color_max = float(finite_distances.max())
    else:
        color_max = float(np.percentile(finite_distances, COLOR_MAX_PERCENTILE))

    if color_max <= color_min:
        color_max = float(finite_distances.max())
    if color_max <= color_min:
        normalized = np.zeros_like(distances, dtype=np.float64)
    else:
        normalized = (distances - color_min) / (color_max - color_min)
        normalized = np.clip(normalized, 0.0, 1.0)

    palette_positions = np.array([0.0, 0.25, 0.50, 0.75, 1.0], dtype=np.float64)
    palette_colors = np.array(
        [
            [0.05, 0.12, 0.75],
            [0.00, 0.65, 1.00],
            [0.20, 0.90, 0.25],
            [1.00, 0.90, 0.00],
            [0.90, 0.05, 0.02],
        ],
        dtype=np.float64,
    )
    colors = np.column_stack(
        [
            np.interp(normalized, palette_positions, palette_colors[:, channel])
            for channel in range(3)
        ]
    )
    return colors, color_min, color_max


def write_visualization_stats(path, mesh_path, point_cloud_path, distances, color_min, color_max, runtime):
    stats_path = Path(path)
    stats_text = "\n".join(
        [
            f"mesh: {mesh_path}",
            f"point_cloud: {point_cloud_path}",
            f"point_count: {len(distances)}",
            f"min_distance: {distances.min():.8f}",
            f"max_distance: {distances.max():.8f}",
            f"mean_distance: {distances.mean():.8f}",
            f"median_distance: {np.median(distances):.8f}",
            f"p90_distance: {np.percentile(distances, 90):.8f}",
            f"p95_distance: {np.percentile(distances, 95):.8f}",
            f"p99_distance: {np.percentile(distances, 99):.8f}",
            f"color_min: {color_min:.8f}",
            f"color_max: {color_max:.8f}",
            f"runtime: {runtime:.4f}",
        ]
    )
    stats_path.write_text(stats_text + "\n", encoding="utf-8")


def visualize_distances(points_np, distances, output_path):
    if len(points_np) != len(distances):
        raise ValueError(
            f"Point count ({len(points_np)}) and distance count ({len(distances)}) do not match."
        )

    colors, color_min, color_max = distance_colors(distances)
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points_np)
    pcd.colors = o3d.utility.Vector3dVector(colors)

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    if not o3d.io.write_point_cloud(str(output_path), pcd):
        raise RuntimeError(f"Failed to write colored point cloud: {output_path}")

    if OPEN_VISUALIZATION_WINDOW:
        vis = o3d.visualization.Visualizer()
        vis.create_window(window_name="Point Distance Visualization")
        vis.add_geometry(pcd)
        render_option = vis.get_render_option()
        render_option.point_size = POINT_SIZE
        render_option.background_color = np.asarray([1.0, 1.0, 1.0])
        vis.run()
        vis.destroy_window()

    return color_min, color_max


def excel_column_name(column_index):
    name = ""
    while column_index:
        column_index, remainder = divmod(column_index - 1, 26)
        name = chr(65 + remainder) + name
    return name


def excel_cell(value, row_index, column_index):
    cell_ref = f"{excel_column_name(column_index)}{row_index}"
    if value is None:
        return f'<c r="{cell_ref}"/>'
    if isinstance(value, (int, float, np.integer, np.floating)):
        return f'<c r="{cell_ref}"><v>{float(value)}</v></c>'
    text = xml_escape(str(value))
    return f'<c r="{cell_ref}" t="inlineStr"><is><t>{text}</t></is></c>'


def write_summary_xlsx(path, rows):
    headers = [
        "Task",
        "Mesh count",
        "Point count",
        "Average distance",
        "Max distance",
        "Runtime",
    ]
    table_rows = [headers] + [
        [
            row["task_id"],
            row["mesh_count"],
            row["point_count"],
            row["average_distance"],
            row["max_distance"],
            row["runtime"],
        ]
        for row in rows
    ]

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

    worksheet_xml = f"""<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">
  <cols>
    <col min="1" max="1" width="12" customWidth="1"/>
    <col min="2" max="3" width="14" customWidth="1"/>
    <col min="4" max="6" width="18" customWidth="1"/>
  </cols>
  <sheetData>{''.join(sheet_rows)}</sheetData>
</worksheet>"""

    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">
  <sheets>
    <sheet name="Summary" sheetId="1" r:id="rId1"/>
  </sheets>
</workbook>"""

    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>"""

    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>"""

    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>"""

    path.parent.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(path, "w", compression=zipfile.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)


def write_summary_csv(path, rows):
    headers = [
        "Task",
        "Mesh count",
        "Point count",
        "Average distance",
        "Max distance",
        "Runtime",
    ]
    lines = [",".join(headers)]
    for row in rows:
        lines.append(
            ",".join(
                [
                    row["task_id"],
                    str(row["mesh_count"]),
                    str(row["point_count"]),
                    f"{row['average_distance']:.8f}",
                    f"{row['max_distance']:.8f}",
                    f"{row['runtime']:.4f}",
                ]
            )
        )
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def compute_one_mesh(mesh_path, point_cloud_path, task_output_dir):
    start_time = time.perf_counter()
    mesh = read_triangle_mesh(str(mesh_path))
    original_points_np = read_point_cloud_points(str(point_cloud_path))
    distances = compute_mesh_distances(
        mesh,
        original_points_np,
        QUERY_CHUNK_SIZE,
        CANDIDATE_TRIANGLE_COUNT,
    )

    task_output_dir.mkdir(parents=True, exist_ok=True)
    distance_output_path = task_output_dir / f"{mesh_path.stem}_distances.npy"
    colored_output_path = task_output_dir / f"{mesh_path.stem}_distance_colored.ply"
    stats_output_path = task_output_dir / f"{mesh_path.stem}_distance_stats.txt"

    np.save(distance_output_path, distances)
    color_min, color_max = visualize_distances(
        original_points_np,
        distances,
        colored_output_path,
    )
    runtime = time.perf_counter() - start_time
    write_visualization_stats(
        stats_output_path,
        mesh_path,
        point_cloud_path,
        distances,
        color_min,
        color_max,
        runtime,
    )

    return {
        "mesh_name": mesh_path.name,
        "point_count": len(distances),
        "distance_sum": float(distances.sum()),
        "average_distance": float(distances.mean()),
        "max_distance": float(distances.max()),
        "runtime": runtime,
        "distance_output_path": distance_output_path,
        "colored_output_path": colored_output_path,
        "stats_output_path": stats_output_path,
    }


def main():
    tasks = discover_mesh_tasks(MESH_DIR, POINT_CLOUD_DIR)
    summary_rows = []

    for task_id, mesh_pairs in tasks.items():
        task_output_dir = RESULT_DIR / task_id
        task_point_count = 0
        task_distance_sum = 0.0
        task_max_distance = 0.0
        task_runtime = 0.0

        for mesh_path, point_cloud_path in mesh_pairs:
            mesh_result = compute_one_mesh(mesh_path, point_cloud_path, task_output_dir)
            task_point_count += mesh_result["point_count"]
            task_distance_sum += mesh_result["distance_sum"]
            task_max_distance = max(task_max_distance, mesh_result["max_distance"])
            task_runtime += mesh_result["runtime"]

            del mesh_result
            gc.collect()
            if DEVICE.type == "cuda":
                torch.cuda.empty_cache()

        summary_rows.append(
            {
                "task_id": task_id,
                "mesh_count": len(mesh_pairs),
                "point_count": task_point_count,
                "average_distance": task_distance_sum / task_point_count,
                "max_distance": task_max_distance,
                "runtime": task_runtime,
            }
        )

    summary_rows.sort(key=lambda row: int(row["task_id"]))
    write_summary_xlsx(SUMMARY_XLSX_PATH, summary_rows)
    write_summary_csv(SUMMARY_CSV_PATH, summary_rows)

    print("Task\tAverage distance\tMax distance\tRuntime")
    for row in summary_rows:
        print(
            f"{row['task_id']}\t"
            f"{row['average_distance']:.8f}\t"
            f"{row['max_distance']:.8f}\t"
            f"{row['runtime']:.4f}"
        )


if __name__ == "__main__":
    main()