import rospy
# ROS Image message
from sensor_msgs.msg import Image, CameraInfo
# ROS Image message -> OpenCV2 image converter
from cv_bridge import CvBridge, CvBridgeError
from gazebo_msgs.msg import LinkStates

# !/usr/bin/env python
import roslib

# roslib.load_manifest('learning_tf')
import rospy

import tf
# import turtlesim.msg

# OpenCV2 for saving an image
import cv2
import numpy as np
from utils.projection import projectToCamera
from utils.transforms import getTransform, get_transform_tree_dict
from utils.draw import drawSquare2D

# Instantiate CvBridge
bridge = CvBridge()
link_states = None
camera_info = None
flip = 0
br = tf.TransformBroadcaster()


def image_callback(msg):
    # print("Received an image!")

    # Convert your ROS Image message to OpenCV2
    cv2_img = bridge.imgmsg_to_cv2(msg, "bgr8")

    camera_info_K = np.array(camera_info.K).reshape([3, 3])
    camera_info_D = np.array(camera_info.D)
    width = camera_info.width
    height = camera_info.height
    global link_states
    global br
    global flip
    # global camera_info

    file_xacro = "xacros/" + "novo.urdf.xacro"
    world_frame = "world"
    transform_dict = get_transform_tree_dict(file_xacro)
    extrinsics = getTransform("camera_2_rgb_optical_frame", world_frame, transform_dict)
    # print(extrinsics)
    # exit(0)

    i = 0
    for name in link_states.name:
        if "actor::" in name:
            joint = name[len("actor::"):]
            if joint != "actor_pose":
                x = link_states.pose[i].position.x
                y = link_states.pose[i].position.y
                z = link_states.pose[i].position.z

                q_x = link_states.pose[i].orientation.x
                q_y = link_states.pose[i].orientation.y
                q_z = link_states.pose[i].orientation.z
                q_w = link_states.pose[i].orientation.w

                if joint == "RightShoulder":
                    if flip < 15:
                        print("inside if ")
                        x = -0.016618864660828205
                        y = 0.26117820144569048
                        z = 0.047082693145915247

                        q_x = -0.024417754149575747
                        q_y = 0.028400752121466692
                        q_z = 0.6233825158909817
                        q_w = 0.7810194680305067

                        flip += 1

                        print(link_states.pose[i])
                    else:
                        print("outside if")


                pts_in_world = np.ndarray((4, 1), dtype=float)
                pts_in_world[0][0] = x
                pts_in_world[1][0] = y
                pts_in_world[2][0] = z
                pts_in_world[3][0] = 1

                br.sendTransform((x, y, z),
                                 (q_x, q_y, q_z, q_w),
                                 rospy.Time.now(),
                                 joint,
                                 "world")

                pts_in_sensor = np.dot(extrinsics, pts_in_world)

                # Project 3D point to camera
                pts_in_image, _, _ = projectToCamera(camera_info_K, camera_info_D,
                                                     width,
                                                     height, pts_in_sensor[0:3, :])

                xpix_projected = pts_in_image[0][0]
                ypix_projected = pts_in_image[1][0]

                drawSquare2D(cv2_img, xpix_projected, ypix_projected, 10, color=(0, 0, 0), thickness=3)

        i = i + 1

    # Save your OpenCV2 image as a jpeg
    cv2.imshow('window', cv2_img)
    cv2.waitKey(30)


def link_states_callback(msg):
    global link_states
    link_states = msg


def main():
    rospy.init_node('image_listener')
    # Define your image topic
    image_topic = "/camera_2/rgb/image_raw"
    topic_link_states = '/gazebo/link_states'
    # Set up your subscriber and define its callback

    rospy.Subscriber(topic_link_states, LinkStates, link_states_callback)
    rospy.Subscriber(image_topic, Image, image_callback)

    global camera_info
    camera_info = rospy.wait_for_message('/camera_2/rgb/camera_info', CameraInfo)

    # Spin until ctrl + c
    rospy.spin()


if __name__ == '__main__':
    main()