import networkx as nx
import numpy as np

class InfluenceGraph(nx.Graph):
    """
    networkX graph with mat and logmat attributes
    """
    def __init__(self, max_proba, *args, **kwargs):
        self.max_proba = max_proba
        super(InfluenceGraph, self).__init__(*args, **kwargs)

    def erdos_init(self, n, p):
        G = nx.erdos_renyi_graph(n, p, directed=True)
        self.add_nodes_from(G.nodes())
        self.add_edges_from(G.edges())

    def createStanfordGraph(self, file):
        """
        Takes a file from the Stanford collection of networks
        Need to remove comments on top of the file
        Graph still needs to be weighted on the edges
        """
        f = open(file, 'r')
        data = f.readlines()
        G = nx.DiGraph()
        for edge in data:
            split1 = edge.split('\t')
            split2 = split1[1].split('\n')
            
            u = int(split1[0])
            v = int(split2[0])
            G.add_edge(u,v)
        self.add_nodes_from(G.nodes())
        self.add_edges_from(G.edges())
        
    @property
    def mat(self):
        if not hasattr(self, '_mat'):
            self._mat = (self.max_proba * np.random.rand(len(self), len(self))
                         * np.asarray(nx.adjacency_matrix(self)))
        return self._mat

    @property
    def logmat(self):
        if not hasattr(self, '_logmat'):
            self._logmat = np.log(1 - self.mat)
        return self._logmat


def icc_cascade(G, p_init):
    """
    Returns boolean vectors for one cascade
    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 generate_cascades(G, p_init, n_cascades):
    """
    returns list of cascades
    """""
    return [icc_cascade(G,p_init) for i in xrange(n_cascades)]


def greedy_prediction(G, cascades):
    """
    returns estimated graph
    """
    G_hat = InfluenceGraph(max_proba=None)
    G.add_nodes_from(G.nodes())

def test():
    """
    unit test
    """
    G = InfluenceGraph(max_proba = .3)
    G.erdos_init(n = 100, p = 1)
    import time
    t0 = time.time()
    print len(icc_cascade(G, p_init=.1))
    t1 = time.time()
    print t1 - t0

if __name__ == "__main__":
    test()