#!/usr/bin/env python3
"""
Superpoint Generation Core — S3DIS
===================================

  knn_graph(points, k) -> segment_point(points, normals, edges, kThresh, segMinVerts)

  超点是点云中的局部过分割单元，类比图像中的 superpixel。
  先把空间上相邻、法向量相近的点分成紧凑的小块，
  再供下游语义网络在超点粒度上做学习与推理。

算法：
  1. 对输入点云构建 k-近邻图 (knn_graph, k=50)
  2. 对每条边计算权重：w = 1 - dot(n1, n2)；凸区域 w = w^2
  3. 按权重升序排列，执行 Felzenszwalb 图分割：
     - 初始每个点为独立分量，阈值 = kThresh
     - 若边权 <= 两端分量阈值，合并，更新阈值 = w + kThresh / size
  4. 合并过小分量 (< segMinVerts)

"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np

DEFAULT_SEGMENTATOR_CPP = Path('segmentator.cpp')
DEFAULT_BUILD_DIR = Path('torch_build')

_cached_seg_module = None


def load_segmentator(cpp_path: Path = DEFAULT_SEGMENTATOR_CPP,
                     build_dir: Path = DEFAULT_BUILD_DIR,
                     verbose: bool = False):
    
    global _cached_seg_module
    if _cached_seg_module is not None:
        return _cached_seg_module

    from torch.utils.cpp_extension import load

    cpp_path = Path(cpp_path)
    build_dir = Path(build_dir)
    if not cpp_path.is_file():
        raise FileNotFoundError(f'segmentator.cpp not found: {cpp_path}')
    build_dir.mkdir(parents=True, exist_ok=True)

    _cached_seg_module = load(
        name='libsegmentator_dyn',
        sources=[str(cpp_path)],
        build_directory=str(build_dir),
        extra_cflags=['-O3'],
        verbose=verbose,
    )
    return _cached_seg_module


def generate_superpoints(
    coords: np.ndarray,
    normals: np.ndarray,
    knn_k: int = 50,
    k_thresh: float = 0.01,
    seg_min_verts: int = 20,
    cpp_path: Path = DEFAULT_SEGMENTATOR_CPP,
    build_dir: Path = DEFAULT_BUILD_DIR,
) -> np.ndarray:
    """基于 knn_graph + Felzenszwalb 图分割生成超点标签。

    Parameters
    ----------
    coords : (N, 3) float32 — 点坐标
    normals : (N, 3) float32 — 点法向量
    knn_k : int — k-近邻数，控制图的连通密度
    k_thresh : float — Felzenszwalb 分割阈值，越小超点越细
    seg_min_verts : int — 最小超点点数，低于此值的分量会被合并

    Returns
    -------
    labels : (N,) int32 — 每个点的超点标签，从 0 开始连续编号
    """
    import torch
    from torch_cluster import knn_graph

    coords = np.ascontiguousarray(np.asarray(coords, dtype=np.float32))
    normals = np.ascontiguousarray(np.asarray(normals, dtype=np.float32))
    if coords.shape != normals.shape:
        raise ValueError(f'coords/normals shape mismatch: {coords.shape} vs {normals.shape}')

    # 归一化法向量
    normals = normals / (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-8)

    seg = load_segmentator(cpp_path, build_dir)

    pts = torch.from_numpy(coords)
    nrm = torch.from_numpy(normals)

    # 构建 k-近邻图：knn_graph 返回 [2, N*k]，转置为 [N*k, 2]
    edges = knn_graph(pts, k=knn_k).T.contiguous().to(dtype=torch.int64, device='cpu')

    # Felzenszwalb 图分割
    labels = seg.segment_point(
        pts.contiguous(), nrm.contiguous(), edges,
        float(k_thresh), int(seg_min_verts),
    )
    labels = labels.cpu().numpy().reshape(-1).astype(np.int64)

    # 重编号为连续整数
    _, inverse = np.unique(labels, return_inverse=True)
    return inverse.astype(np.int32)


def iter_rooms(data_root: Path):
    for area in sorted(data_root.glob('Area_*')):
        if not area.is_dir():
            continue
        for room in sorted(area.iterdir()):
            if room.is_dir() and (room / 'coord.npy').is_file():
                yield room


def process_room(room_dir: Path, args) -> dict:
    """对单个房间生成超点标签"""
    coord = np.load(room_dir / 'coord.npy').astype(np.float32)
    normal_path = room_dir / 'normal.npy'
    if not normal_path.is_file():
        raise FileNotFoundError(f'normal.npy required: {room_dir}')
    normal = np.load(normal_path).astype(np.float32)

    t0 = time.time()
    labels = generate_superpoints(
        coord, normal,
        knn_k=args.knn_k,
        k_thresh=args.k_thresh,
        seg_min_verts=args.seg_min_verts,
        cpp_path=Path(args.segmentator_cpp),
        build_dir=Path(args.build_dir),
    )
    elapsed = time.time() - t0

    uniq, cnt = np.unique(labels, return_counts=True)
    rel_path = room_dir.relative_to(args.data_root)

    info = {
        'room': str(rel_path),
        'points': int(coord.shape[0]),
        'num_superpoints': int(uniq.size),
        'mean_pts_per_sp': round(float(cnt.mean()), 1),
        'max_pts_per_sp': int(cnt.max()),
        'min_pts_per_sp': int(cnt.min()),
        'time_sec': round(elapsed, 1),
        'params': {
            'knn_k': args.knn_k,
            'k_thresh': args.k_thresh,
            'seg_min_verts': args.seg_min_verts,
        },
    }


    disk_sp = room_dir / 'superpoint.npy'
    if disk_sp.is_file():
        old = np.load(disk_sp).reshape(-1)
        if len(old) == len(labels):
            old_uniq = np.unique(old)
            info['existing_num_superpoints'] = int(old_uniq.size)

    
    if args.write:
        if args.output_root:
            out_path = Path(args.output_root) / rel_path / 'superpoint.npy'
        else:
            out_path = room_dir / 'superpoint.npy'
        out_path.parent.mkdir(parents=True, exist_ok=True)
        np.save(str(out_path), labels)
        info['output'] = str(out_path)

    return info


def main():
    ap = argparse.ArgumentParser(description='S3DIS Superpoint Generation (segmentator / Felzenszwalb)')
    ap.add_argument('--data_root', type=Path, default=Path('/mnt/data/AODUOLI/_work_biptv3/pointcept_framework/data/s3dis_official'))
    ap.add_argument('--room', type=str, default=None, help=' Area_1/office_1')
    ap.add_argument('--all', action='store_true', help='ALL')
    ap.add_argument('--knn_k', type=int, default=50)
    ap.add_argument('--k_thresh', type=float, default=0.01)
    ap.add_argument('--seg_min_verts', type=int, default=20)
    ap.add_argument('--segmentator_cpp', type=str, default=str(DEFAULT_SEGMENTATOR_CPP))
    ap.add_argument('--build_dir', type=str, default=str(DEFAULT_BUILD_DIR))
    ap.add_argument('--write', action='store_true', help='写出 superpoint.npy')
    ap.add_argument('--output_root', type=str, default=None, help='output')
    args = ap.parse_args()

    if not args.data_root.is_dir():
        print(f'data_root not found: {args.data_root}', file=sys.stderr)
        sys.exit(1)

    rooms = []
    if args.room:
        rooms = [args.data_root / args.room]
    elif args.all:
        rooms = list(iter_rooms(args.data_root))
    else:
        print('Please specify --room <path> or --all', file=sys.stderr)
        sys.exit(1)

    print(f'Rooms: {len(rooms)}, knn_k={args.knn_k}, k_thresh={args.k_thresh}, seg_min_verts={args.seg_min_verts}')
    if not args.write:
        print('[DRY-RUN] Add --write to save files.\n')

    for room_dir in rooms:
        if not (room_dir / 'coord.npy').is_file():
            print(f'SKIP: {room_dir}')
            continue
        info = process_room(room_dir, args)
        print(json.dumps(info, ensure_ascii=False))


if __name__ == '__main__':
    main()