from scipy import linalg as la
import matplotlib.pyplot as pl
import numpy as np

import quadrotor as quad
import formation_distance as form
import quadlog
import animation as ani



#def listen_rel_xyz(quad1: quadrotor, quad2; quadrotor)



# Quadrotor
m = 0.65  # Kg
l = 0.23  # m
Jxx = 7.5e-3  # Kg/m^2
Jyy = Jxx
Jzz = 1.3e-2
Jxy = 0
Jxz = 0
Jyz = 0
J = np.array([[Jxx, Jxy, Jxz],
              [Jxy, Jyy, Jyz],
              [Jxz, Jyz, Jzz]])
CDl = 9e-3
CDr = 9e-4
kt = 3.13e-5  # Ns^2
km = 7.5e-7   # Ns^2
kw = 1/0.18   # rad/s

# Initial conditions
att_0 = np.array([0.0, 0.0, 0.0])
pqr_0 = np.array([0.0, 0.0, 0.0])
xyzt_0 = np.array([0.0, 0.0, 0.0])
xyz1_0 = np.array([1.0, 1.2, -0.0])
xyz2_0 = np.array([1.2, 2.0, -0.0])
xyz3_0 = np.array([-1.1, 2.6, -0.0])
v_ned_0 = np.array([0.0, 0.0, 0.0])
w_0 = np.array([0.0, 0.0, 0.0, 0.0])

# Setting quads
qt = quad.quadrotor(0, m, l, J, CDl, CDr, kt, km, kw,
                    att_0, pqr_0, xyzt_0, v_ned_0, w_0)

q1 = quad.quadrotor(1, m, l, J, CDl, CDr, kt, km, kw,
                    att_0, pqr_0, xyz1_0, v_ned_0, w_0)

q2 = quad.quadrotor(2, m, l, J, CDl, CDr, kt, km, kw,
                    att_0, pqr_0, xyz2_0, v_ned_0, w_0)

q3 = quad.quadrotor(3, m, l, J, CDl, CDr, kt, km, kw,
                    att_0, pqr_0, xyz3_0, v_ned_0, w_0)


# Simulation parameters
tf = 400
dt = 5e-2
time = np.linspace(0, tf, tf/dt)
it = 0
frames = 100

# Data log
qt_log = quadlog.quadlog(time)
q1_log = quadlog.quadlog(time)
q2_log = quadlog.quadlog(time)
q3_log = quadlog.quadlog(time)

# Plots
quadcolor = ['m', 'r', 'g', 'b']
pl.close("all")
pl.ion()
fig = pl.figure(0)
axis3d = fig.add_subplot(111, projection='3d')

init_area = 5
s = 2

# Desired altitude and heading
alt_d = 2
qt.yaw_d = 0
q1.yaw_d = 0
q2.yaw_d = 0
q3.yaw_d = 0


def impPath3(qt, q1, q2, q3, r, a1, a2, a3):
    q1xyz = np.array([qt.xyz[0], qt.xyz[1] - r])
    q1xyz = q1xyz - np.array([q1.xyz[0], q1.xyz[1]])
    q2xyz = np.array([np.cos(a1-90)*r+qt.xyz[0], np.sin(a1-90)*r+qt.xyz[1]])
    q3xyz = np.array([np.cos((a1+a2)-90)*r+qt.xyz[0], np.sin((a1+a2)-90)*r+qt.xyz[1]])
    print("new: ", q3xyz)
    return (q1xyz, q2xyz, q3xyz)

def impPath2(qt, q1, q2, q3, r):
    q1xyz = np.array([qt.xyz[0], qt.xyz[1] - r])
    q1xyz = q1xyz - np.array([q1.xyz[0], q1.xyz[1]])
    q2xyz = np.array([qt.xyz[0] - 1.4142, qt.xyz[1] + 1.4142])
    q2xyz = q2xyz - np.array([q2.xyz[0], q2.xyz[1]])
    q3xyz = np.array([qt.xyz[0] + 1.4142, qt.xyz[1] + 1.4142])
    q3xyz = q3xyz - np.array([q3.xyz[0], q3.xyz[1]])
    print("old: ", q3xyz)
    return (q1xyz, q2xyz, q3xyz)


for t in time:

    print(q1.broadcast_rel_xyz(qt))

    # Simulation
    X = np.append(q1.xyz[0:2], np.append(q2.xyz[0:2], q3.xyz[0:2]))
    V = np.append(q1.v_ned[0:2], np.append(q2.v_ned[0:2], q3.v_ned[0:2]))

    gradF1 = np.array([X[0]*2, X[1]*2])
    gradF2 = np.array([X[2]*2, X[3]*2])
    gradF3 = np.array([X[4]*2, X[5]*2])

    tangentHelp = np.array([[0, -1], [1, 0]])
    tangent1 = tangentHelp.dot((gradF1/la.norm(gradF1)))
    tangent2 = tangentHelp.dot((gradF2/la.norm(gradF2)))
    tangent3 = tangentHelp.dot((gradF3/la.norm(gradF3)))

    #e1 = impPath(X[0], X[1], 2)
    #e2 = impPath(X[2], X[3], 2)
    #e3 = impPath(X[4], X[5], 2)
    ke = 0.012  # Gain for going towards circle
    kt = 0.035  # Gain for velocity
    kv = 0.02   # QT speed
    kvt = [kv,kv]
    kp = 0.1    # Proportional of error

    # This small u is our circle following control input (velocity)
    #u1 = ke*(-e1)*gradF1 + kt*tangent1
    #u2 = ke*(-e2)*gradF2 + kt*tangent2
    #u3 = ke*(-e3)*gradF3 + kt*tangent3
    u1, u2, u3 = impPath2(qt,q1,q2,q3,2)
    #u1, u2, u3 = impPath3(qt,q1,q2,q3,2,120,120,120)
    u1 = u1 * kp
    u2 = u2 * kp
    u3 = u3 * kp



    qt.set_v_2D_alt_lya(kvt, -alt_d)
    q1.set_v_2D_alt_lya(u1, -alt_d)
    q2.set_v_2D_alt_lya(u2, -alt_d)
    q3.set_v_2D_alt_lya(u3, -alt_d)

    qt.step(dt)
    q1.step(dt)
    q2.step(dt)
    q3.step(dt)

    # Animation
    if it % frames == 0:

        pl.figure(0)
        axis3d.cla()
        ani.draw3d(axis3d, q1.xyz, q1.Rot_bn(), quadcolor[0])
        ani.draw3d(axis3d, q2.xyz, q2.Rot_bn(), quadcolor[1])
        ani.draw3d(axis3d, q3.xyz, q3.Rot_bn(), quadcolor[2])
        ani.draw3d(axis3d, qt.xyz, qt.Rot_bn(), quadcolor[3])
        axis3d.set_xlim(-5, 5)
        axis3d.set_ylim(-5, 5)
        axis3d.set_zlim(0, 10)
        axis3d.set_xlabel('South [m]')
        axis3d.set_ylabel('East [m]')
        axis3d.set_zlabel('Up [m]')
        axis3d.set_title("Time %.3f s" % t)
        pl.pause(0.001)
        pl.draw()
        # namepic = '%i'%it
        # digits = len(str(it))
        # for j in range(0, 5-digits):
        #    namepic = '0' + namepic
        # pl.savefig("./images/%s.png"%namepic)


    # Log
    q1_log.xyz_h[it, :] = q1.xyz
    q1_log.att_h[it, :] = q1.att
    q1_log.w_h[it, :] = q1.w
    q1_log.v_ned_h[it, :] = q1.v_ned

    q2_log.xyz_h[it, :] = q2.xyz
    q2_log.att_h[it, :] = q2.att
    q2_log.w_h[it, :] = q2.w
    q2_log.v_ned_h[it, :] = q2.v_ned

    q3_log.xyz_h[it, :] = q3.xyz
    q3_log.att_h[it, :] = q3.att
    q3_log.w_h[it, :] = q3.w
    q3_log.v_ned_h[it, :] = q3.v_ned

    qt_log.xyz_h[it, :] = qt.xyz
    qt_log.att_h[it, :] = qt.att
    qt_log.w_h[it, :] = qt.w
    qt_log.v_ned_h[it, :] = qt.v_ned

    it += 1

    # Stop if crash
    if (qt.crashed == 1 or q1.crashed == 1 or q2.crashed == 1 or q3.crashed == 1):
        break

pl.figure(1)
pl.title("2D Position [m]")
pl.plot(q1_log.xyz_h[:, 0], q1_log.xyz_h[:, 1], label="q1", color=quadcolor[0])
pl.plot(q2_log.xyz_h[:, 0], q2_log.xyz_h[:, 1], label="q2", color=quadcolor[1])
pl.plot(q3_log.xyz_h[:, 0], q3_log.xyz_h[:, 1], label="q3", color=quadcolor[2])
pl.plot(qt_log.xyz_h[:, 0], qt_log.xyz_h[:, 1], label="qt", color=quadcolor[3])
pl.xlabel("East")
pl.ylabel("South")
pl.legend()

pl.figure(2)
pl.plot(time, q1_log.att_h[:, 2], label="yaw q1")
pl.plot(time, q2_log.att_h[:, 2], label="yaw q2")
pl.plot(time, q3_log.att_h[:, 2], label="yaw q3")
pl.plot(time, qt_log.att_h[:, 2], label="yaw qt")
pl.xlabel("Time [s]")
pl.ylabel("Yaw [rad]")
pl.grid()
pl.legend()

pl.figure(3)
pl.plot(time, -q1_log.xyz_h[:, 2], label="$q_1$")
pl.plot(time, -q2_log.xyz_h[:, 2], label="$q_2$")
pl.plot(time, -q3_log.xyz_h[:, 2], label="$q_3$")
pl.plot(time, -qt_log.xyz_h[:, 2], label="$q_t$")
pl.xlabel("Time [s]")
pl.ylabel("Altitude [m]")
pl.grid()
pl.legend(loc=2)


pl.pause(0)