'''
Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.

NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
'''
import argparse
import numpy as np
import cv2    # for 2d minimum rotated rectangle fitting and similarity color assignment in LAB
import quaternion    # for convert cv rotation matrix to quaternion

from isaac import Application, Codelet
'''
Contains PyCodelets and helper functions to compute 3d pose of colored cubes from RGB-D image
from a near-top-down camera. It uses a geometry method based on surflets generated by the
superpixel algorithm to compute 3d poses.

This is intended as an example for fast-prototyping using Isaac's python API.
'''

ELEMENT_TYPE_INT = 'int32'
COLOR_TEMPLATES = {
    "red": [128, 0, 0],
    "green": [0, 128, 0],
    "blue": [0, 0, 128],
    "yellow": [128, 128, 0]
}
CUBE_LENGTH = 0.048
PLANE_DISTANCE_THRESHOLD = 0.005


def fit_plane_to_points(data):
    '''
    Fits a plane normal to a set of 3d points using SVD, and return the normal and offset of the
    plane model.

    Input:
      data: float32 2D (Nx3) numpy array of the input 3D point cloud

    Output:
      normal, offset of the plane model
    '''
    centroid = np.mean(data, axis=0)
    X = data - centroid
    u, s, vh = np.linalg.svd(X, full_matrices=True)
    normal = vh[-1, :]
    offset = -np.dot(normal, centroid)
    return normal, offset


def detect_plane(data, distance_threshold=0.005, max_iterations=10):
    '''
    Given a point cloud, return a mask for inliers and the plane model. This method iteractively
    fits plane to inliers points until the set of inliers are unchanged

    Input:
      data: float32 2D (Nx3) numpy array of the input 3D point cloud
      distance_threshold: threshold for point to plane distance for inliers

    Output:
      inliers: boolean numpy array of size N, mask for inliers
      model: plane model of (normal, offset)
    '''
    new_inliers = np.ones((data.shape[0]), dtype=bool)
    for i in range(max_iterations):
        inliers = new_inliers
        if np.sum(inliers) < 3:
            break
        normal, offset = fit_plane_to_points(data[inliers, :])
        dist = np.abs(np.matmul(data, normal) + offset)
        new_inliers = dist <= distance_threshold
        if np.all(inliers == new_inliers):
            break
    return inliers, np.append(normal, offset)


def estimate_pose(points,
                  block_size,
                  distance_threshold,
                  fit_rotation=True,
                  model=None,
                  min_size_ratio=0.5,
                  max_size_ratio=1.5):
    '''
    Estimates the pose of a block from a near-top-down viewpoint by first fitting the top plane
    of the block, then project points to the top plane and use in-plane minimum rectangle to
    estimate centroid position and in-plane rotation, and finally combine the top plane model
    with in-plane pose to output 3d pose

    Inputs:
      points: Nx3 numpy array of point clouds that belongs to the block
      block_size: edge length of the block
      distance_threshold: point-to-plane distance threshold for fitting the top surface plane
      fit_rotation: if True, use minimal rectangle to estimate in-plane rotation. Otherwise
                    default to 0 rotatino angle
      model: 3d plane model for the top surface. if None, estimate the plane model by fitting to
             points
      min(max)_size_ratio: if using the minimal rectangle, the ratio of the fitted bounding box
                           size to the block_size must fall between this two

    Output:
      centroid: 3d position of the block's center
      Q: rotation quaternion
    '''
    # get plane model for top plane of the block cluster
    if not model:
        _, model = detect_plane(points, distance_threshold)
        if np.count_nonzero(np.isnan(model)) > 0:
            raise ValueError("Unable to fit plane model")

    # get 3d axis of the block, making sure the offset is positive, e.g., the plane normal is
    # pointing towards the camera
    if model[3] < 0:
        model = [-x for x in model]
    az = model[:3]
    az /= np.linalg.norm(az)
    ax = np.cross(np.array([-1.0, 0.0, 0.0]), az)
    ax /= np.linalg.norm(ax)
    ay = np.cross(az, ax)

    # project to inplane (2d) to get in-plane centroid, extend and rotation
    projected = np.vstack([points.dot(ax), points.dot(ay)]).transpose().astype(np.float32)
    if fit_rotation:
        box_center, box_size, box_angle = cv2.minAreaRect(projected)
        # limit in-plane rotation to [-45, 45]
        while box_angle > 45.0:
            box_angle -= 90.0
        while box_angle < -45.0:
            box_angle += 90.0
        box_angle = np.deg2rad(box_angle)
    else:
        box_center = (np.max(projected, axis=0) + np.min(projected, axis=0)) * 0.5
        box_size = np.max(projected, axis=0) - np.min(projected, axis=0)

    # check box size is reasonable
    max_size = max_size_ratio * block_size
    min_size = min_size_ratio * block_size
    if not (min_size <= box_size[0] <= max_size) or not (min_size <= box_size[1] <= max_size):
        raise ValueError("Box size unexpected")

    # get 3d centroid of box
    projected_z = np.max(points.dot(az))
    centroid = box_center[0] * ax + box_center[1] * ay + (projected_z - block_size * 0.5) * az
    # align the in-plane axis and get the 3d rotation as quaternion
    if fit_rotation:
        ax, ay = ax * np.cos(box_angle) + ay * np.sin(box_angle), ay * np.cos(box_angle) - ax * \
                 np.sin(box_angle)
    M = np.vstack([ax, ay, az]).transpose()
    Q = quaternion.from_rotation_matrix(M)
    return centroid, Q


class ClusterPoseEstimation(Codelet):
    '''
    Estimates the 3d pose of blocks given their clusters of surflets, using estimation of the top
    plane of the block and in-plane projections of the ponits.
    '''

    def start(self):
        # Receives surflets
        self.rx_surflets = self.isaac_proto_rx("CompositeProto", "surflets")
        # Receives cluster assignment
        self.rx_assignment = self.isaac_proto_rx("TensorProto", "assignment")
        # Receives cluster assignment
        self.rx_plane_model = self.isaac_proto_rx("PlaneProto", "plane_model")
        # Publishes detections3 (prediction + pose)
        self.tx_output_poses = self.isaac_proto_tx("Detections3Proto", "output_poses")
        # Publishes filtered assignment
        self.tx_output_assignment = self.isaac_proto_tx("TensorProto", "output_assignment")

        self.tick_on_message(self.rx_assignment)
        self.synchronize(self.rx_assignment, self.rx_plane_model, self.rx_surflets)

    def tick(self):
        # parse surflets
        acqtime = self.rx_surflets.message.acqtime
        sfl = self.rx_surflets.message.tensor
        if sfl is None:
            self.log_info("Surflect tensor is None")
            return
        points = sfl.reshape((-1, 11))[:, 2:5].astype(np.float32)

        # get cluster assignment
        clusters = self.rx_assignment.message.tensor
        if clusters is None:
            self.log_info("Assignment tensor is None")
            return
        num_cluster = clusters.shape[0]
        assert points.shape[0] == clusters.shape[1], "Size of surflets do not match size of " \
                                                    "assignment"
        # get plane model
        if self.config.fit_plane:
            plane_normal = self.rx_plane_model.message.proto.normal
            plane_model = [
                plane_normal.x, plane_normal.y, plane_normal.z,
                self.rx_plane_model.message.proto.offset
            ]
        else:
            plane_model = None

        # get pose estimation for each cluster
        valid_clusters = []
        results = []
        for i in range(num_cluster):
            try:
                t, q = estimate_pose(points[clusters[i] > 0], self.config.block_length,
                                     self.config.distance_threhold, self.config.fit_rotation,
                                     plane_model)
                valid_clusters.append(i)
                results.append((t, q))
            except ValueError:
                pass

        num_cluster = len(valid_clusters)
        self.show("cluster", num_cluster)

        # publish updated cluster assignment
        tx_message = self.tx_output_assignment.init()
        tx_message.proto.elementType = ELEMENT_TYPE_INT
        tensor_size = tx_message.proto.init('sizes', 2)
        tensor_size[0] = num_cluster
        tensor_size[1] = clusters.shape[1]
        tx_message.proto.dataBufferIndex = 0
        tx_message.buffers = [clusters[valid_clusters]]
        tx_message.acqtime = acqtime
        self.tx_output_assignment.publish()

        # publish detection3 (output_poses)
        tx_message = self.tx_output_poses.init()
        predictions = tx_message.proto.init("predictions", num_cluster)
        poses = tx_message.proto.init("poses", num_cluster)
        for i in range(num_cluster):
            predictions[i].label = self.config.label
            # estimate poses for each cluster
            t, q = results[i]
            predictions[i].confidence = 1.0
            poses[i].translation.x = t[0].item()
            poses[i].translation.y = t[1].item()
            poses[i].translation.z = t[2].item()
            poses[i].rotation.q.w = q.w
            poses[i].rotation.q.x = q.x
            poses[i].rotation.q.y = q.y
            poses[i].rotation.q.z = q.z
        tx_message.acqtime = acqtime
        self.tx_output_poses.publish()


def rgb2lab(color):
    '''
    Converts array of rgb color (in uint8) to lab color

    Input:
      color: input numpy array of rgb colors in [0, 255], show be shape of [N, 3]
    Output:
      color converted to LAB colorspace
    '''
    num_pixel = int(color.size / 3)
    # resize to rank-3
    lab_cv = cv2.cvtColor(color.reshape((1, num_pixel, 3)), cv2.COLOR_RGB2LAB).reshape(num_pixel, 3)
    # cv2.cvtColor rescales LAB to 255, scale it back to L [0,100[, a/b [-128, 128[
    return np.vstack((lab_cv[:, 0] / 255.0 * 100.0, lab_cv[:, 1] - 128.0, lab_cv[:, 2] - 128.0)).T


def deltaE_cie76(colors_1, colors_2):
    '''
    Computes the CIE76 distance of two set of LAB colors

    Input:
      colors_1: a list of LAB colors
      colors_2: another list of LAB colors
    Output:
      numpy array of CIE76 distance between each pair of colors in colors_1 and colors_2
    '''
    colors_1 = colors_1.reshape(-1, 3)
    colors_2 = colors_2.reshape(-1, 3)
    row, col = colors_1.shape[0], colors_2.shape[0]
    result = np.zeros((row, col))
    for i in range(row):
        x = colors_1[i, :]
        for j in range(col):
            y = colors_2[j, :]
            result[i, j] = np.linalg.norm(x-y)
    # reshape output to lower rank if input colors_1 or colors_2 has only a single color
    if col == 1 and row == 1:
        return result[0, 0]
    if col == 1:
        return result[:, 0]
    if row == 1:
        return result[0, :]
    return result


class ColorMatcher(object):
    '''
    Helper class to query the closest color in a fixed set of color templates to any given input.
    Color distance is computed in the lab color space using deltaE_cie76
    '''

    def __init__(self, color_templates, threshold=60):
        '''
        Inputs:
          color_templates: a dictionary of color name (str) to rgb value (list of three ints)
          threshold: distance between a query color and template color to be considered a match
        '''
        self._color_labels = list(color_templates.keys())
        color_values = [color_templates[c] for c in self._color_labels]
        self._color_values = rgb2lab(np.asarray(color_values, dtype=np.uint8))
        self._threshold = threshold

    def match(self, color):
        '''
        Input:
          color: rgb value of the query color

        Output:
          Name (str) of the matched template color. None if the distance to all template colors are
          greater than the threshold
        '''
        query = rgb2lab(np.asarray(color, dtype=np.uint8))
        delta = deltaE_cie76(self._color_values, query)
        idx = np.argmin(delta)
        if delta[idx] < self._threshold:
            return self._color_labels[idx]
        return None


class ClusterColorAssignment(Codelet):
    '''
    This codelet subfixes detections3 labels with color label by matching the color for the
    corresponding surflet clusters to a set of color templates.
    '''

    def start(self):
        # Receives surflets
        self.rx_surflets = self.isaac_proto_rx("CompositeProto", "surflets")
        # Receives cluster assignment
        self.rx_assignment = self.isaac_proto_rx("TensorProto", "assignment")
        # Receives detections3 (prediction + pose)
        self.rx_input_poses = self.isaac_proto_rx("Detections3Proto", "input_poses")
        # Publishes detections3 (prediction + pose) with relabeld color
        self.tx_output_poses = self.isaac_proto_tx("Detections3Proto", "output_poses")

        self.tick_on_message(self.rx_input_poses)
        self.synchronize(self.rx_input_poses, self.rx_assignment, self.rx_surflets)

        # initialize the color matcher with color templates from config
        self._color_matcher = ColorMatcher(self.config.color_templates)

    def tick(self):
        # parse surflets
        acqtime = self.rx_surflets.message.acqtime
        sfl = self.rx_surflets.message.tensor
        if sfl is None:
            self.log_info("Surflect tensor is None")
            return
        sfl_colors = sfl.reshape((-1, 11))[:, 8:11]

        # get cluster assignment
        clusters = self.rx_assignment.message.tensor
        num_cluster = clusters.shape[0]
        # assign color
        colors = [np.median(sfl_colors[clusters[i] > 0], axis=0) for i in range(num_cluster)]
        color_assignment = [self._color_matcher.match(c) for c in colors]
        indices = [i for i in range(num_cluster) if color_assignment[i] is not None]
        num_valid = len(indices)

        # publish detection3 (output_poses)
        tx_message = self.tx_output_poses.init()
        predictions = tx_message.proto.init("predictions", num_valid)
        poses = tx_message.proto.init("poses", num_valid)
        input_proto = self.rx_input_poses.message.proto
        for i in range(num_valid):
            j = indices[i]
            # assign detection color as subfix to label
            predictions[i].label = "{0}_{1}".format(input_proto.predictions[j].label,
                                                    color_assignment[j])
            predictions[i].confidence = 1.0
            # estimate poses for each cluster
            poses[i] = input_proto.poses[j]
        tx_message.acqtime = acqtime
        self.tx_output_poses.publish()


def create_block_pose_estimation(app,
                                 block_length=CUBE_LENGTH,
                                 distance_threhold=PLANE_DISTANCE_THRESHOLD,
                                 use_refinement=False):
    '''
    Creates the block pose estimation subgraph, adds ClusterPoseEstimation and
    ClusterColorAssignment PyCodelets and connects the appropriate edges with or without refinement
    '''
    app.load(
        filename="apps/samples/pick_and_place/block_pose_estimation.subgraph.json",
        prefix="detection_pose_estimation")
    perception_interface = app.nodes["detection_pose_estimation.interface"]["Subgraph"]

    cluster_pose = app.nodes["detection_pose_estimation.block_detection"].add(
        ClusterPoseEstimation, "ClusterPoseEstimation")
    cluster_pose.config.label = "block"
    cluster_pose.config.block_length = block_length
    cluster_pose.config.fit_rotation = not use_refinement
    cluster_pose.config.distance_threhold = distance_threhold

    color_assignment = app.nodes["detection_pose_estimation.block_detection"].add(
        ClusterColorAssignment, "ClusterColorAssignment")
    color_assignment.config.color_templates = COLOR_TEMPLATES
    if use_refinement:
        pose_refinement_interface = app.nodes[
            'detection_pose_estimation.object_pose_refinement.subgraph']['interface']
        app.connect(pose_refinement_interface, "output_poses", color_assignment, "input_poses")
    else:
        app.nodes['detection_pose_estimation.object_pose_refinement' \
                  '.object_pose_refinement'].config['disable_automatic_start'] = True
        app.connect(cluster_pose, "output_poses", color_assignment, "input_poses")

    return perception_interface


def main(args):
    '''
    Creates and runs block pose estimation app on RGBD image. Input image may come from sim (
    default), cask recording or realsense camera

    To run with sim (sim needs to be running first):
      bazel run apps/samples/pick_and_place:block_pose_estimation
    To see all options:
      bazel run apps/samples/pick_and_place:block_pose_estimation -- --help
    '''
    app = Application(name="block_pose_estimation")
    perception_interface = create_block_pose_estimation(app, use_refinement=args.refinement)

    if args.mode == 'cask':
        assert args.cask is not None, "cask path not specified. use --cask to set."
        # Load replay subgraph and configure interface node
        app.load("packages/cask/apps/replay.subgraph.json", prefix="replay")
        source = app.nodes["replay.interface"].components["output"]
        source.config.cask_directory = args.cask
        source.config.loop = True
    elif args.mode == "realsense":
        # Create realsense camera codelet
        app.load_module("realsense")
        source = app.add("camera").add(app.registry.isaac.RealsenseCamera)
        source.config.rows = 720
        source.config.cols = 1280
        source.config.color_framerate = 15
        source.config.depth_framerate = 15
    elif args.mode == "sim":
        app.load(filename='packages/navsim/apps/navsim_tcp.subgraph.json', prefix='simulation')
        source = app.nodes["simulation.interface"]["output"]

    app.connect(source, "color", perception_interface, "color")
    app.connect(source, "depth", perception_interface, "depth")
    app.connect(source, "color_intrinsics", perception_interface, "intrinsics")

    app.load_module("utils")
    pose_node = app.add("pose")
    pose_injector = pose_node.add(app.registry.isaac.utils.DetectionsToPoseTree)
    pose_injector.config.detection_frame = "camera"
    app.connect(perception_interface, "output_poses", pose_injector, "detections")

    app.run()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--mode",
        help="Mode for data source",
        type=str,
        choices=["sim", "realsense", "cask"],
        default="sim")
    parser.add_argument("--cask", help="Name of the cask recording base directory", type=str)
    parser.add_argument("--refinement", help="In franka, use pose refinement", action='store_true')
    args = parser.parse_args()
    main(args)
