#!/usr/bin/env python
import rospy
import tf
import tf2_ros
from std_msgs.msg import Empty
import geometry_msgs.msg
from math import radians as rad
from rospy import loginfo, logdebug, logwarn, logerr, logfatal
import sys
import time
from fiducial_msgs.msg import *

letani = True
takeoff = rospy.Publisher("/tello/takeoff", Empty, queue_size=1)
land = rospy.Publisher("/tello/land", Empty, queue_size=1)
cmd_vel = rospy.Publisher("/tello/cmd_vel", geometry_msgs.msg.Twist, queue_size=1)
fastmode = rospy.Publisher("/tello/fastmode", Empty, queue_size=1)

global marks
global vel_nul
global waypoint
global yaw


def marks_count(msg):
    # loginfo("Number of tags detected: " + str(len(msg.transforms)))
    global marks
    marks = len(msg.transforms)


def navigation(data, set_x, set_y, set_z, set_angle):
    # if data.transforms.child_frame_id == "tello":
    global marks
    global vel_nul
    global yaw
    yaw = 0.0
    cmd_vel2 = rospy.Publisher("/tello/cmd_vel", geometry_msgs.msg.Twist, queue_size=1)
    if marks < 1:
        if letani:
            vel = geometry_msgs.msg.Twist()
            vel.linear.x = 0.0
            vel.linear.y = 0.0
            vel.linear.z = 0.0
            if set_angle > yaw:
                vel.angular.z = -0.7
            else:
                vel.angular.z = 0.7
            cmd_vel2.publish(vel)
    else:
        try:
            zf = data.lookup_transform("map", "tello", rospy.Time())

            explicit_quat = [zf.transform.rotation.x, zf.transform.rotation.y, zf.transform.rotation.z,
                             zf.transform.rotation.w]
            roll, pitch, yaw = tf.transformations.euler_from_quaternion(explicit_quat)
            if yaw < 0:
                yaw2 = yaw + 3.14
            if yaw > 0:
                yaw2 = yaw - 3.14
            Kp = 1.0
            if abs(set_x - zf.transform.translation.x) > 0.02:
                final_x = Kp * (set_x - zf.transform.translation.x)
            else:
                final_x = 0.0

            if abs(set_y - zf.transform.translation.y) > 0.02:
                final_y = Kp * (set_y - zf.transform.translation.y)
            else:
                final_y = 0.0

            if abs(set_z - zf.transform.translation.z) > 0.02:
                final_z = Kp * 2.0 * (set_z - zf.transform.translation.z)
            else:
                final_z = 0.0

            if abs(set_angle - yaw2) > 0.03:
                if set_angle > yaw2:
                    final_rx = -0.5
                else:
                    final_rx = 0.5
                # final_rx = Kp * (set_angle - yaw)
                #if final_rx > 3.14 or abs(set_angle - yaw) > 3.14:
                #    final_rx = 0.0
            else:
                final_rx = 0.0

            if letani:
                vel = geometry_msgs.msg.Twist()
                if -0.3925 < yaw < 0.3925:
                    vel.linear.x = final_x
                    vel.linear.y = final_y
                    vel.linear.z = final_z
                    vel.angular.z = final_rx

                if -1.9625 < yaw < -1.1775:
                    vel.linear.x = - final_y
                    vel.linear.y = final_x
                    vel.linear.z = final_z
                    vel.angular.z = final_rx

                if 1.1775 < yaw < 1.9625:
                    vel.linear.x = final_y
                    vel.linear.y = - final_x
                    vel.linear.z = final_z
                    vel.angular.z = final_rx

                if yaw > 2.7475 or yaw < -2.7475:
                    vel.linear.x = - final_x
                    vel.linear.y = - final_y
                    vel.linear.z = final_z
                    vel.angular.z = final_rx

                if 0.3925 < yaw < 1.1775:
                    vel.linear.x = 0.5 * (final_x + final_y)
                    vel.linear.y = 0.5 * (final_y - final_x)
                    vel.linear.z = final_z
                    vel.angular.z = final_rx

                if 1.9625 < yaw < 2.7475:
                    vel.linear.x = 0.5 * (final_y - final_x)
                    vel.linear.y = 0.5 * (- final_x - final_y)
                    vel.linear.z = final_z
                    vel.angular.z = final_rx

                if -2.7475 < yaw < -1.9625:
                    vel.linear.x = 0.5 * (- final_y - final_x)
                    vel.linear.y = 0.5 * (final_x - final_y)
                    vel.linear.z = final_z
                    vel.angular.z = final_rx

                if -1.1775 < yaw < -0.3925:
                    vel.linear.x = 0.5 * (final_x - final_y)
                    vel.linear.y = 0.5 * (final_y + final_x)
                    vel.linear.z = final_z
                    vel.angular.z = final_rx

                cmd_vel2.publish(vel)

            loginfo("x: " + str(zf.transform.translation.x) + " vel: " + str(vel.linear.x))
            loginfo("y: " + str(zf.transform.translation.y) + " vel: " + str(vel.linear.y))
            loginfo("z: " + str(zf.transform.translation.z) + " vel: " + str(vel.linear.z))
            loginfo("angle: " + str(yaw2) + " vel: " + str(vel.angular.z))
        except:
            pass


def set_waypoint(data):
    global waypoint
    waypoint = data


def main():
    """
    Main
    """

    # Init node
    loginfo("init")
    rospy.init_node('navigation', anonymous=False)

    global vel_nul
    global waypoint
    global marks

    # Subscibers setup
    rospy.Subscriber('fiducial_transforms', FiducialTransformArray, marks_count)

    tfBuffer = tf2_ros.Buffer()
    listener = tf2_ros.TransformListener(tfBuffer)

    vel_nul = geometry_msgs.msg.Twist()
    waypoint = geometry_msgs.msg.Twist()
    rate = rospy.Rate(30)  # Hz
    # Infinite localization loop
    connect_timeout = 0
    loginfo("Connecting...")
    while cmd_vel.get_num_connections() == 0 or takeoff.get_num_connections() == 0 or land.get_num_connections() == 0:
        rospy.loginfo("Waiting for subscriber to connect")
        connect_timeout = connect_timeout + 1
        if connect_timeout > 20:
            sys.exit("ERROR: Dron is not connected or tello driver is not running")
        rospy.sleep(1)

    if letani:
        loginfo("Taking off")
        takeoff.publish(Empty())

        fastmode.publish(Empty())
        time.sleep(4)

        # Initial speed is set to 0
        vel_nul.linear.x = 0.0  # doprava
        vel_nul.linear.y = 0.0  # vpred
        vel_nul.linear.z = 0.0  # nahoru
        vel_nul.angular.x = 0.0  # nic
        vel_nul.angular.y = 0.0  # nic
        vel_nul.angular.z = 0.0  # doprava rotace
        cmd_vel.publish(vel_nul)

        # Default waypoint after taking off
        waypoint.linear.x = 0.0
        waypoint.linear.y = 1.5
        waypoint.linear.z = 1.5
        waypoint.angular.z = 0.0 # 0.55

    lf = 0

    while not rospy.is_shutdown():
        if letani:
            # Looking for marks
            if marks < 1:
                loginfo("Looking for marks. Timeout: " + str(20 - lf) + " s")
                cmd_vel.publish(vel_nul)
                if lf > 1:
                    navigation(tfBuffer, 0.0, 0.0, 0.0, 0.6)
                time.sleep(1)
                lf += 1
                if lf > 20:
                    loginfo("No mark found! Landing.")
                    cmd_vel.publish(vel_nul)
                    land.publish(Empty())
                    time.sleep(4)
                    sys.exit("Exit")
            # Subscriber and calling of regulator loop
            else:
                loginfo("Marks found. Current waypoint: x:" + str(waypoint.linear.x) + ", y: " + str(waypoint.linear.y) + ", z: " + str(waypoint.linear.z) + ", angle: " + str(waypoint.angular.z))
                lf = 0
                rospy.Subscriber("tello/command", geometry_msgs.msg.Twist, set_waypoint)
                if waypoint.linear.x == 0.0 and waypoint.linear.y == 0.0 and waypoint.linear.z == 0.0:
                    cmd_vel.publish(vel_nul)
                    land.publish(Empty())
                    time.sleep(4)
                    sys.exit("Exit")
                else:
                    navigation(tfBuffer, waypoint.linear.x, waypoint.linear.y, waypoint.linear.z, waypoint.angular.z)
        rate.sleep()


if __name__ == '__main__':
    main()

