#!/usr/bin/env python3
"""
无头环境生成超点 / 语义 GT 点云预览 PNG（matplotlib，非 Blender）。

若需要 Blender（bpy）真渲染，请用：
  PAMI2026/run_blender_superpoints.sh
  或直接 blender --background --python PAMI2026/blender_visualize_superpoints.py -- ...

输出目录默认：PAMI2026/outputs/superpoint_vis/

用法：
  /mnt/data/AODUOLI/miniconda_envs/Aoduo/bin/python \\
    /mnt/data/AODUOLI/PAMI2026/render_superpoints_image.py \\
    --room Area_1/office_1
"""
from __future__ import annotations

import argparse
import math
import os
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


def _hash_color(uid: int):
    h = (int(uid) * 1103515245 + 12345) & 0x7FFFFFFF
    r = ((h >> 0) & 255) / 255.0
    g = ((h >> 8) & 255) / 255.0
    b = ((h >> 16) & 255) / 255.0
    return r, g, b


def _class_color(cid: int, n_cls: int = 13):
    t = (int(cid) % max(n_cls, 1)) / max(n_cls, 1)
    return (
        0.5 + 0.5 * math.cos(2 * math.pi * t),
        0.5 + 0.5 * math.cos(2 * math.pi * t + 2.09),
        0.5 + 0.5 * math.cos(2 * math.pi * t + 4.18),
    )


def _labels_to_rgb(labels: np.ndarray, mode: str) -> np.ndarray:
    labels = labels.reshape(-1).astype(np.int64)
    n = labels.shape[0]
    rgb = np.zeros((n, 3), dtype=np.float64)
    for lid in np.unique(labels):
        if mode == "superpoint":
            c = _hash_color(int(lid))
        else:
            c = _class_color(int(lid))
        m = labels == lid
        rgb[m] = c
    return rgb


def _subsample(coord: np.ndarray, labels: np.ndarray, max_points: int, seed: int):
    n = coord.shape[0]
    if n <= max_points:
        return coord, labels, np.arange(n)
    rng = np.random.default_rng(seed)
    idx = rng.choice(n, size=max_points, replace=False)
    return coord[idx], labels[idx], idx


def render_png(
    coord: np.ndarray,
    rgb: np.ndarray,
    out_path: Path,
    title: str,
    elev: float = 20.0,
    azim: float = -60.0,
    point_size: float = 0.15,
):
    fig = plt.figure(figsize=(12, 10), dpi=150)
    ax = fig.add_subplot(111, projection="3d")
    ax.scatter(
        coord[:, 0],
        coord[:, 1],
        coord[:, 2],
        c=rgb,
        s=point_size,
        linewidths=0,
        alpha=0.85,
        depthshade=False,
    )
    ax.set_title(title, fontsize=12)
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")
    try:
        ax.set_box_aspect(
            (
                float(coord[:, 0].ptp()),
                float(coord[:, 1].ptp()),
                float(coord[:, 2].ptp()),
            )
        )
    except Exception:
        pass
    ax.view_init(elev=elev, azim=azim)
    ax.set_axis_off()
    plt.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, bbox_inches="tight", facecolor="white")
    plt.close(fig)


def main():
    root = Path(__file__).resolve().parent
    default_data = (
        root.parent
        / "_work_biptv3"
        / "pointcept_framework"
        / "data"
        / "s3dis_official"
    )
    default_out = root / "outputs" / "superpoint_vis"

    ap = argparse.ArgumentParser()
    ap.add_argument("--data_root", type=Path, default=default_data)
    ap.add_argument("--room", type=str, default="Area_1/office_1")
    ap.add_argument("--out_dir", type=Path, default=default_out)
    ap.add_argument("--max_points", type=int, default=120000)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--elev", type=float, default=22.0)
    ap.add_argument("--azim", type=float, default=-58.0)
    args = ap.parse_args()

    room_dir = args.data_root / args.room
    coord = np.load(room_dir / "coord.npy")
    sp = np.load(room_dir / "superpoint.npy").reshape(-1)
    seg = np.load(room_dir / "segment.npy").reshape(-1)
    assert len(sp) == len(coord) == len(seg)

    room_tag = args.room.replace("/", "_")
    out_dir = args.out_dir.resolve()

    # 超点
    c1, l1, _ = _subsample(coord, sp, args.max_points, args.seed)
    rgb_sp = _labels_to_rgb(l1, "superpoint")
    path_sp = out_dir / f"{room_tag}_superpoint.png"
    render_png(
        c1,
        rgb_sp,
        path_sp,
        title="Superpoint ids (geometry, subsampled)",
        elev=args.elev,
        azim=args.azim,
    )

    # 语义 GT
    c2, l2, _ = _subsample(coord, seg, args.max_points, args.seed + 1)
    rgb_seg = _labels_to_rgb(l2, "segment")
    path_seg = out_dir / f"{room_tag}_segment_gt.png"
    render_png(
        c2,
        rgb_seg,
        path_seg,
        title="Semantic GT (segment.npy, subsampled)",
        elev=args.elev,
        azim=args.azim,
    )

    print("WROTE_SUPERPOINT", str(path_sp))
    print("WROTE_SEGMENT", str(path_seg))


if __name__ == "__main__":
    main()