#!/usr/bin/env python

import chainer
import chainer.functions as F
import numpy as np
import trimesh
import trimesh.transformations as tf

import imgviz

import objslampp


np.random.seed(0)

pcd_file = objslampp.datasets.YCBVideoModels()\
    .get_model(class_id=20)['points_xyz']
points = np.loadtxt(pcd_file, dtype=np.float32)
indices = np.random.permutation(len(points))[:1000]
points = points[indices]

dim = 16
pitch = max(points.max(axis=0) - points.min(axis=0)) * 1.1 / dim
origin = (- pitch * dim / 2,) * 3

points_reference = points

print('voxelization start')
print(f'pitch: {pitch}')
print(f'dim: {dim}')
grid_reference = objslampp.functions.occupancy_grid_3d(
    points,
    pitch=pitch,
    origin=origin,
    dimension=(dim, dim, dim),
    connectivity=2,
).array
if 0:
    trimesh.Scene(trimesh.voxel.Voxel(grid_reference, pitch, origin).as_boxes()).show()

index = 0
quaternion = None
while True:
    index += 1

    print('>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>')
    if quaternion is None:
        quaternion = tf.random_quaternion().astype(np.float32)
        quaternion = chainer.Variable(quaternion)
    else:
        quaternion = chainer.Variable(quaternion.array)

    print(type(quaternion), quaternion, quaternion.grad)

    transform = objslampp.functions.quaternion_matrix(quaternion[None])
    points_observed = objslampp.functions.transform_points(points_reference, transform)[0]

    add = objslampp.metrics.average_distance(
        [points_reference], [np.eye(4)], [transform.array[0]]
    )[0]
    print('add:', add)

    if 1:
        geom1 = trimesh.PointCloud(vertices=points_reference, colors=[(1., 0, 0)]*len(points_reference))
        geom2 = trimesh.PointCloud(vertices=points_observed.array, colors=[(0, 1., 0)]*len(points_reference))
        scene = trimesh.Scene([geom1, geom2])
        scene.set_camera(angles=(np.deg2rad(0), np.deg2rad(0), 0), distance=0.3)
        image = objslampp.extra.trimesh.save_image(scene)
        imgviz.io.imsave(f'points_{index:08d}.jpg', image[:, :, :3])

    grid_observed = objslampp.functions.occupancy_grid_3d(
        points_observed,
        pitch=pitch,
        origin=origin,
        dimension=(dim, dim, dim),
        connectivity=2,
    )
    if 1:
        grid = grid_observed.array
        colors = imgviz.depth2rgb(grid.reshape(1, -1), min_value=0, max_value=1)
        colors = colors.reshape(dim, dim, dim, 3)
        colors = np.concatenate((colors, np.full((dim, dim, dim, 1), 127)), axis=3)
        geom = trimesh.voxel.Voxel(grid, pitch, origin).as_boxes()
        geom.apply_translation((pitch / 2, pitch / 2, pitch / 2))
        I, J, K = zip(*np.argwhere(grid))
        geom.visual.face_colors = colors[I, J, K].repeat(12, axis=0)
        geom1 = geom

        grid = grid_reference
        colors = imgviz.depth2rgb(grid.reshape(1, -1), min_value=0, max_value=1)
        colors = colors.reshape(dim, dim, dim, 3)
        colors = np.concatenate((colors, np.full((dim, dim, dim, 1), 127)), axis=3)
        geom = trimesh.voxel.Voxel(grid, pitch, origin).as_boxes()
        geom.apply_translation((pitch / 2, pitch / 2, pitch / 2))
        I, J, K = zip(*np.argwhere(grid))
        geom.visual.face_colors = colors[I, J, K].repeat(12, axis=0)
        geom2 = geom

        scene = trimesh.Scene([geom1, geom2])
        scene.set_camera(angles=(np.deg2rad(0), np.deg2rad(0), 0), distance=0.3)
        image = objslampp.extra.trimesh.save_image(scene)
        imgviz.io.imsave(f'occupancy_{index:08d}.jpg', image[:, :, :3])

    loss = F.mean_squared_error(grid_reference, grid_observed)
    loss.backward()
    print('loss:', float(loss.array))
    del loss, grid_observed

    print(quaternion.grad)
    lr = 1
    quaternion = quaternion - lr * quaternion.grad
    quaternion.cleargrad()
    print('<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<')