#!python

import logging

logger = logging.getLogger(f"rosout.{__name__}")
import sys
import asyncio
import time
import base64

from officebots import Robot

officebots_logger = logging.getLogger("officebots")

from math import pi

import cv2
import numpy as np
from cv_bridge import CvBridge

import rospy
from geometry_msgs.msg import Twist, PoseStamped
from nav_msgs.msg import Odometry
from sensor_msgs.msg import LaserScan, Image
from hri_msgs.msg import IdsList, LiveSpeech
from std_msgs.msg import String, Float32

import random

import tf

OPTICAL_FRAME_ROTATION = tf.transformations.quaternion_from_euler(-pi / 2, 0, -pi / 2)

argv = rospy.myargv(argv=sys.argv)

if len(argv) < 2:
    print("Usage: %s <robot name>" % argv[0])
    sys.exit(1)

robot_name = argv[1]


def random_id():
    return "".join(random.choices("abcdefghijklmnopqrstuvxyz", k=5))


def restore_logging():
    handler = logging.StreamHandler()
    handler.setFormatter(
        logging.Formatter("[%(levelname)s] %(asctime)s - %(name)s: %(message)s")
    )
    officebots_logger.addHandler(handler)
    officebots_logger.setLevel(logging.INFO)


class RosOfficeBots(Robot):

    base_frame = "base_link"
    sellion_frame = "sellion_link"

    def __init__(self):
        super().__init__()
        rospy.init_node("officebots")
        restore_logging()
        rospy.Subscriber("cmd_vel", Twist, self.on_cmd_vel)
        rospy.Subscriber("move_base_simple/goal", PoseStamped, self.on_nav_goal)
        rospy.Subscriber("screen/raw_image", Image, self.on_screen_image)
        rospy.Subscriber("say", String, self.on_say)

        self.supports_camera = True  # will be automatically set to false if the game reports no camera support

        self.cv_bridge = CvBridge()
        self.br = tf.TransformBroadcaster()

        self.odom_pub = rospy.Publisher("odom", Odometry, queue_size=1)
        self.odom_msg = Odometry()
        self.odom_msg.header.frame_id = "odom"
        self.odom_msg.child_frame_id = self.base_frame
        self.odom_msg.pose.pose.position.z = 0.0
        self.odom_msg.pose.pose.orientation.x = 0.0
        self.odom_msg.pose.pose.orientation.y = 0.0

        # laser scan parameters

        nb_rays = 50
        self.angle_increment = pi / (nb_rays - 1)
        self.angle_min = -pi / 2
        self.angle_max = pi / 2
        self.range_min = 0.0  # m
        self.range_max = 20.0  # m

        self.scan_pub = rospy.Publisher("scan", LaserScan, queue_size=1)
        self.scan_msg = LaserScan()
        self.scan_msg.header.frame_id = self.base_frame
        self.scan_msg.angle_min = self.angle_min
        self.scan_msg.angle_max = self.angle_max
        self.scan_msg.angle_increment = self.angle_increment
        self.scan_msg.time_increment = 0.0
        self.scan_msg.range_min = self.range_min
        self.scan_msg.range_max = self.range_max

        self.image_pub = rospy.Publisher("camera/rgb/image_raw", Image, queue_size=1)

        self.last = time.time()

        self.humans = {}
        self.face_tracked_pub = rospy.Publisher(
            f"/humans/faces/tracked", IdsList, queue_size=1
        )
        self.body_tracked_pub = rospy.Publisher(
            f"/humans/bodies/tracked", IdsList, queue_size=1
        )
        self.voices_tracked_pub = rospy.Publisher(
            f"/humans/voices/tracked", IdsList, queue_size=1
        )
        self.person_tracked_pub = rospy.Publisher(
            f"/humans/persons/tracked", IdsList, queue_size=1
        )
        self.person_known_pub = rospy.Publisher(
            f"/humans/persons/known", IdsList, queue_size=1
        )

        self.voices = {}

    def safe_name(self, name):
        for c in " -,.'\";+_":
            name = name.replace(c, "")
        return name.lower()

    def add_human(self, name):
        # person_id = "person_" + random_id()
        person_id = "person_" + self.safe_name(name)

        face_id_pub = rospy.Publisher(
            f"/humans/persons/{person_id}/face_id", String, queue_size=1
        )
        body_id_pub = rospy.Publisher(
            f"/humans/persons/{person_id}/body_id", String, queue_size=1
        )
        voice_id_pub = rospy.Publisher(
            f"/humans/persons/{person_id}/voice_id", String, queue_size=1
        )

        name_pub = rospy.Publisher(
            f"/humans/persons/{person_id}/name", String, queue_size=1, latch=True
        )

        self.humans[name] = {
            "id": person_id,
            "visible": False,
            "face_id_pub": face_id_pub,
            "body_id_pub": body_id_pub,
            "voice_id_pub": voice_id_pub,
            "name_pub": name_pub,
        }

        return person_id

    def add_voice(self):
        voice_id = random_id()

        speech_pub = rospy.Publisher(
            f"/humans/voices/{voice_id}/speech", LiveSpeech, queue_size=1
        )

        self.voices[voice_id] = {"speech_pub": speech_pub}

        return voice_id

    def remove_human(self, name):
        if name not in self.humans:
            return

        p = self.humans[name]

        # unregister all ROS publishers
        for _, v in p:
            if hasattr(v, "unregister"):
                v.unregister()

        del self.humans[name]

    def update_humans(self, seen_humans):

        seen = {p[0]: p[1:] for p in seen_humans}

        for name in seen.keys():
            if name not in self.humans:
                self.add_human(name)

            # the human is either new or just re-appeared
            if not self.humans[name]["visible"]:
                self.humans[name]["visible"] = True
                face_id = random_id()
                self.humans[name]["face"] = face_id

            self.humans[name]["pos"] = seen[name][0:3]
            self.humans[name]["face_transform"] = seen[name][3:9]

        not_visible_anymore = []

        for name in self.humans.keys():
            if name not in seen.keys():
                not_visible_anymore.append(name)

        for name in not_visible_anymore:
            self.humans[name]["visible"] = False
            self.humans[name]["face"] = ""
            self.humans[name]["gaze"] = ""
            self.humans[name]["body"] = ""

    def publish_humans(self):

        faces_list = IdsList()
        bodies_list = IdsList()
        voices_list = IdsList()
        tracked_persons_list = IdsList()
        known_persons_list = IdsList()

        for name, p in self.humans.items():

            known_persons_list.ids.append(p["id"])

            if "pos" in p:
                x, y, theta = p["pos"]

                # for 'known' persons (even if not currently visible), always
                # publish the map <-> person transform
                qx, qy, qz, qw = tf.transformations.quaternion_from_euler(0, 0, theta)
                self.br.sendTransform(
                    (x, y, 0),
                    (qx, qy, qz, qw),
                    rospy.Time.now(),
                    p["id"],
                    "map",
                )

            if p["visible"]:
                tracked_persons_list.ids.append(p["id"])

            msg = String()
            if "face" in p and p["face"] != "":
                faces_list.ids.append(p["face"])
                msg.data = p["face"]

                x, y, z, rx, ry, rz = p["face_transform"]
                # publish the 'face' and 'gaze' tf frames
                qx, qy, qz, qw = tf.transformations.quaternion_from_euler(rx, ry, rz)
                self.br.sendTransform(
                    (x, y, z),
                    (qx, qy, qz, qw),
                    rospy.Time.now(),
                    "face_" + p["face"],
                    "map",
                )
                # gaze follows the optical frame convention: z-forward
                self.br.sendTransform(
                    (0, 0, 0),
                    OPTICAL_FRAME_ROTATION,
                    rospy.Time.now(),
                    "gaze_" + p["face"],
                    "face_" + p["face"],
                )
            # else, no face -> publish empty face id
            p["face_id_pub"].publish(msg)

            msg = String()
            if "body" in p and p["body"] != "":
                bodies_list.ids.append(p["body"])
                msg.data = p["body"]
            # else, no body -> publish empty body id
            p["body_id_pub"].publish(msg)

            msg = String()
            if "voice" in p and p["voice"] != "":
                voices_list.ids.append(p["voice"])
                msg.data = p["voice"]
            # else, no voice -> publish empty voice id
            p["voice_id_pub"].publish(msg)

            msg = String()
            msg.data = name
            p["name_pub"].publish(msg)

        self.face_tracked_pub.publish(faces_list)
        self.body_tracked_pub.publish(bodies_list)
        self.voices_tracked_pub.publish(voices_list)
        self.person_tracked_pub.publish(tracked_persons_list)
        self.person_known_pub.publish(known_persons_list)

    async def run(self):
        res = await self.execute([robot_name, "create"])
        if res[0] != self.OK:
            logger.warning(res[1])
        else:
            logger.info(f"Created the robot {robot_name}")

        logger.info("ROS Officebots bridge ready")

        # simply keeps python from exiting until this node is stopped
        while not rospy.is_shutdown():
            if self.supports_camera:
                status, jpg_data = await self.execute([robot_name, "export-camera"])
                if status == Robot.OK:
                    # await asyncio.sleep(0.1)
                    array = np.asarray(bytearray(jpg_data), dtype=np.uint8)
                    image = cv2.imdecode(array, cv2.IMREAD_COLOR)
                    image_msg = self.cv_bridge.cv2_to_imgmsg(
                        image, encoding="passthrough"
                    )
                    rospy.logwarn("publishing image")
                    self.image_pub.publish(image_msg)
                else:
                    if "Unknown command" in jpg_data:
                        self.supports_camera = False

            status, humans = await self.execute([robot_name, "get-humans"])
            if status == Robot.OK:
                self.update_humans(humans)

        self.stop()
        logger.info("Bye")

    async def on_robot_update(self, data):

        now = time.time()
        dt = now - self.last
        self.last = now

        x, y, theta, v, w = data["odom"]

        qx, qy, qz, qw = tf.transformations.quaternion_from_euler(0, 0, theta)

        self.odom_msg.header.stamp = rospy.Time.now()
        self.odom_msg.pose.pose.position.x = x
        self.odom_msg.pose.pose.position.y = y

        self.odom_msg.pose.pose.orientation.z = qz
        self.odom_msg.pose.pose.orientation.w = qw

        self.odom_msg.twist.twist.linear.x = v
        self.odom_msg.twist.twist.angular.z = w

        self.odom_pub.publish(self.odom_msg)

        self.br.sendTransform((0, 0, 0), (0, 0, 0, 1), rospy.Time.now(), "odom", "map")

        self.br.sendTransform(
            (x, y, 0), (qx, qy, qz, qw), rospy.Time.now(), self.base_frame, "odom"
        )
        self.br.sendTransform(
            (0, 0, 1.6),
            (0, 0, 0, 1),
            rospy.Time.now(),
            self.sellion_frame,
            self.base_frame,
        )

        self.scan_msg.scan_time = dt
        self.scan_msg.ranges = data["laserscan"]
        self.scan_msg.header.stamp = rospy.Time.now()

        self.scan_pub.publish(self.scan_msg)

        speech = data["speech_heard"]
        if speech:
            text, name = speech

            if name not in self.humans:
                self.add_human(name)

            if "voice" not in self.humans[name] or (not self.humans[name]["voice"]):

                voice_id = self.add_voice()
                self.humans[name]["voice"] = voice_id

            voice_id = self.humans[name]["voice"]
            speech_msg = LiveSpeech()
            speech_msg.final = text
            speech_msg.confidence = 1.0
            self.voices[voice_id]["speech_pub"].publish(speech_msg)

        self.publish_humans()

    def on_cmd_vel(self, twist):
        x = twist.linear.x
        w = twist.angular.z

        self._event_loop.create_task(self.execute([robot_name, "cmd-vel", [x, w]]))

    def on_nav_goal(self, pose):
        if pose.header.frame_id != "map":
            rospy.logwarn(
                "The navigation goal PoseStamped must currently be set in the 'map' reference frame"
            )
            return

        x = pose.pose.position.x
        y = pose.pose.position.y

        self._event_loop.create_task(self.execute([robot_name, "navigate-to", [x, y]]))

    def on_screen_image(self, image):

        cv_img = self.cv_bridge.imgmsg_to_cv2(image, desired_encoding="passthrough")

        cv_img = cv2.resize(cv_img, (320, 240))
        _, buf = cv2.imencode(".jpg", cv_img)
        b64 = base64.b64encode(buf).decode("ascii")
        rospy.logdebug("Sending image size: %skb" % (len(b64) / 1024))

        self._event_loop.create_task(self.execute([robot_name, "set-screen", [b64]]))

    def on_say(self, msg):

        self._event_loop.create_task(self.execute([robot_name, "say", [msg.data]]))


if __name__ == "__main__":

    RosOfficeBots().start()
