"""Offscreen renderer for point clouds and meshes (Open3D EGL, headless).

Produces clean multi-view sheets so reconstruction results can be judged
visually. Framing uses robust (percentile) bounds so a handful of stray
outlier points cannot shrink the subject to a speck.
"""
import os, argparse
import numpy as np
import open3d as o3d
from open3d.visualization import rendering
from PIL import Image, ImageDraw

W = H = 560
VIEWS = [("front", 0, 0), ("right", 90, 0), ("back", 180, 0),
         ("left", 270, 0), ("top", 0, 88), ("iso", 40, 25)]
BG = (0.12, 0.13, 0.15, 1.0)


def robust_bounds(pts, lo=0.5, hi=99.5):
    mn = np.percentile(pts, lo, axis=0)
    mx = np.percentile(pts, hi, axis=0)
    return mn, mx


def geom_points(geom):
    if isinstance(geom, o3d.geometry.PointCloud):
        return np.asarray(geom.points)
    return np.asarray(geom.vertices)


def render_geometry(geom, out_prefix, point_size=2.5, fill=0.85, label=None):
    pts = geom_points(geom)
    mn, mx = robust_bounds(pts)
    center = (mn + mx) / 2.0
    radius = float(np.linalg.norm(mx - mn)) / 2.0
    dist = radius / (np.tan(np.radians(30.0)) * fill)

    rend = rendering.OffscreenRenderer(W, H)
    rend.scene.set_background(np.array(BG, dtype=np.float32))
    mat = rendering.MaterialRecord()
    if isinstance(geom, o3d.geometry.PointCloud):
        mat.shader = "defaultUnlit"
        mat.point_size = point_size
    else:
        mat.shader = "defaultLit"
        mat.base_color = [0.78, 0.80, 0.85, 1.0]
    rend.scene.add_geometry("g", geom, mat)
    rend.scene.scene.enable_sun_light(True)
    rend.scene.scene.set_sun_light([-0.3, -0.9, -0.5], [1.0, 1.0, 1.0], 75000)

    tiles = []
    for name, az, el in VIEWS:
        a, e = np.radians(az), np.radians(el)
        eye = center + dist * np.array([np.cos(e) * np.sin(a), np.sin(e), np.cos(e) * np.cos(a)])
        up = [0, 1, 0] if el < 80 else [0, 0, -1]
        rend.scene.camera.look_at(center, eye, up)
        tiles.append((name, Image.fromarray(np.asarray(rend.render_to_image()))))
    del rend

    sheet = Image.new("RGB", (W * 3, H * 2), (30, 33, 38))
    d = ImageDraw.Draw(sheet)
    for i, (name, im) in enumerate(tiles):
        x, y = (i % 3) * W, (i // 3) * H
        sheet.paste(im, (x, y))
        d.text((x + 10, y + 8), name, fill=(210, 215, 225))
    if label:
        d.text((10, H * 2 - 20), label, fill=(160, 168, 180))
    sheet.save(out_prefix + "_sheet.png")
    tiles[5][1].save(out_prefix + "_iso.png")
    return out_prefix + "_sheet.png"


def load(path):
    if path.endswith((".ply", ".obj")):
        m = o3d.io.read_triangle_mesh(path)
        if len(m.triangles) > 0:
            m.compute_vertex_normals()
            return m
    if path.endswith(".obj"):
        raise ValueError("empty obj: " + path)
    return o3d.io.read_point_cloud(path)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("input"); ap.add_argument("out_prefix")
    ap.add_argument("--point-size", type=float, default=2.5)
    ap.add_argument("--label", default=None)
    a = ap.parse_args()
    print(render_geometry(load(a.input), a.out_prefix,
                          point_size=a.point_size, label=a.label))