import copy
import random


def nx_color_spread(init_pt, graph,maxtsteps):
    """ Function to perform the random simulation of spreading the GHZ state on
    a weighted graph

    Parameters
    ----------
    init_pt, any type:
        Label for the initial point we are spreading from. Should be present in
        the graph.
    graph, Graph:
        Weighted graph describing connections between nodes. Weight of each node
        will be the probability of success
    maxtsteps, int:
        Maximum number of timesteps to run simulation for

    Returns
    -------
    color_count, list of ints:
        List whose ith element is the number of converted nodes at the ith
        timestep
    boundary_count, list of ints:
        List whose ith element is the size of the boundary at the ith timestep
    graph, Graph:
        The graph (complete with converted nodes recolored blue) after the final
        timestep
    """
    # Try to probabilistically color the graph and see how long it takes
    # Function will return a list of how many points were colored at each time
    # step
    color_count = [1]
    color_dict = {}
    for x in graph.nodes():
        color_dict[x] = 'white'
    color_dict[init_pt] = 'blue' # I like blue.
    nodes = [init_pt]
    boundary = list(graph[init_pt].keys())
    boundary_count = [len(boundary)]
    for t in range(maxtsteps):
        new_boundary = copy.deepcopy(boundary) # Make a copy of the current
                                               # boundary to update as we go
        for b in boundary: # Look at each node b on the boundary
            for n in graph[b]: 
                # Look at each node connected to those nodes
                if color_dict[n] == 'blue': 
                # If they are blue, there is chance they spread the entanglement
                    if random.uniform(0,1) <= graph[n][b]['weight']:
                        # weight provides the probability of spread
                        color_dict[b] = 'blue' # color that node
                        nodes.append(b)
                        # Now that we have added the node b, we need to remove
                        # it from the list of the boundary
                        new_boundary.remove(b)
                        for new_b in graph[b]:
                            # Don't add already-added nodes to boundary!
                            if color_dict[new_b] != 'blue':
                                # Don't double up in the boundary
                                if new_b not in new_boundary: 
                                    new_boundary.append(new_b)
                if color_dict[b] == 'blue':
                    break # Break out of this loop
        boundary = new_boundary
        color_count.append(len(nodes))
        boundary_count.append(len(boundary))
    return color_count