#!/usr/bin/env python3
"""
S3DIS 超点 / 语义 GT 可视化（Mitsuba 3）。

模式 superpoint_per_class：对 13 个语义类分别只画该类的点，按超点 id 上色，各输出一张 PNG。

用法见 --help。
"""
from __future__ import annotations

import argparse
import math
import os
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)


def _pick_variant():
    """scalar_rgb 与 mi_objects 中 Transform 序列化兼容。"""
    import mitsuba as mi

    for v in ("scalar_rgb", "llvm_ad_rgb", "cuda_ad_rgb"):
        try:
            mi.set_variant(v)
            return v
        except Exception:
            continue
    raise RuntimeError("无法设置任何 Mitsuba variant")


_GOLDEN = 0.618033988749895


def _hsv_to_rgb(h: float, s: float, v: float):
    import colorsys

    r, g, b = colorsys.hsv_to_rgb(h % 1.0, s, v)
    return r, g, b


def _hash_color(uid: int):
    h = (int(uid) * _GOLDEN) % 1.0
    return _hsv_to_rgb(h, 0.82, 0.96)


S3DIS_SEGMENT_RGB = [
    [0.90, 0.90, 0.92],
    [0.72, 0.52, 0.30],
    [0.42, 0.72, 0.48],
    [0.35, 0.55, 0.85],
    [0.95, 0.75, 0.25],
    [0.45, 0.85, 0.95],
    [0.92, 0.35, 0.35],
    [0.95, 0.55, 0.20],
    [0.75, 0.45, 0.92],
    [0.55, 0.40, 0.88],
    [0.50, 0.82, 0.45],
    [0.88, 0.50, 0.65],
    [0.55, 0.55, 0.58],
]

S3DIS_CLASS_NAMES = [
    "ceiling",
    "floor",
    "wall",
    "beam",
    "column",
    "window",
    "door",
    "table",
    "chair",
    "sofa",
    "bookcase",
    "board",
    "clutter",
]


def _class_color(cid: int, n_cls: int = 13):
    i = int(cid) % max(n_cls, 1)
    return tuple(S3DIS_SEGMENT_RGB[i])


def _maybe_crop_png(
    path: str, crop_wh: tuple[int, int] | None, center: bool = False
) -> str | None:
    if not crop_wh:
        return None
    try:
        import time
        from PIL import Image
    except ImportError:
        print("未安装 Pillow，跳过裁剪", file=sys.stderr)
        return None
    w, h = crop_wh
    im = None
    for _ in range(40):
        try:
            im = Image.open(path)
            im.load()
            break
        except (OSError, Exception):
            time.sleep(0.05)
    if im is None:
        print("裁剪失败：无法读取刚写入的 PNG", file=sys.stderr)
        return None
    im = im.convert("RGB")
    cw, ch = min(w, im.size[0]), min(h, im.size[1])
    if center:
        left = max(0, (im.size[0] - cw) // 2)
        top = max(0, (im.size[1] - ch) // 2)
        im = im.crop((left, top, left + cw, top + ch))
    else:
        im = im.crop((0, 0, cw, ch))
    base, ext = os.path.splitext(path)
    out = base + "_crop" + ext
    im.save(out)
    return out


def _render_one(
    mi,
    points: "np.ndarray",
    label_for_color: "np.ndarray",
    color_fn,
    out_png: str,
    sphere_radius: float,
    film_size: int,
    spp: int,
    fov: float,
    use_plastic: bool,
    light_intensity: float,
    fill_irradiance: float,
):
    """points: (N,3) 已归一化前原始坐标；label_for_color 与 points 等长。"""
    import numpy as np

    from lib_render import center_normalize
    from mi_objects import MiFloor, MiScene, MiSensor, MiSoftlight, MiSphere

    n = points.shape[0]
    if n == 0:
        return False

    pts = center_normalize(points.astype(np.float64))
    pts[:, 2] += sphere_radius / 2.0

    lds_spp = max(16, 2 ** int(round(math.log(max(int(spp), 16), 2))))
    lds_spp = min(lds_spp, 1024)

    scene = MiScene()
    scene.add(
        "sensor",
        MiSensor(
            origin=[2, 2, 2],
            target=[0, 0, 0],
            fov=int(round(fov)),
            sample_count=lds_spp,
            film_width=int(film_size),
            film_height=int(film_size),
        ),
    )
    scene.add("floor", MiFloor(width=10, height=10, color=[0.96, 0.96, 0.98]))
    scene.add(
        "soft_light",
        MiSoftlight(
            origin=[-4, 4, 20],
            target=[0, 0, 0],
            intensity=float(light_intensity),
        ),
    )

    for i in range(n):
        c = color_fn(int(label_for_color[i]))
        clr = [float(c[0]), float(c[1]), float(c[2])]
        scene.add(
            f"sphere{i}",
            MiSphere(pts[i].tolist(), float(sphere_radius), clr, plastic=use_plastic),
        )

    scene_dict = scene.dict()
    fi = float(fill_irradiance)
    scene_dict["fill_sun"] = {
        "type": "directional",
        "direction": [0.35, -0.55, 0.75],
        "irradiance": {"type": "rgb", "value": [fi, fi * 0.98, fi * 0.95]},
    }

    scene_mi = mi.load_dict(scene_dict)
    img = mi.render(scene_mi, spp=int(spp))
    out = os.path.abspath(out_png)
    os.makedirs(os.path.dirname(out) or ".", exist_ok=True)
    mi.util.write_bitmap(out, img)
    print("MITSUBA_PNG", out)
    return True


def main():
    _pick_variant()
    import numpy as np
    import mitsuba as mi

    default_root = os.path.normpath(
        os.path.join(_HERE, "..", "_work_biptv3", "pointcept_framework", "data", "s3dis_official")
    )

    ap = argparse.ArgumentParser()
    ap.add_argument("--data_root", type=str, default=default_root)
    ap.add_argument("--room", type=str, required=True)
    ap.add_argument(
        "--out_png",
        type=str,
        default=None,
        help="单张输出路径（superpoint / segment 模式必填）",
    )
    ap.add_argument(
        "--out_dir",
        type=str,
        default=None,
        help="superpoint_per_class 时输出目录，默认 outputs/superpoint_vis/<room>_per_class_sp/",
    )
    ap.add_argument(
        "--mode",
        choices=("superpoint", "segment", "superpoint_per_class"),
        default="superpoint",
    )
    ap.add_argument("--max_points", type=int, default=40000)
    ap.add_argument("--sphere_radius", type=float, default=0.008)
    ap.add_argument("--film_size", type=int, default=2560)
    ap.add_argument("--spp", type=int, default=256)
    ap.add_argument("--fov", type=float, default=25.0)
    ap.add_argument("--crop", type=str, default=None)
    ap.add_argument("--crop_center", action="store_true")
    ap.add_argument("--no_plastic_spheres", action="store_true")
    ap.add_argument("--light_intensity", type=float, default=14.0)
    ap.add_argument("--fill_irradiance", type=float, default=3.5)
    ap.add_argument(
        "--min_class_points",
        type=int,
        default=50,
        help="superpoint_per_class：点数少于此的类跳过",
    )
    args = ap.parse_args()

    room = os.path.join(args.data_root, args.room)
    coord = np.load(os.path.join(room, "coord.npy"))
    n_all = coord.shape[0]

    use_plastic = not args.no_plastic_spheres
    room_tag = args.room.replace("/", "_")

    if args.mode == "superpoint_per_class":
        segment = np.load(os.path.join(room, "segment.npy")).reshape(-1).astype(np.int64)
        superpoint = np.load(os.path.join(room, "superpoint.npy")).reshape(-1).astype(np.int64)
        if len(segment) != n_all or len(superpoint) != n_all:
            raise ValueError("segment/superpoint 与 coord 长度不一致")

        out_dir = args.out_dir
        if not out_dir:
            out_dir = os.path.join(_HERE, "outputs", "superpoint_vis", f"{room_tag}_per_class_sp")
        out_dir = os.path.abspath(out_dir)
        os.makedirs(out_dir, exist_ok=True)

        written = []
        for cls in range(13):
            mask = segment == cls
            cnt = int(mask.sum())
            if cnt < args.min_class_points:
                print(f"skip cls{cls:02d} {S3DIS_CLASS_NAMES[cls]}: n={cnt} (<{args.min_class_points})")
                continue
            coord_c = coord[mask]
            sp_c = superpoint[mask]
            if cnt > args.max_points:
                rng = np.random.default_rng(cls)
                idx = rng.choice(cnt, size=args.max_points, replace=False)
                coord_c = coord_c[idx]
                sp_c = sp_c[idx]

            name = S3DIS_CLASS_NAMES[cls]
            out_png = os.path.join(out_dir, f"{room_tag}_cls{cls:02d}_{name}_superpoint.png")
            ok = _render_one(
                mi,
                coord_c,
                sp_c,
                _hash_color,
                out_png,
                args.sphere_radius,
                args.film_size,
                args.spp,
                args.fov,
                use_plastic,
                args.light_intensity,
                args.fill_irradiance,
            )
            if ok:
                written.append(out_png)
                if args.crop:
                    wh = args.crop.lower().replace(" ", "").split("x")
                    if len(wh) == 2:
                        cw, ch = int(wh[0]), int(wh[1])
                        cpath = _maybe_crop_png(out_png, (cw, ch), center=args.crop_center)
                        if cpath:
                            print("MITSUBA_PNG_CROP", cpath)

        print("PER_CLASS_DONE", len(written), "files in", out_dir)
        return

    if not args.out_png:
        sys.exit("非 superpoint_per_class 模式必须指定 --out_png")

    if args.mode == "superpoint":
        labels = np.load(os.path.join(room, "superpoint.npy")).reshape(-1).astype(np.int64)
        color_fn = _hash_color
    else:
        labels = np.load(os.path.join(room, "segment.npy")).reshape(-1).astype(np.int64)
        color_fn = _class_color

    if len(labels) != n_all:
        raise ValueError("标签与 coord 长度不一致")

    coord_u = coord
    labels_u = labels
    if n_all > args.max_points:
        rng = np.random.default_rng(0)
        idx = rng.choice(n_all, size=args.max_points, replace=False)
        coord_u = coord[idx]
        labels_u = labels[idx]

    _render_one(
        mi,
        coord_u,
        labels_u,
        color_fn,
        args.out_png,
        args.sphere_radius,
        args.film_size,
        args.spp,
        args.fov,
        use_plastic,
        args.light_intensity,
        args.fill_irradiance,
    )

    if args.crop:
        wh = args.crop.lower().replace(" ", "").split("x")
        if len(wh) == 2:
            cw, ch = int(wh[0]), int(wh[1])
            cropped = _maybe_crop_png(args.out_png, (cw, ch), center=args.crop_center)
            if cropped:
                print("MITSUBA_PNG_CROP", cropped)


if __name__ == "__main__":
    main()