from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import _init_paths

import os
import json
import cv2
import cv2.aruco as aruco
import numpy as np
import sys

import rospy
from std_msgs.msg import Bool
from std_msgs.msg import Float64MultiArray
from sensor_msgs.msg import Image, CameraInfo
from cv_bridge import CvBridge, CvBridgeError
import message_filters

import torch

from external.nms import soft_nms
from opts import opts
from logger import Logger
from utils.utils import AverageMeter
from datasets.dataset_factory import dataset_factory
from detectors.detector_factory import detector_factory
# transformation from the robot base to aruco tag
M_BL = np.array([[1., 0., 0.,  0.30000],
                 [0., 1., 0.,  0.32000],
                 [0., 0., 1.,  -0.0450],
                 [0., 0., 0.,  1.00000]])

# default transformation from the camera to aruco tag
default_M_CL = np.array([[-0.07134498, -0.99639369,  0.0459293,  -0.13825178],
                         [-0.8045912,   0.03027403, -0.59305689,  0.08434352],
                         [ 0.58952768, -0.07926594, -0.8038495,   0.66103522],
                         [ 0.,          0.,          0.,          1.        ]]
                        )

# camera intrinsic matrix of Realsense D435
cameraMatrix = np.array([[607.47165, 0.0,  325.90064],
                         [0.0, 606.30420, 240.91934],
                         [0.0, 0.0, 1.0]])

# distortion of Realsense D435
distCoeffs = np.array([0.08847, -0.04283, 0.00134, -0.00102, 0.0])

# initialize GKNet Detector
opt = opts().parse()
Dataset = dataset_factory[opt.dataset]
opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
print(opt)
Detector = detector_factory[opt.task]
detector = Detector(opt)

# Publisher of perception result
pub_res = rospy.Publisher('/result', Float64MultiArray, queue_size=10)

def get_M_CL_info(gray, image_init, visualize=False):
    # parameters
    markerLength_CL = 0.093
    aruco_dict_CL = aruco.Dictionary_get(aruco.DICT_ARUCO_ORIGINAL)
    # aruco_dict_CL = aruco.Dictionary_get(aruco.DICT_6X6_250)
    parameters = aruco.DetectorParameters_create()

    corners_CL, ids_CL, rejectedImgPoints = aruco.detectMarkers(gray, aruco_dict_CL, parameters=parameters)

    # for the first frame, it may contain nothing
    if ids_CL is None:
        return default_M_CL, None

    rvec_CL, tvec_CL, _objPoints_CL = aruco.estimatePoseSingleMarkers(corners_CL[0], markerLength_CL,
                                                                      cameraMatrix, distCoeffs)
    dst_CL, jacobian_CL = cv2.Rodrigues(rvec_CL)
    M_CL = np.zeros((4, 4))
    M_CL[:3, :3] = dst_CL
    M_CL[:3, 3] = tvec_CL
    M_CL[3, :] = np.array([0, 0, 0, 1])

    if visualize:
        # print('aruco is located at mean position (%d, %d)' %(mean_x ,mean_y))
        aruco.drawAxis(image_init, cameraMatrix, distCoeffs, rvec_CL, tvec_CL, markerLength_CL)
    return M_CL, corners_CL[0][0, :, :]


def aruco_tag_remove(rgb_image, corners):
    img_out = rgb_image.copy()

    # find the top-left and right-bottom corners
    min = sys.maxsize
    max = -sys.maxsize
    tl_pxl = None
    br_pxl = None
    for corner in corners:
        if corner[0] + corner[1] < min:
            min = corner[0] + corner[1]
            tl_pxl = [int(corner[0]), int(corner[1])]

        if corner[0] + corner[1] > max:
            max = corner[0] + corner[1]
            br_pxl = [int(corner[0]), int(corner[1])]

    # get the replacement pixel value
    rep_color = img_out[tl_pxl[0] - 10, tl_pxl[1] - 10, :]

    for h in range(tl_pxl[1] - 45, br_pxl[1] + 46):
        for w in range(tl_pxl[0] - 45, br_pxl[0] + 46):
            img_out[h, w, :] = rep_color

    return img_out

def project(pixel, depth_image, M_CL, M_BL, cameraMatrix):
    '''
     project 2d pixel on the image to 3d by depth info
     :param pixel: x, y
     :param M_CL: trans from camera to aruco tag
     :param cameraMatrix: camera intrinsic matrix
     :param depth_image: depth image
     :param depth_scale: depth scale that trans raw data to mm
     :return:
     q_B: 3d coordinate of pixel with respect to base frame
     '''
    depth = depth_image[pixel[1], pixel[0]]

    # if the depth of the detected pixel is 0, check the depth of its neighbors
    # by counter-clock wise
    nei_range = 1
    while depth == 0:
        for delta_x in range(-nei_range, nei_range + 1):
            for delta_y in range(-nei_range, nei_range + 1):
                nei = [pixel[0] + delta_x, pixel[1] + delta_y]
                depth = depth_image[nei[1], nei[0]]

                if depth != 0:
                    break

            if depth != 0:
                break

        nei_range += 1

    pxl = np.linalg.inv(cameraMatrix).dot(
        np.array([pixel[0] * depth, pixel[1] * depth, depth]))
    q_C = np.array([pxl[0], pxl[1], pxl[2], 1])
    q_L = np.linalg.inv(M_CL).dot(q_C)
    q_B = M_BL.dot(q_L)

    return q_B

def pre_process(rgb_img, depth_img):
    inp_image = rgb_img
    inp_image[:, :, 0] = depth_img

    inp_image = cv2.resize(inp_image, (256, 256))

    return inp_image

def kinect_rgbd_callback(rgb_data, depth_data):
    """
    Save raw RGB and depth input from Kinect V1
    :param rgb_data: RGB image
    :param depth_data: raw depth image
    :return: None
    """
    try:
        cv_rgb = cv_bridge.imgmsg_to_cv2(rgb_data, "bgr8")
        cv_depth = cv_bridge.imgmsg_to_cv2(depth_data, "32FC1")

        cv_rgb_arr = np.array(cv_rgb, dtype=np.uint8)
        cv_depth_arr = np.array(cv_depth, dtype=np.float32)
        # cv_depth_arr = np.nan_to_num(cv_depth_arr)

        cv2.imshow("Depth", cv_depth)
        cv2.imshow("RGB", cv_rgb)

        img = cv_rgb_arr.copy()
        depth_raw = cv_depth_arr.copy()

        gray = img.astype(np.uint8)
        depth = (depth_raw * 1000).astype(np.uint8)

        # get the current transformation from the camera to aruco tag
        M_CL, corners = get_M_CL_info(gray, img, False)

        # remove aruco tag from input image to avoid mis-detection
        if corners is not None:
            img_wo_at = aruco_tag_remove(img, corners)

        # replace blue channel with the depth channel
        inp_image = pre_process(img_wo_at, depth)

        # pass the image into the network
        ret = detector.run(inp_image[:, :, :])
        ret = ret["results"]

        loc_ori = KpsToGrasppose(ret, img, depth_raw, M_CL, M_BL, cameraMatrix)
        pub_res.publish(loc_ori)

    except CvBridgeError as e:
        print(e)

def isWithinRange(pxl, w, h):
    x, y = pxl[:]

    return w/12. <= x <= 11*w/12 and h/12. <= y <= 11*h/12

def KpsToGrasppose(net_output, rgb_img, depth_map, M_CL, M_BL, cameraMatrix, visualize=True):
    kps_pr = []
    for category_id, preds in net_output.items():
        if len(preds) == 0:
            continue

        for pred in preds:
            kps = pred[:4]
            score = pred[-1]
            kps_pr.append([kps[0], kps[1], kps[2], kps[3], score])

    # no detection
    if len(kps_pr) == 0:
        return [0, 0, 0, 0]

    # sort by the confidence score
    kps_pr = sorted(kps_pr, key=lambda x: x[-1], reverse=True)
    # select the top 1 grasp prediction within the workspace
    res = None
    for kp_pr in kps_pr:
        f_w, f_h = 640. / 512., 480. / 512.
        kp_lm = (int(kp_pr[0] * f_w), int(kp_pr[1] * f_h))
        kp_rm = (int(kp_pr[2] * f_w), int(kp_pr[3] * f_h))

        if isWithinRange(kp_lm, 640, 480) and isWithinRange(kp_rm, 640, 480):
            res = kp_pr
            break

    if res is None:
        return [0, 0, 0, 0]

    f_w, f_h = 640./512., 480./512.
    kp_lm = (int(res[0]*f_w), int(res[1]*f_h))
    kp_rm = (int(res[2]*f_w), int(res[3]*f_h))
    center = (int((kp_lm[0]+kp_rm[0])/2), int((kp_lm[1]+kp_rm[1])/2))

    kp_lm_3d = project(kp_lm, depth_map, M_CL, M_BL, cameraMatrix)
    kp_rm_3d = project(kp_rm, depth_map, M_CL, M_BL, cameraMatrix)
    center_3d = project(center, depth_map, M_CL, M_BL, cameraMatrix)

    orientation = np.arctan2(kp_rm_3d[1] - kp_lm_3d[1], kp_rm_3d[0] - kp_lm_3d[0])
    # motor 7 is clockwise
    if orientation > np.pi / 2:
        orientation = np.pi - orientation
    elif orientation < -np.pi / 2:
        orientation = -np.pi - orientation
    else:
        orientation = -orientation

    # compute the open width
    dist = np.linalg.norm(kp_lm_3d[:2] - kp_rm_3d[:2])

    # draw arrow for left-middle and right-middle key-points
    lm_ep = (int(kp_lm[0] + (kp_rm[0] - kp_lm[0]) / 5.), int(kp_lm[1] + (kp_rm[1] - kp_lm[1]) / 5.))
    rm_ep = (int(kp_rm[0] + (kp_lm[0] - kp_rm[0]) / 5.), int(kp_rm[1] + (kp_lm[1] - kp_rm[1]) / 5.))
    rgb_img = cv2.arrowedLine(rgb_img, kp_lm, lm_ep, (0, 0, 0), 2)
    rgb_img = cv2.arrowedLine(rgb_img, kp_rm, rm_ep, (0, 0, 0), 2)
    # draw left-middle, right-middle and center key-points
    rgb_img = cv2.circle(rgb_img, (int(kp_lm[0]), int(kp_lm[1])), 2, (0, 0, 255), 2)
    rgb_img = cv2.circle(rgb_img, (int(kp_rm[0]), int(kp_rm[1])), 2, (0, 0, 255), 2)
    rgb_img = cv2.circle(rgb_img, (int(center[0]), int(center[1])), 2, (0, 0, 255), 2)

    if visualize:
        cv2.namedWindow('visual', cv2.WINDOW_AUTOSIZE)
        cv2.imshow('visual', rgb_img)

    return [center_3d[0], center_3d[1], center_3d[2], orientation, dist]

if __name__ == '__main__':
    # initialize ros node
    rospy.init_node("Static_grasping")

    # Bridge to convert ROS Image type to OpenCV Image type
    cv_bridge = CvBridge()
    cv2.WITH_QT = False
    # Get camera calibration parameters
    cam_param = rospy.wait_for_message('/camera/rgb/camera_info', CameraInfo, timeout=None)

    # Subscribe to rgb and depth channel
    image_sub = message_filters.Subscriber("/camera/rgb/image_rect_color", Image)
    depth_sub = message_filters.Subscriber("/camera/depth_registered/image", Image)
    ts = message_filters.ApproximateTimeSynchronizer([image_sub, depth_sub], 1, 0.1)
    ts.registerCallback(kinect_rgbd_callback)

    rospy.spin()
[((30, 7, 33, 41), 'numpy.array', 'np.array', ({(30, 16, 33, 40): '[[1.0, 0.0, 0.0, 0.3], [0.0, 1.0, 0.0, 0.32], [0.0, 0.0, 1.0, -0.045], [0.0,\n    0.0, 0.0, 1.0]]'}, {}), '([[1.0, 0.0, 0.0, 0.3], [0.0, 1.0, 0.0, 0.32], [0.0, 0.0, 1.0, -\n    0.045], [0.0, 0.0, 0.0, 1.0]])', True, 'import numpy as np\n'), ((36, 15, 40, 25), 'numpy.array', 'np.array', ({(36, 24, 39, 78): '[[-0.07134498, -0.99639369, 0.0459293, -0.13825178], [-0.8045912, \n    0.03027403, -0.59305689, 0.08434352], [0.58952768, -0.07926594, -\n    0.8038495, 0.66103522], [0.0, 0.0, 0.0, 1.0]]'}, {}), '([[-0.07134498, -0.99639369, 0.0459293, -0.13825178], [-0.8045912, \n    0.03027403, -0.59305689, 0.08434352], [0.58952768, -0.07926594, -\n    0.8038495, 0.66103522], [0.0, 0.0, 0.0, 1.0]])', True, 'import numpy as np\n'), ((43, 15, 45, 42), 'numpy.array', 'np.array', ({(43, 24, 45, 41): '[[607.47165, 0.0, 325.90064], [0.0, 606.3042, 240.91934], [0.0, 0.0, 1.0]]'}, {}), '([[607.47165, 0.0, 325.90064], [0.0, 606.3042, 240.91934], [0.0, \n    0.0, 1.0]])', True, 'import numpy as np\n'), ((48, 13, 48, 66), 'numpy.array', 'np.array', ({(48, 22, 48, 65): '[0.08847, -0.04283, 0.00134, -0.00102, 0.0]'}, {}), '([0.08847, -0.04283, 0.00134, -0.00102, 0.0])', True, 'import numpy as np\n'), ((59, 10, 59, 70), 'rospy.Publisher', 'rospy.Publisher', (), '', False, 'import rospy\n'), ((64, 20, 64, 67), 'cv2.aruco.Dictionary_get', 'aruco.Dictionary_get', ({(64, 41, 64, 66): 'aruco.DICT_ARUCO_ORIGINAL'}, {}), '(aruco.DICT_ARUCO_ORIGINAL)', True, 'import cv2.aruco as aruco\n'), ((66, 17, 66, 50), 'cv2.aruco.DetectorParameters_create', 'aruco.DetectorParameters_create', ({}, {}), '()', True, 'import cv2.aruco as aruco\n'), ((68, 44, 68, 107), 'cv2.aruco.detectMarkers', 'aruco.detectMarkers', (), '', True, 'import cv2.aruco as aruco\n'), ((74, 38, 75, 95), 'cv2.aruco.estimatePoseSingleMarkers', 'aruco.estimatePoseSingleMarkers', ({(74, 70, 74, 83): 'corners_CL[0]', (74, 85, 74, 100): 'markerLength_CL', (75, 70, 75, 82): 'cameraMatrix', (75, 84, 75, 94): 'distCoeffs'}, {}), '(corners_CL[0], markerLength_CL,\n    cameraMatrix, distCoeffs)', True, 'import cv2.aruco as aruco\n'), ((76, 26, 76, 48), 'cv2.Rodrigues', 'cv2.Rodrigues', ({(76, 40, 76, 47): 'rvec_CL'}, {}), '(rvec_CL)', False, 'import cv2\n'), ((77, 11, 77, 27), 'numpy.zeros', 'np.zeros', ({(77, 20, 77, 26): '(4, 4)'}, {}), '((4, 4))', True, 'import numpy as np\n'), ((80, 17, 80, 39), 'numpy.array', 'np.array', ({(80, 26, 80, 38): '[0, 0, 0, 1]'}, {}), '([0, 0, 0, 1])', True, 'import numpy as np\n'), ((146, 10, 146, 47), 'numpy.array', 'np.array', ({(146, 19, 146, 46): '[pxl[0], pxl[1], pxl[2], 1]'}, {}), '([pxl[0], pxl[1], pxl[2], 1])', True, 'import numpy as np\n'), ((156, 16, 156, 49), 'cv2.resize', 'cv2.resize', ({(156, 27, 156, 36): 'inp_image', (156, 38, 156, 48): '(256, 256)'}, {}), '(inp_image, (256, 256))', False, 'import cv2\n'), ((249, 18, 249, 82), 'numpy.arctan2', 'np.arctan2', ({(249, 29, 249, 54): 'kp_rm_3d[1] - kp_lm_3d[1]', (249, 56, 249, 81): 'kp_rm_3d[0] - kp_lm_3d[0]'}, {}), '(kp_rm_3d[1] - kp_lm_3d[1], kp_rm_3d[0] - kp_lm_3d[0])', True, 'import numpy as np\n'), ((259, 11, 259, 54), 'numpy.linalg.norm', 'np.linalg.norm', ({(259, 26, 259, 53): 'kp_lm_3d[:2] - kp_rm_3d[:2]'}, {}), '(kp_lm_3d[:2] - kp_rm_3d[:2])', True, 'import numpy as np\n'), ((264, 14, 264, 66), 'cv2.arrowedLine', 'cv2.arrowedLine', ({(264, 30, 264, 37): 'rgb_img', (264, 39, 264, 44): 'kp_lm', (264, 46, 264, 51): 'lm_ep', (264, 53, 264, 62): '(0, 0, 0)', (264, 64, 264, 65): '2'}, {}), '(rgb_img, kp_lm, lm_ep, (0, 0, 0), 2)', False, 'import cv2\n'), ((265, 14, 265, 66), 'cv2.arrowedLine', 'cv2.arrowedLine', ({(265, 30, 265, 37): 'rgb_img', (265, 39, 265, 44): 'kp_rm', (265, 46, 265, 51): 'rm_ep', (265, 53, 265, 62): '(0, 0, 0)', (265, 64, 265, 65): '2'}, {}), '(rgb_img, kp_rm, rm_ep, (0, 0, 0), 2)', False, 'import cv2\n'), ((279, 4, 279, 38), 'rospy.init_node', 'rospy.init_node', ({(279, 20, 279, 37): '"""Static_grasping"""'}, {}), "('Static_grasping')", False, 'import rospy\n'), ((282, 16, 282, 26), 'cv_bridge.CvBridge', 'CvBridge', ({}, {}), '()', False, 'from cv_bridge import CvBridge, CvBridgeError\n'), ((285, 16, 285, 91), 'rospy.wait_for_message', 'rospy.wait_for_message', (), '', False, 'import rospy\n'), ((288, 16, 288, 81), 'message_filters.Subscriber', 'message_filters.Subscriber', ({(288, 43, 288, 73): '"""/camera/rgb/image_rect_color"""', (288, 75, 288, 80): 'Image'}, {}), "('/camera/rgb/image_rect_color', Image)", False, 'import message_filters\n'), ((289, 16, 289, 83), 'message_filters.Subscriber', 'message_filters.Subscriber', ({(289, 43, 289, 75): '"""/camera/depth_registered/image"""', (289, 77, 289, 82): 'Image'}, {}), "('/camera/depth_registered/image', Image)", False, 'import message_filters\n'), ((290, 9, 290, 84), 'message_filters.ApproximateTimeSynchronizer', 'message_filters.ApproximateTimeSynchronizer', ({(290, 53, 290, 75): '[image_sub, depth_sub]', (290, 77, 290, 78): '1', (290, 80, 290, 83): '0.1'}, {}), '([image_sub, depth_sub], 1, 0.1)', False, 'import message_filters\n'), ((293, 4, 293, 16), 'rospy.spin', 'rospy.spin', ({}, {}), '()', False, 'import rospy\n'), ((51, 6, 51, 12), 'opts.opts', 'opts', ({}, {}), '()', False, 'from opts import opts\n'), ((53, 6, 53, 12), 'opts.opts', 'opts', ({}, {}), '()', False, 'from opts import opts\n'), ((84, 8, 84, 95), 'cv2.aruco.drawAxis', 'aruco.drawAxis', ({(84, 23, 84, 33): 'image_init', (84, 35, 84, 47): 'cameraMatrix', (84, 49, 84, 59): 'distCoeffs', (84, 61, 84, 68): 'rvec_CL', (84, 70, 84, 77): 'tvec_CL', (84, 79, 84, 94): 'markerLength_CL'}, {}), '(image_init, cameraMatrix, distCoeffs, rvec_CL, tvec_CL,\n    markerLength_CL)', True, 'import cv2.aruco as aruco\n'), ((145, 8, 145, 61), 'numpy.array', 'np.array', ({(145, 17, 145, 60): '[pixel[0] * depth, pixel[1] * depth, depth]'}, {}), '([pixel[0] * depth, pixel[1] * depth, depth])', True, 'import numpy as np\n'), ((171, 21, 171, 53), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((172, 23, 172, 59), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((175, 8, 175, 37), 'cv2.imshow', 'cv2.imshow', ({(175, 19, 175, 26): '"""Depth"""', (175, 28, 175, 36): 'cv_depth'}, {}), "('Depth', cv_depth)", False, 'import cv2\n'), ((176, 8, 176, 33), 'cv2.imshow', 'cv2.imshow', ({(176, 19, 176, 24): '"""RGB"""', (176, 26, 176, 32): 'cv_rgb'}, {}), "('RGB', cv_rgb)", False, 'import cv2\n'), ((272, 8, 272, 54), 'cv2.namedWindow', 'cv2.namedWindow', ({(272, 24, 272, 32): '"""visual"""', (272, 34, 272, 53): 'cv2.WINDOW_AUTOSIZE'}, {}), "('visual', cv2.WINDOW_AUTOSIZE)", False, 'import cv2\n'), ((273, 8, 273, 37), 'cv2.imshow', 'cv2.imshow', ({(273, 19, 273, 27): '"""visual"""', (273, 29, 273, 36): 'rgb_img'}, {}), "('visual', rgb_img)", False, 'import cv2\n'), ((144, 10, 144, 37), 'numpy.linalg.inv', 'np.linalg.inv', ({(144, 24, 144, 36): 'cameraMatrix'}, {}), '(cameraMatrix)', True, 'import numpy as np\n'), ((147, 10, 147, 29), 'numpy.linalg.inv', 'np.linalg.inv', ({(147, 24, 147, 28): 'M_CL'}, {}), '(M_CL)', True, 'import numpy as np\n')]