"""Re-render all existing superpoint results with improved circular splatting."""
import os
import sys
import numpy as np
import open3d as o3d

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, SCRIPT_DIR)

from run_pycut_s3dis import render_perspective_cutaway, OUT_DIR

S3DIS_ROOT = "/mnt/data/AODUOLI/_work_biptv3/pointcept_framework/data/s3dis_official"
SCANNET_PLY_DIR = os.path.join(SCRIPT_DIR, "paper_assets_grouped_by_dataset", "scannet", "ply")


def rerender_s3dis():
    """Re-render S3DIS rooms using saved superpoint labels."""
    rooms = [
        ("Area_1", "office_1"),
        ("Area_1", "conferenceRoom_1"),
    ]
    for area, room in rooms:
        room_dir = os.path.join(S3DIS_ROOT, area, room)
        coord = np.load(os.path.join(room_dir, "coord.npy")).astype(np.float32)

        label_path = os.path.join(OUT_DIR, f"{area}_{room}_superpoint.npy")
        if not os.path.exists(label_path):
            print(f"  SKIP {area}/{room}: no saved labels")
            continue
        labels = np.load(label_path)
        n_sp = int(labels.max()) + 1
        print(f"\n{'='*60}")
        print(f"Re-rendering {area}/{room}  ({coord.shape[0]} pts, {n_sp} sp)")
        print(f"{'='*60}")

        prefix = f"{area}_{room}_pycut"
        for azim, tag in [(225, "a"), (135, "b"), (315, "c"), (45, "d")]:
            render_perspective_cutaway(
                coord, labels,
                os.path.join(OUT_DIR, f"{prefix}_paper_{tag}.png"),
                azim_deg=azim,
            )


def rerender_scannet():
    """Re-render ScanNet scenes using saved superpoint labels."""
    scenes = [
        ("scannet_scene0050_00", "scannet_scene0050_00_ours.ply",
         "scannet_scene0050_00_spfinal.npy"),
        ("scannet_scene0217_00", "scannet_scene0217_00_ours.ply",
         "scannet_scene0217_00_spfinal.npy"),
        ("scannet_scene0568_01", "scannet_scene0568_01_ours.ply",
         "scannet_scene0568_01_spfinal.npy"),
    ]

    for scene_name, ply_file, sp_file in scenes:
        ply_path = os.path.join(SCANNET_PLY_DIR, ply_file)
        sp_path = os.path.join(OUT_DIR, sp_file)
        if not os.path.exists(sp_path):
            print(f"  SKIP {scene_name}: no saved labels at {sp_path}")
            continue

        pcd = o3d.io.read_point_cloud(ply_path)
        xyz = np.asarray(pcd.points, dtype=np.float32)
        labels = np.load(sp_path)
        n_sp = int(labels.max()) + 1
        print(f"\n{'='*60}")
        print(f"Re-rendering {scene_name}  ({xyz.shape[0]} pts, {n_sp} sp)")
        print(f"{'='*60}")

        for azim, tag in [(225, "a"), (135, "b")]:
            render_perspective_cutaway(
                xyz, labels,
                os.path.join(OUT_DIR, f"{scene_name}_v3_{tag}.png"),
                azim_deg=azim,
            )


if __name__ == "__main__":
    rerender_s3dis()
    rerender_scannet()
    print("\nDone!")