#a node is the basic object that is saved in the tree

import matplotlib.pylab as plt
import numpy as np

#function to plot nodes on graph
def plot_nodes(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]
    
    
    #Create the figure
    plt.figure()
    plt.title("State of nodes generated by search algorithm")
    plt.ylabel("Z position")
    plt.xlabel("X position")
    
    #plot all nodes
    plt.scatter(x,z, c = 'C0', label = 'quad unvisited')
    plt.scatter(x_mass_a,z_mass_a, c = 'C1', label = 'mass unvisited')

    #plot visited nodes
    plt.scatter(x_visited,z_visited, c = 'C2', label = 'quad visited')
    plt.scatter(x_mass_visited,z_mass_visited, c = 'C3', label = 'mass visited')
    plt.legend(loc = 'best')

#basic object that stores information about the node
class Node:
    def __init__(self,parent_id,node_id, state, input_type, cost_to_come = 0, cost_to_go = 0):
        self.parent_id = parent_id
        self.node_id = node_id
        self.state = state
        self.input_type = input_type
        self.cost_to_come = cost_to_come
        self.cost_to_go = cost_to_go
        self.visited = False
        
    def __str__(self):
        return "ID = %s, P_ID = %s, State = %s, Cost2c = %s, Cost2g = %s, Input = %s" % (self.node_id,self.parent_id,self.state,self.cost_to_come,self.cost_to_go,self.input_type)
    def __repr__(self):
        return "ID = %s, P_ID = %s, State = %s, Cost2c = %s, Cost2g = %s, Input = %s" % (self.node_id,self.parent_id,self.state,self.cost_to_come,self.cost_to_go,self.input_type)
    
    