import matplotlib.pyplot as plt
import copy
import sys


def compare(li1, li2):
    for i, j in zip(li1, li2):
        if i != j:
            return False
    return True


def set_node_position(G, **kwargs):
    """
    Moves node to position

    Parameters
    ----------
    G        : Graph whose nodes are to be moved

    kwargs
    ------
    node     : The node to move if not specified moves all the nodes to the
               corresponding position in position
    position : The new position of the node. If node specified should be in
               form (x,y) or [x,y]. If node not specified should be in form
               [(x1,y1), (x2,y2), (x3,y3)]. It must have new position for
               each node in the graph.

    Example
    -------
    set_node_position(G, node=1, position=[0.5, 0.5])
    set_node_position(G, node=1, position=(0.5, 0.5))
    #Considering graph G has 3 nodes
    set_node_position(G, position=[(0.1,0.1), (0.2, 0.2), (0.3, 0.3)]
    """
    try:
        node = kwargs['node']
    except KeyError:
        pass
    try:
        position = kwargs['position']
    except KeyError:
        print("position argument required")
        sys.exit(0)

    # TODO: Recursive call the below for each node for it to work when no node is specified

    axes = plt.gca()
    node_index = G.nodes().index(node)
    no_of_nodes = G.number_of_nodes()

    if node:
        # move node
        x, y = position
        nodes_collection = axes.get_children()[no_of_nodes + 2]
        prev_positions = nodes_collection.get_offsets()
        prev_position = copy.deepcopy(prev_positions[node_index])
        prev_positions[node_index] = [x, y]

        # move edge
        edges_collection = axes.get_children()[no_of_nodes + 3]
        for edge in edges_collection.get_paths():
            if compare(edge.vertices[0], prev_position):
                edge.vertices[0] = position
            elif compare(edge.vertices[1], prev_position):
                edge.vertices[1] = position

        # move label
        label = axes.get_children()[node_index + 2]
        label.set_position([x,y])

    plt.draw()

def set_node_color(G, **kwargs):
    """
    Changes node colors

    Parameters
    ----------
    G     : Graph whose's node's color needs to be changed

    kwargs
    ------
    node  : The node whose color needs to be changed. If node not specified
            and only one color is specified changes all the nodes to that
            color. If node not specified and color is a list or tuple of
            colors, each node will be changed to its corresponding color.

    color : Can be a single value of a list of colors.kwargs

    The following color abbreviations are supported:
    ==========  ========
    character   color
    ==========  ========
    'b'         blue
    'g'         green
    'r'         red
    'c'         cyan
    'm'         magenta
    'y'         yellow
    'k'         black
    'w'         white
    ==========  ========

    Examples:
    set_node_color(G,
    """
    try:
        node = kwargs['node']
    except KeyError:
        node = None
    try:
        color = kwargs['position']
    except KeyError:
        print("color argument required")
        sys.exit(0)

    colors_dict = {'b': [0.,    0.,    1.,  1.],
                   'g': [0.,    0.5,   0.,  1.],
                   'r': [1.,    0.,    0.,  1.],
                   'c': [0.,    0.75,  0.75,  1.],
                   'm': [0.75,  0.,    0.75,  1.],
                   'y': [0.75,  0.75,  0.,  1.],
                   'k': [0.,    0.,    0.,  1.],
                   'w': [1.,    1.,    1.,  1.]}

    fig = plt.gcf()
    axes = plt.gca()
    no_of_nodes = G.number_of_nodes()
    nodes_collection = axes.get_children()[no_of_nodes + 2]

    # if node is specified manually changing the facecolor and edgecolor array so that
    # the colors for other nodes are retained
    if node:
        node_index = G.nodes().index(node)
        facecolor_array = nodes_collection.get_facecolor()
        edgecolor_array = nodes_collection.get_edgecolor()
        facecolor_array[node_index] = colors_dict[color]
        edgecolor_array[node_index] = colors_dict[color]

    # if node not specified call the matplotlibs set_color function
    else:
        nodes_collection.set_color(color)
@@ -3027,16 +3027,211 @@
 color(G,
+ node=1, color='b')%0A    set_node_color(G, node=1, color=%5B'b'%5D)%0A    set_node_color(G, color=%5B'r', 'b', 'g', 'y', 'm', 'k', 'c'%5D)%0A    set_node_color(G, color='b')%0A    set_node_color(G, color=('b'))
 %0A    %22%22%22
@@ -3434,24 +3434,132 @@
 ys.exit(0)%0A%0A
+    if (isinstance(color, list) or isinstance(color, tuple)) and len(color) == 1:%0A        color = color%5B0%5D%0A%0A
     colors_d
@@ -4552,16 +4552,17 @@
 tplotlib
+'
 s set_co

add check for list or tuple in argument