#!/usr/bin/env python

import math
import rospy
import geometry_msgs
import tf_conversions
from crazyflie_driver.msg import Position
from crazyflie_driver.msg import FullState
from std_msgs.msg import Empty
from geometry_msgs.msg import PoseStamped, Twist
from crazyflie_driver.msg import VelocityWorld, FullState
from crazyflie_init import crazyflie_init
from crazyflie_driver.srv import UpdateParams
from drone_model import drone_model_t 
from drone_control import drone_controller_t
import numpy as np

#noramalize vector
def normal(ax, ay, bx, by):
    result = math.sqrt((ax - bx) * (ax - bx) + (ay - by) * (ay - by))
    if result < 0.001:
        return 1 
    else:
        return result

#distance between two agents
def dist(ax, ay, bx, by):
    dist_square = (ax - bx) * (ax - bx) + (ay - by) * (ay - by)
    return math.sqrt(dist_square)

#takeoff
def takeoff(rospy, posi_pub, posi_msg, rate, x_init, y_init, height, frequency):
    while not rospy.is_shutdown():
        for i in range(int(0.3 * frequency)):
            posi_msg.x = x_init
            posi_msg.y = y_init
            posi_msg.yaw = 0.0
            posi_msg.z = i * 0.3 / (0.3 * frequency)
            posi_msg.header.seq += 1
            posi_msg.header.stamp = rospy.Time.now()
            posi_pub.publish(posi_msg)
            rate.sleep()
        for i in range(int(2 * frequency)):
            posi_msg.x = x_init
            posi_msg.y = y_init
            posi_msg.yaw = 0.0
            posi_msg.z = 0.3 + (height - 0.3) * i / (2 * frequency)
            posi_msg.header.seq += 1
            posi_msg.header.stamp = rospy.Time.now()
            posi_pub.publish(posi_msg)
            rate.sleep()
        break

#hover
def hover(rospy, posi_pub, posi_msg, rate, x_init, y_init, height, time):
    while not rospy.is_shutdown():
        for i in range(int(time)):
            posi_msg.x = x_init
            posi_msg.y = y_init
            posi_msg.yaw = 0.0
            posi_msg.z = height
            posi_msg.header.seq += 1
            posi_msg.header.stamp = rospy.Time.now()
            posi_pub.publish(posi_msg)
            rate.sleep()
        break

#land
def land(rospy, posi_pub, posi_msg, rate, x_end, y_end, z_end, frequency):
    while not rospy.is_shutdown():
        for i in range(int(2 * frequency)):
            posi_msg.x = x_end
            posi_msg.y = y_end
            posi_msg.yaw = 0.0
            posi_msg.z = z_end * (2 * frequency - i) / (2 * frequency)
            posi_msg.header.seq += 1
            posi_msg.header.stamp = rospy.Time.now()
            posi_pub.publish(posi_msg)
            rate.sleep()
        break

#position callback
def agent1_cb(pose):
    global x1
    global y1
    global z1
    #update the position of agent1    
    x1 = pose.pose.position.x
    y1 = pose.pose.position.y 
    z1 = pose.pose.position.z

def agent2_cb(pose):
    global x2
    global y2
    global z2
    #update the position of agent1    
    x2 = pose.pose.position.x
    y2 = pose.pose.position.y 
    z2 = pose.pose.position.z

def agent3_cb(pose):
    global x3
    global y3
    global z3
    #update the position of agent1    
    x3 = pose.pose.position.x
    y3 = pose.pose.position.y 
    z3 = pose.pose.position.z

def updatePIDparmeters(rospy):
    rospy.wait_for_service('update_params')
    rospy.loginfo("found update_params service")
    update_params = rospy.ServiceProxy('update_params', UpdateParams)
    xki = rospy.get_param("~xki", 0)
    yki = rospy.get_param("~yki", 0)
    rospy.set_param("posCtlPid/xKi", xki)
    rospy.set_param("posCtlPid/yKi", yki)
    update_params(["posCtlPid/xKi", "posCtlPid/yKi"])
    vxki = rospy.get_param("~vxki", 3)
    vyki = rospy.get_param("~vyki", 3)
    vxkp = rospy.get_param("~vxkp", 25)
    vykp = rospy.get_param("~vykp", 25)
    vxkd = rospy.get_param("~vxkd", 1)
    vykd = rospy.get_param("~vykd", 1)
    rospy.set_param("velCtlPid/vxKi", vxki)
    rospy.set_param("velCtlPid/vyKi", vyki)
    rospy.set_param("velCtlPid/vxKp", vxkp)
    rospy.set_param("velCtlPid/vyKp", vykp)
    rospy.set_param("velCtlPid/vxKd", vxkd)
    rospy.set_param("velCtlPid/vyKd", vykd)
    update_params(["velCtlPid/vxKi", "velCtlPid/vyKi", "velCtlPid/vxKp", "velCtlPid/vyKp", "velCtlPid/vxKd", "velCtlPid/vyKd"])

#initialize global variable
count_dist = 0
def distanceTopic(rospy, dist_pub, x1, y1, x2, y2, x3, y3, dp1x, dp1y, dp2x, dp2y, dp3x, dp3y):
    dist12 = dist(x1, y1, x2, y2)
    dist13 = dist(x1, y1, x3, y3)
    dist23 = dist(x2, y2, x3, y3)
    dist_des12 = dist(dp1x, dp1y, dp2x, dp2y)
    dist_des13 = dist(dp1x, dp1y, dp3x, dp3y)
    dist_des23 = dist(dp2x, dp2y, dp3x, dp3y)

    global count_dist
    dist_msg = PoseStamped()
    count_dist += 1
    dist_msg.header.seq = count_dist
    dist_msg.header.frame_id = "/world"
    dist_msg.header.stamp = rospy.Time.now()
    dist_msg.pose.position.x = dist12
    dist_msg.pose.position.y = dist13
    dist_msg.pose.position.z = dist23
    dist_msg.pose.orientation.x = dist_des12
    dist_msg.pose.orientation.y = dist_des13
    dist_msg.pose.orientation.z = dist_des23
    dist_pub.publish(dist_msg)

count_full_state = 0
def publish_full_state(rospy, full_state_pub, worldFrame, pos_cmd, vel_cmd):
    global count_full_state
    count_full_state += 1
    msg = FullState()
    msg.header.seq = count_full_state
    msg.header.stamp = rospy.Time.now()
    msg.header.frame_id = worldFrame
    msg.pose.position.x    = pos_cmd[0]
    msg.pose.position.y    = pos_cmd[1]
    msg.pose.position.z    = pos_cmd[2]
    msg.twist.linear.x     = vel_cmd[0]
    msg.twist.linear.y     = vel_cmd[1]
    msg.twist.linear.z     = vel_cmd[2]
    msg.acc.x              = 0.0
    msg.acc.y              = 0.0
    msg.acc.z              = 0.0
    msg.pose.orientation   = geometry_msgs.msg.Quaternion(*tf_conversions.transformations.quaternion_from_euler(0, 0, 0))
    msg.twist.angular.x    = 0.0
    msg.twist.angular.y    = 0.0
    msg.twist.angular.z    = 0.0
    full_state_pub.publish(msg)

#electrical fence
def now_in_elctricial_fence(x1, y1, x2, y2, x3, y3):
    a1_state = x1 < 1.5 and x1 > -1.5 and y1 < 1.5 and y1 > -1
    a2_state = x2 < 1.5 and x2 > -1.5 and y2 < 1.5 and y2 > -1
    a3_state = x3 < 1.5 and x3 > -1.5 and y3 < 1.5 and y3 > -1
    if (a1_state and a2_state and a3_state):
        return True

if __name__ == '__main__':
    print 'start'
    rospy.init_node('position', anonymous=True)
    
    if not crazyflie_init(rospy):
        rospy.logwarn(rospy.get_name()+" init failed !")
        exit(1)
    worldFrame = rospy.get_param("~worldFrame", "/world")
    #initialize position of agent 
    x1 = 0
    y1 = 0
    z1 = 0
    x2 = 0
    y2 = 0
    z2 = 0
    x3 = 0
    y3 = 0
    z3 = 0
    #initial value
    topic = rospy.get_param("~topic", "/crazyflie1/external_pose")
    pose = rospy.wait_for_message(topic, PoseStamped)
    x_init = pose.pose.position.x
    y_init = pose.pose.position.y 
    z_init = pose.pose.position.z
    #maximum veloctiy limit
    threh_max_vel = 1
    #permutation based on the number of agent
    if topic == "/crazyflie1/external_pose":
        topic_1 = "/crazyflie1/external_pose"
        topic_2 = "/crazyflie2/external_pose"
        topic_3 = "/crazyflie3/external_pose"
    elif topic == "/crazyflie2/external_pose":
        topic_1 = "/crazyflie2/external_pose"
        topic_2 = "/crazyflie3/external_pose"
        topic_3 = "/crazyflie1/external_pose"
    else:
        topic_1 = "/crazyflie3/external_pose"
        topic_2 = "/crazyflie1/external_pose"
        topic_3 = "/crazyflie2/external_pose"
    #define the position callback
    pose1 = rospy.Subscriber(topic_1, PoseStamped, agent1_cb)
    pose2 = rospy.Subscriber(topic_2, PoseStamped, agent2_cb)
    pose3 = rospy.Subscriber(topic_3, PoseStamped, agent3_cb)
    #update PID parameters
    updatePIDparmeters(rospy)
    #desired value 
    dp1x = rospy.get_param("~a1x")
    dp1y = rospy.get_param("~a1y")
    dp2x = rospy.get_param("~a2x")
    dp2y = rospy.get_param("~a2y")
    dp3x = rospy.get_param("~a3x")
    dp3y = rospy.get_param("~a3y")
    height = rospy.get_param("~height", 0.6)
    coef = rospy.get_param("~coef", 1.8)
    coef_integral = rospy.get_param("~coef_integral", 0.3)
    frequency = rospy.get_param("~frequency", 30)
    rate = rospy.Rate(frequency)
    #position command 
    posi_msg = Position()
    posi_msg.header.seq = 0
    posi_msg.header.stamp = rospy.Time.now()
    posi_msg.header.frame_id = worldFrame
    posi_msg.x = 0.0
    posi_msg.y = 0.0
    posi_msg.z = 0.0
    posi_msg.yaw = 0.0

    #stop command
    stop_msg = Empty()
    
    takeoff_flag = True

    if takeoff_flag:
        posi_pub = rospy.Publisher("cmd_position", Position, queue_size=1)
        #stop command
        stop_msg = Empty()
        stop_pub = rospy.Publisher("cmd_stop", Empty, queue_size=1)
        #full state command
        full_state_pub = rospy.Publisher("cmd_full_state", FullState, queue_size=1)
    else:
        posi_pub = rospy.Publisher("test_cmd_position", Position, queue_size=1)

        stop_pub = rospy.Publisher("test_cmd_stop", Empty, queue_size=1)
        #full state command
        full_state_pub = rospy.Publisher("test_cmd_full_state", FullState, queue_size=1)

    #publish three distances among agent1, agent2, agent3
    dist_pub = rospy.Publisher("info_dist", PoseStamped, queue_size=1)

    takeoff(rospy, posi_pub, posi_msg, rate, x_init, y_init, height, frequency)
    hover(rospy, posi_pub, posi_msg, rate, x_init, y_init, height, 3 * frequency)

    #init controller and model
    kp = 0.1
    kv = 0
    init_pos_group = [\
    [x1, y1],\
    [x2, y2],\
    [x3, y3]\
    ]

    start_vel = np.array([0,0])

    destination_pos_group = [\
    [dp1x, dp1y],\
    [dp2x, dp2y],\
    [dp3x, dp3y]\
    ]

    id_group = range(len(init_pos_group))

    drones = []
    for id in id_group:
        drone = drone_model_t(id, np.array(init_pos_group[id]), start_vel, destination_pos_group[id])
        drones.append(drone)

    controllers = []
    for id in id_group:
        controller = drone_controller_t(id, kp, kv)
        controllers.append(controller)

    #the time of begin
    begin_time = rospy.get_time()
    #convergence
    while not rospy.is_shutdown():
        now = rospy.get_time()
        t = now - begin_time
        if (t > 20.0):
            break 

        if False == now_in_elctricial_fence(x1, y1, x2, y2, x3, y3):
            break

        neis = [{}, {}, {}]
        for id in id_group:
            neis[id]['vel'] = np.array(drones[id].get_vel())
            neis[id]['des'] = np.array(drones[id].get_destination())
        
        neis[0]['pos'] = np.array([x1, y1])
        neis[1]['pos'] = np.array([x2, y2])
        neis[2]['pos'] = np.array([x3, y3])
        
        u = [[], [], []]
        u[0] = controllers[0].update(neis[0], [neis[1], neis[2]])
        u[1] = controllers[1].update(neis[1], [neis[0], neis[2]])
        u[2] = controllers[2].update(neis[2], [neis[0], neis[1]])

        for id in id_group:
            drones[id].update(u[id], rospy.get_time())

        pos_cmd = [drones[0].get_pos()[0], drones[0].get_pos()[1], height]
        # vel_cmd = [u[0][0], u[0][1], 0]
        vel_cmd = [0, 0, 0]
        publish_full_state(rospy, full_state_pub, worldFrame, pos_cmd, vel_cmd)
        distanceTopic(rospy, dist_pub, x1, y1, x2, y2, x3, y3, dp1x, dp1y, dp2x, dp2y, dp3x, dp3y)

    # control break
    land(rospy, posi_pub, posi_msg, rate, x1, y1, z1, frequency)    