"""Regenerate ScanNet superpoints with tuned per-scene lambdas + circular splatting render."""
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 (
    generate_superpoints_pycut, render_perspective_cutaway,
    create_labeled_point_cloud, OUT_DIR,
)

SCANNET_PLY_DIR = os.path.join(SCRIPT_DIR, "paper_assets_grouped_by_dataset", "scannet", "ply")

# Per-scene tuned parameters
SCENES = [
    {
        "name": "scannet_scene0050_00",
        "ply": "scannet_scene0050_00_ours.ply",
        "lam": 1.0, "merge_min_size": 150,
    },
    {
        "name": "scannet_scene0217_00",
        "ply": "scannet_scene0217_00_ours.ply",
        "lam": 10.0, "merge_min_size": 200,
    },
    {
        "name": "scannet_scene0568_01",
        "ply": "scannet_scene0568_01_ours.ply",
        "lam": 0.5, "merge_min_size": 100,
    },
]

for sc in SCENES:
    ply_path = os.path.join(SCANNET_PLY_DIR, sc["ply"])
    print(f"\n{'='*60}")
    print(f"Processing {sc['name']}  (lam={sc['lam']})")
    print(f"{'='*60}")

    pcd = o3d.io.read_point_cloud(ply_path)
    xyz = np.asarray(pcd.points, dtype=np.float32)

    # Estimate normals (ScanNet PLY has no normals)
    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=30))
    normals = np.asarray(pcd.normals, dtype=np.float32)

    print(f"  points={xyz.shape[0]}")

    labels = generate_superpoints_pycut(
        xyz, normals=normals,
        k_feat=10, k_adj=10, chunk_size=8192,
        use_input_normals=True, use_xyz=False,
        xyz_scale=0.10, normal_scale=0.25,
        lam=sc["lam"], sigma=0.5,
        mutual=False, undirected=True,
        min_comp_weight=20, weight_decay=0.7,
        merge_min_size=sc["merge_min_size"],
        verbose=True,
    )

    n_sp = int(labels.max()) + 1
    print(f"  Final: {n_sp} superpoints")

    # Save labels
    sp_path = os.path.join(OUT_DIR, f"{sc['name']}_sp_tuned.npy")
    np.save(sp_path, labels)

    # Save PLY
    create_labeled_point_cloud(xyz, labels, f"{sc['name']}_tuned", normals=normals)

    # Render with circular splatting
    for azim, tag in [(225, "a"), (135, "b")]:
        render_perspective_cutaway(
            xyz, labels,
            os.path.join(OUT_DIR, f"{sc['name']}_v3_{tag}.png"),
            azim_deg=azim,
        )

print("\nDone!")