import sys
import threading
import time

import cv2
from matplotlib import pyplot as plt
import numpy as np
import rclpy
from cv_bridge import CvBridge
from rclpy.qos import QoSHistoryPolicy, QoSProfile, QoSReliabilityPolicy
from sensor_msgs.msg import Image

from rpt_msgs.msg import Poses
from skelda import utils_view

# ==================================================================================================

bridge = CvBridge()
node = None
publisher_img = None

img_input_topic = "/{}/pylon_ros2_camera_node/image_raw"
pose_input_topic = "/poses/{}"
img_output_topic = "/{}/img_with_pose"

last_input_image = None
lock = threading.Lock()

# ==================================================================================================


def bayer2rgb(bayer):
    img = cv2.cvtColor(bayer, cv2.COLOR_BayerBG2RGB)
    return img


# ==================================================================================================


def callback_images(image_data):
    global last_input_image, lock

    # Convert into cv images from image string
    if image_data.encoding == "bayer_rggb8":
        bayer_image = bridge.imgmsg_to_cv2(image_data, "bayer_rggb8")
        color_image = bayer2rgb(bayer_image)
    elif image_data.encoding == "mono8":
        gray_image = bridge.imgmsg_to_cv2(image_data, "mono8")
        color_image = cv2.cvtColor(gray_image, cv2.COLOR_GRAY2RGB)
    elif image_data.encoding == "rgb8":
        color_image = bridge.imgmsg_to_cv2(image_data, "rgb8")
    else:
        raise ValueError("Unknown image encoding:", image_data.encoding)

    time_stamp = image_data.header.stamp.sec + image_data.header.stamp.nanosec / 1.0e9

    with lock:
        last_input_image = (color_image, time_stamp)


# ==================================================================================================


def callback_poses(pose_data):
    global last_input_image, lock

    ptime = time.time()
    if last_input_image is None:
        return

    # Extract pose data
    joint_names = pose_data.joint_names
    bodies2D = np.array(pose_data.bodies_flat).reshape(pose_data.bodies_shape).tolist()

    # Collect inputs
    images_2d = []
    timestamps = []
    with lock:
        img = np.copy(last_input_image[0])
        ts = last_input_image[1]
        images_2d.append(img)
        timestamps.append(ts)

    # Visualize 2D poses
    colors = plt.cm.hsv(np.linspace(0, 1, len(bodies2D), endpoint=False)).tolist()
    colors = [[int(c[0] * 255), int(c[1] * 255), int(c[2] * 255)] for c in colors]
    for i, body in enumerate(bodies2D):
        color = list(reversed(colors[i]))
        img = utils_view.draw_body_in_image(img, body, joint_names, color)

    # Publish image with poses
    publish(img)

    msg = "Visualization time: {:.3f}s"
    print(msg.format(time.time() - ptime))


# ==================================================================================================


def publish(img):
    # Publish image data
    msg = bridge.cv2_to_imgmsg(img, "rgb8")
    publisher_img.publish(msg)


# ==================================================================================================


def main():
    global node, publisher_img

    # Start node
    rclpy.init(args=sys.argv)
    cam_id = sys.argv[1]
    node = rclpy.create_node("pose2d_visualizer")

    # Quality of service settings
    qos_profile = QoSProfile(
        reliability=QoSReliabilityPolicy.RELIABLE,
        history=QoSHistoryPolicy.KEEP_LAST,
        depth=1,
    )

    # Create subscribers
    _ = node.create_subscription(
        Image,
        img_input_topic.format(cam_id),
        callback_images,
        qos_profile,
    )
    _ = node.create_subscription(
        Poses,
        pose_input_topic.format(cam_id),
        callback_poses,
        qos_profile,
    )

    # Create publishers
    publisher_img = node.create_publisher(
        Image,
        img_output_topic.format(cam_id),
        qos_profile,
    )

    node.get_logger().info("Finished initialization of pose visualizer")

    # Run ros update thread
    rclpy.spin(node)

    node.destroy_node()
    rclpy.shutdown()


# ==================================================================================================

if __name__ == "__main__":
    main()
