#!/usr/bin/env python

import numpy as np
import trimesh

import objslampp

models = objslampp.datasets.YCBVideoModels()
dataset = objslampp.datasets.YCBVideoDataset('train')
frame = dataset[0]

rgb = frame['color']
depth = frame['depth']
label = frame['label']
K = frame['meta']['intrinsic_matrix']
class_ids = frame['meta']['cls_indexes']

class_id = class_ids[0]
mask = label == class_id
cad_file = models.get_cad_model(class_id=class_id)
dim = 32
pitch = models.get_bbox_diagonal(cad_file) / dim

pcd = objslampp.geometry.pointcloud_from_depth(
    depth, fx=K[0, 0], fy=K[1, 1], cx=K[0, 2], cy=K[1, 2]
)

pcd_ins = pcd[mask]
centroid = np.nanmean(pcd_ins, axis=0)
aabb_min = centroid - pitch * dim / 2
aabb_max = aabb_min + pitch * dim

camera = trimesh.scene.Camera(
    resolution=(rgb.shape[1], rgb.shape[0]),
    focal=(K[0, 0], K[1, 1]),
)

centers = trimesh.voxel.matrix_to_points(
    np.ones((dim, dim, dim)),
    pitch=pitch,
    origin=aabb_min + pitch,
)

r, c = objslampp.geometry.project_to_camera(
    centers,
    fx=K[0, 0],
    fy=K[1, 1],
    cx=K[0, 2],
    cy=K[1, 2],
    image_shape=rgb.shape[:2],
)
r, c = r.round().astype(int), c.round().astype(int)
is_target = mask[r, c]
z_depth = depth[r, c]
z_voxel = centers[:, 2]

nonnan = ~np.isnan(z_depth)
# empty: 1, target: 2, other: 3, unknown: 4
labels = np.full((centers.shape[0],), 4, dtype=np.int32)
labels[np.greater((z_depth - z_voxel), pitch / 2, where=nonnan)] = 1
is_occupied = np.less(np.abs(z_depth - z_voxel), pitch / 2, where=nonnan)
labels[is_occupied & is_target] = 2
labels[is_occupied & (~is_target)] = 3

# -----------------------------------------------------------------------------

scene = trimesh.Scene()

nonnan = ~np.isnan(pcd).any(axis=2)
geom = trimesh.PointCloud(vertices=pcd[mask & nonnan], colors=rgb[mask & nonnan])
scene.add_geometry(geom)

import imgviz
# keep = labels == 4  # unknown
# keep = np.isin(labels, [2, 3, 4])
keep = np.isin(labels, [3])
centers = centers[keep]
labels = labels[keep]

colors = imgviz.label2rgb(labels[None])[0]
colors = np.hstack((colors, [[127]] * len(colors)))
boxes = trimesh.voxel.multibox(centers, pitch=pitch, colors=colors)
scene.add_geometry(boxes)

geom = trimesh.creation.box(extents=aabb_max - aabb_min)
geom.apply_translation(aabb_min + (aabb_max - aabb_min) / 2)
# geom.visual.face_colors = (0, 1., 0, 0.5)
# scene.add_geometry(geom)
geom = objslampp.extra.trimesh.box_to_wired_box(geom)
scene.add_geometry(geom)

scene.camera.resolution = camera.resolution
scene.camera.focal = camera.focal
scene.camera.transform = objslampp.extra.trimesh.camera_transform(np.eye(4))

scene.show()