import os
import sys

import mitsuba as mi
import numpy as np
import open3d as o3d

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

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

def load_ply(path, MAX_PTS=4096):
    pcd = o3d.io.read_point_cloud(path)
    xyz = np.asarray(pcd.points)
    rgb = np.asarray(pcd.colors) if pcd.has_colors() else np.full_like(xyz, 0.0)
    if len(xyz) > MAX_PTS:
        idx = np.random.choice(len(xyz), MAX_PTS, replace=True)
        xyz, rgb = xyz[idx], rgb[idx]
    return xyz, rgb



if __name__ == '__main__':
    mi.set_variant("scalar_rgb")

    INPUT_DIR = "vis_results/pcd"
    OUTPUT_DIR = "vis_results/rd_img"
    SPHERE_RADIUS = 0.005
    MAX_POINTS = 25000
    SAVE_XML = False

    os.makedirs(OUTPUT_DIR, exist_ok=True)

    for file_name in os.listdir(INPUT_DIR):
        if not file_name.endswith('.ply'):
            continue

        print(f"Rendering {file_name}...")
        file_path = os.path.join(INPUT_DIR, file_name)
        output_img_path = os.path.join(OUTPUT_DIR, file_name.replace('.ply', '.png'))

        points, colors = load_ply(file_path, MAX_POINTS)
        points = center_normalize(points)
        points[:, 2] += SPHERE_RADIUS / 2

        scene = MiScene()
        scene.add("sensor", MiSensor(
            origin=[2, 2, 2],
            target=[0, 0, 0],
            fov=25,
            sample_count=256,
            film_width=2000,
            film_height=2000,
        ))
        scene.add("floor", MiFloor(width=10, height=10, color=[1, 1, 1]))
        scene.add("soft_light", MiSoftlight(origin=[-4, 4, 20], target=[0, 0, 0], intensity=6))

        for i, (pt, clr) in enumerate(zip(points, colors)):
            scene.add(f"sphere{i}", MiSphere(pt.tolist(), SPHERE_RADIUS, clr.tolist()))

        if SAVE_XML:
            xml_path = os.path.join(OUTPUT_DIR, file_name.replace('.ply', '.xml'))
            mi.xml.dict_to_xml(scene.dict(), xml_path)
        else:
            scene_mi = mi.load_dict(scene.dict())
            image = mi.render(scene_mi, spp=256)
            mi.util.write_bitmap(output_img_path, image)

    print("Rendering completed.")