import networkx as nx
import numpy as np

def icc_cascade(G, p_init):
    """
    input: graph with prob as edge attr
    returns: 2D boolean matrix with indep. casc.
        where True means node was active at that time step
    p_init: proba that node in seed set
    """
    susceptible = np.ones(G.number_of_nodes(), dtype=bool)
    active = np.random.rand(G.number_of_nodes()) < p_init
    susceptible = susceptible - active
    cascade = []
    while sum(active) and sum(susceptible):
        cascade.append(active)
        tmp = np.zeros(G.number_of_nodes(), dtype=bool)
        for node in np.where(active)[0]:
            for edge in G.edges(node, data=True):
                tmp[edge[1]] += np.random.rand() < edge[2]["weight"] \
                                and susceptible[edge[1]]
        active = tmp
        susceptible = susceptible - active
        cascade.append(active)
    return cascade


def icc_cascade_2(G, p_init):
    """
    input: graph with prob as edge attr
    returns: 2D boolean matrix with indep. casc.
        where True means node was active at that time step
    p_init: proba that node in seed set
    """
    susceptible = np.ones(G.number_of_nodes(), dtype=bool)
    active = np.random.rand(G.number_of_nodes()) < p_init
    susceptible = susceptible - active
    cascade = []
    while active.any() and susceptible.any():
        cascade.append(active)
        active = np.exp(np.dot(G.logmat, active)) \
                    < np.random.rand(G.number_of_nodes())
        active = active & susceptible
        susceptible = susceptible - active
    return cascade


def test():
    """
    unit test
    """
    G = nx.erdos_renyi_graph(1000, 1, directed=True)
    G.logmat = np.zeros((G.number_of_nodes(), G.number_of_nodes()))
    for edge in G.edges(data=True):
        edge[2]['weight'] = .3*np.random.rand()
        G.logmat[edge[0],edge[1]] = np.log(1 - edge[2]["weight"])
    import time
    t0 = time.time()
    print len(icc_cascade(G, p_init=.1))
    t1 = time.time()
    print t1 - t0
    print len(icc_cascade_2(G, p_init=.1))
    print time.time() - t1

if __name__ == "__main__":
    test()