#!/usr/bin/env python3
"""
S3DIS 超点可视化（编辑节点 / poplab）

数据侧 superpoint 来源（与预处理一致）：
  - 每个房间目录下的 superpoint.npy：与 coord.npy 等长的 int32，表示每个点所属超点 id。
  - 生成逻辑（无外部分割器时）：体素坐标 +（可选）法向粗分箱 -> np.unique 得到逆映射作为 id。
  - 见 pointcept/datasets/preprocessing/s3dis/preprocess_s3dis.py 中 generate_superpoint_labels。

用法示例：
  python scripts/visualize_s3dis_superpoints.py \\
    --room Area_1/office_1 \\
    --out /tmp/superpoints_vis.ply

  python scripts/visualize_s3dis_superpoints.py --room Area_1/office_1 --recompute --voxel 0.12 --normal-bins 8
"""

from __future__ import annotations

import argparse
import os
from pathlib import Path

import numpy as np


def generate_superpoint_labels(coords, normals=None, voxel_size=0.12, normal_bins=8):
    """与 preprocess_s3dis.generate_superpoint_labels 一致的几何 fallback。"""
    coords = np.asarray(coords, dtype=np.float32)
    coord_min = coords.min(axis=0, keepdims=True)
    voxel_coord = np.floor((coords - coord_min) / max(float(voxel_size), 1e-4)).astype(np.int64)

    if normals is not None and len(normals) == len(coords):
        normals = np.asarray(normals, dtype=np.float32)
        normals = normals / (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-8)
        normal_q = np.floor((normals + 1.0) * 0.5 * normal_bins).astype(np.int64)
        normal_q = np.clip(normal_q, 0, normal_bins)
        tokens = np.concatenate([voxel_coord, normal_q], axis=1)
    else:
        tokens = voxel_coord

    _, inverse = np.unique(tokens, axis=0, return_inverse=True)
    return inverse.astype(np.int32)


def sp_ids_to_colors(sp: np.ndarray) -> np.ndarray:
    """稳定伪彩色：对每个 superpoint id 做哈希映射到 [0,1]^3。"""
    sp = sp.astype(np.int64).reshape(-1)
    n = len(sp)
    colors = np.zeros((n, 3), dtype=np.float64)
    for uid in np.unique(sp):
        h = (int(uid) * 1103515245 + 12345) & 0x7FFFFFFF
        r = ((h >> 0) & 255) / 255.0
        g = ((h >> 8) & 255) / 255.0
        b = ((h >> 16) & 255) / 255.0
        colors[sp == uid] = (r, g, b)
    return np.clip(colors, 0.0, 1.0)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--data_root",
        type=Path,
        default=Path(__file__).resolve().parents[1] / "data" / "s3dis_official",
        help="S3DIS 处理根目录（含 Area_*/*/）",
    )
    ap.add_argument("--room", type=str, required=True, help="例如 Area_1/office_1")
    ap.add_argument("--out", type=Path, default=Path("superpoints_vis.ply"))
    ap.add_argument(
        "--recompute",
        action="store_true",
        help="不读 superpoint.npy，按 generate_superpoint_labels 现场重算（用于对照）",
    )
    ap.add_argument("--voxel", type=float, default=0.12)
    ap.add_argument("--normal-bins", type=int, default=8)
    args = ap.parse_args()

    room_dir = args.data_root / args.room
    coord_p = room_dir / "coord.npy"
    if not coord_p.is_file():
        raise FileNotFoundError(coord_p)

    coord = np.load(coord_p)
    normal = None
    npy_n = room_dir / "normal.npy"
    if npy_n.is_file():
        normal = np.load(npy_n)

    if args.recompute:
        sp = generate_superpoint_labels(
            coord, normals=normal, voxel_size=args.voxel, normal_bins=args.normal_bins
        )
        tag = "recomputed"
    else:
        sp_p = room_dir / "superpoint.npy"
        if not sp_p.is_file():
            raise FileNotFoundError(
                f"{sp_p} 不存在，可先跑预处理加 --generate_superpoint，或改用 --recompute"
            )
        sp = np.load(sp_p).reshape(-1).astype(np.int32)
        tag = "disk"

    if len(sp) != len(coord):
        raise ValueError(f"superpoint 长度 {len(sp)} != coord 长度 {len(coord)}")

    colors = sp_ids_to_colors(sp)
    try:
        import open3d as o3d

        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(coord.astype(np.float64))
        pcd.colors = o3d.utility.Vector3dVector(colors)
        args.out.parent.mkdir(parents=True, exist_ok=True)
        o3d.io.write_point_cloud(str(args.out), pcd)
        print(f"[{tag}] wrote {args.out}  points={len(coord)}  unique_superpoints={len(np.unique(sp))}")
        print("若在本机有显示器，可取消注释 o3d.visualization.draw_geometries([pcd])")
        # o3d.visualization.draw_geometries([pcd])
    except ImportError:
        print("未安装 open3d，无法写 PLY；请: pip install open3d")


if __name__ == "__main__":
    main()