'''
Just a file containing some of the plotting functions used by the notebook
'''

import numpy as np
import matplotlib.pylab as plt
import Node
#function to plot the nodes generated by the algorithm
def plot_nodes_astar(node_list, ymin = None, ymax = None):
    L = 0.5

    #Add limits if necessary
    if ymin != None:
        plt.ylim(ymin , ymax) 
    z = []
    x = []
    a = []
    visited = []
    
    #get z,x,alpha for all nodes
    for node in node_list:
        state = node.state
        visited.append(node.visited)
        z.append(state[0])
        x.append(state[2])
        a.append(state[4])
        
    #calculate pos of mas
    z_mass_a = z-L*np.cos(a)
    x_mass_a = x-L*np.sin(a)
    
    x_visited = np.array(x)[visited]
    z_visited = np.array(z)[visited]
    x_mass_visited = np.array(x_mass_a)[visited]
    z_mass_visited = np.array(z_mass_a)[visited]
    
    plt.title("Nodes Generated")
    plt.ylabel("Z Position")
    plt.xlabel("X Position")
    plt.axis('equal')
    #plot all nodes
    plt.scatter(x,z, c = 'C0', label = 'Nodes Not Visited', marker = '.')

    #plot visited nodes
    plt.scatter(x_visited,z_visited, c = 'C2', label = 'Nodes Visited', marker = '.')
    plt.legend(loc = 'best')
    
   #Function for plotting the position of the quadrotor woth suspended payload
def plot_position(result_a,t_a,positions):

 
    z_quad_a = result_a[0,:]
    x_quad_a = result_a[2,:]
    alpha_a = result_a[4,:]
    L = 0.5

    z_mass_a = z_quad_a-L*np.cos(alpha_a)
    x_mass_a = x_quad_a-L*np.sin(alpha_a)                 
    plt.title("Path Followed by Quadrotor")
    plt.plot(x_mass_a, z_mass_a, label = 'Mass', ls ='--', lw = 0.5)
    plt.plot(x_quad_a,z_quad_a, label = 'Quadrotor', ls ='--', color = 'tab:blue', lw = 0.5)           
    for p in positions:
        link = plt.Line2D((x_quad_a[p], x_mass_a[p]), (z_quad_a[p], z_mass_a[p]), lw=0.5, color = 'k')
        plt.gca().add_line(link)
        mass = plt.Circle((x_mass_a[p], z_mass_a[p]), 0.07, fc='k')
        plt.gca().add_patch(mass)

    plt.xlabel('X Position')
    plt.ylabel('Y Position')
    #plt.legend(loc = 'best')
    #plt.show()      

def plot_input_seq(ts, input_seq_a, U, path):
    '''
    This method plots a corresponding input sequence
    '''
    #plot the input sequence
    k = 0
    t = []
    vz = []
    vx = []
    for u in input_seq_a:
        t.append(k*Tsample)
        t.append((k+1)*Tsample)
        vz.append(u[0])
        vz.append(u[0])
        vx.append(u[1])
        vx.append(u[1])
        k += 1

    fig = plt.figure(figsize= [6, 3])
    ax = fig.add_subplot(111)
    fig.subplots_adjust(top=0.85)
    #ax.set_title('Desired Input Sequence')
    ax.set_title(r'Input Commands')
    ax.set_xlabel(r'Time $t_{s}$ (s)')
    ax.set_ylabel('Input')

    #ax.axis([-0.1, 3.5, -5.9,  5.9])


    t_pos = 0
    plt.plot([0,0],[-6,6], color = 'grey', ls = ':', lw = 0.5)
    for n in path:
        duration = U[n.input_type].shape[0]
        plt.plot([t_pos + ts*duration,t_pos + ts*duration],[-6,6], color = 'grey', ls = ':', lw = 0.5)
        t_pos += ts*duration

    plt.plot(t, vz, label = r"$\dot{z}_{ref}$")
    plt.plot(t, vx, label = r"$\dot{x}_{ref}$")
    plt.legend(loc = 'best')
    plt.tight_layout()
    plt.show()