Module tsp.tsp_plot

"""
    TSP - plotting routines.
"""
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

__fig = None
__ax = None
__figNum = 0

__canvas = None
__tk_window = None

#---------------------------------------------------------------------

def __plotEdge( A, B, special=False ):
    """
    Plot line connecting coordinate points A and B.
    
    Parameters:
        A : list of float
            Coordinates of point A
        B : list of float
            Coordinates of point B
        special : bool, optional
            If true, a red line otherwise a blue line is drawn.
            Default=False.
    """
    col = 'r-' if special else 'b-'
    plt.plot( [A[0],B[0]], [A[1],B[1]], col )
    
#---------------------------------------------------------------------

def __initFig( ):
    global __fig
    global __ax
    global __figNum
    global __canvas

    if plt.isinteractive() or (not __tk_window is None):
        if __fig is None:
            __fig, __ax = plt.subplots()
        else:
            plt.clf()
    else:
        __figNum += 1
        __fig = plt.figure(__figNum)
        __ax  = __fig.subplots()
    
    if not __tk_window is None and __canvas is None:
        from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
        import tkinter as tk
        __canvas = FigureCanvasTkAgg(__fig, master=__tk_window)  
        __canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=1)
    
#---------------------------------------------------------------------

def set_tk_window( root ):
    global __tk_window
    
    __tk_window = root
    
#---------------------------------------------------------------------

def __setFigLims( X, Y ):
    """
    Set bounding box for a figure.
    
    Parameters:
        ax : object (matplotlib)
            axis object.
        X  : list of X-coordinates.
        Y  : list of Y-coordinates.   
    """
    global __ax
    xmin, ymin = min(X), min(Y)
    xmax, ymax = max(X), max(Y)
    xsep, ysep = 0.1*(xmax-xmin), 0.1*(ymax-ymin)
    __ax.set_xlim(xmin-xsep, xmax+xsep)
    __ax.set_ylim(ymin-ysep, ymax+ysep)

#---------------------------------------------------------------------

def displayPoints( problem, fontsize=6, wait_func=None ):
    """
    Display nodes of the TSP problem  
    
    Parameters:
        **problem** : object  
            TSP problem object as returned by function load_problem
            of the package tsplib95.  
        **fontsize**: int  
            size of the font used to display node numbers  
        **wait_func**: function, optional  
            function to be called (without any arguments)
            after the drawing was made. Usually used when
            matplotlib.pyplot is not in interactive mode
            so that execution halts before proceeding.
            Default: None    
    """
    global __fig
    global __ax
    global __canvas
    
    no_coords = problem.node_coords is None or len(problem.node_coords)==0
    dis_data = problem.display_data if no_coords else problem.node_coords
    X = [ dis_data[i][0] for i in problem.get_nodes()]
    Y = [ dis_data[i][1] for i in problem.get_nodes()]
    
    __initFig()
   
    __setFigLims( X, Y )
    __ax.plot(X,Y,'bo',markersize=2)
    for i in problem.get_nodes():
        __ax.text( dis_data[i][0],dis_data[i][1],str(i),fontsize=fontsize)
    plt.title("Node locations")
    if __canvas is None:
        plt.show()
    else:
        __canvas.draw()
    
    if not wait_func is None: wait_func()
            
#---------------------------------------------------------------------

def displayRoute( problem, route, routeLen, title=None, animate=True,
                  wait_func = None, fontsize=6 ):
    """ 
    Display (animated) route over nodes.
    
    Parameters:
        **coords** : list of coordinates  
            Node coordinates (attribute problem.node_coords of class
            problem in the package tsplib95).  
        **route** : list of int  
            The route to be displayed.  
        **routeLen** : int  
            The length of the route.  
        **title** : str  
            String used as a header for the figure.  
        **animate** : bool, optional  
            If true, the default, the route is displayed in animated way.        
        **wait_func** : function, optional  
            function to be called (without any arguments)
            after the drawing was made. Usually used when
            matplotlib.pyplot is not in interactive mode
            so that execution halts before proceeding.  
            Default: None    
        **fontsize** : int  
            Size of the font used to display node numbers          
    """
    global __fig
    global __ax
    global __canvas
 
    no_coords = problem.node_coords is None or len(problem.node_coords)==0
    coords = problem.display_data if no_coords else problem.node_coords
    X = [ coords[i][0] for i in route]
    Y = [ coords[i][1] for i in route] 
    X.append( X[0] )
    Y.append( Y[0] )
  
    __initFig()  
        
    ln, = plt.plot([], [], 'b-', markersize=2,animated=True)
    plt.plot(X,Y,'bo',markersize=2)
    if title != None: plt.title(title+"(Length="+str(routeLen)+")")
    for i in range(len(route)):
        plt.text(X[i],Y[i],str(route[i]),fontsize=fontsize)

    def init():
        __setFigLims( X, Y )
        return ln,

    def update(j):
        ln.set_data(X[0:j], Y[0:j])
        return ln,
   
    if animate:
        anim = FuncAnimation(__fig, update, frames=range(1,len(route)+2),\
                   init_func=init, blit=True, repeat=False)
    else:
        __setFigLims( X, Y )
        plt.plot( X, Y, 'b-')
    
    if __canvas is None:    
        plt.show()
    else:
        __canvas.draw() 

    if not wait_func is None: wait_func()
    
#---------------------------------------------------------------------

def plotEdgeList( problem, edges, specialNodes=None, specialEdges=None,\
                  title=None, fontsize=6, wait_func=None ):
    """
    Displays a graph of a list of edges.
    
    Parameters:
        **problem** : class  
            The problem object as returned by function load_problem of
            package tsplib95.  
        **edges** : list of pairs of int  
            The edges to be displayed.  
        **specialNodes** : list of int  
            A list of nodes seen as special and drawn in red color.
            Default=None.  
        **specialEdges** : list of pairs of int  
            List of special edges to be drawn in red color. Default=None.  
        **title** : str  
            String used as header for the figure. Default=None.    
        **fontsize** : int  
            Size of the font used to display node numbers.  
        **wait_func** : function, optional  
            function to be called (without any arguments)
            after the drawing was made. Usually used when
            matplotlib.pyplot is not in interactive mode
            so that execution halts before proceeding.
            Default: None                              
    """
    global __fig
    global __ax
    global __canvas
    
    no_coords = problem.node_coords is None or len(problem.node_coords)==0
    coords = problem.display_data if no_coords else problem.node_coords
    X = [ coords[i][0] for i in problem.get_nodes()]
    Y = [ coords[i][1] for i in problem.get_nodes()]    
    
    __initFig()
        
    plt.plot(X,Y,'bo',markersize=2)
    for e in edges: __plotEdge( coords[e[0]], coords[e[1]] )
    if specialEdges != None:
        for e in specialEdges:
            __plotEdge( coords[e[0]], coords[e[1]], special=True )
    if specialNodes != None:
        XS = [coords[i][0] for i in specialNodes]
        YS = [coords[i][1] for i in specialNodes]
        plt.plot(XS,YS,'ro',markersize=4)
    for i in problem.get_nodes():
        plt.text( coords[i][0], coords[i][1], str(i),fontsize=fontsize)

    if title != None: plt.title(title)
    
    if __canvas is None:
        plt.show()
    else:
        __canvas.draw()

    if not wait_func is None: wait_func()