import numpy as np
import numpy.random as nr


def likelihood(p, cascades, node):
    x, y = build_matrix(cascades, node)
    a = 1 - np.exp(-np.dot(x, p))
    return np.inner(y, np.log(a)) + np.inner((1 - y), np.log(1-a))


def build_matrix(cascades, node):

    def aux(cascade, node):
        try:
            m = np.vstack(x for x, s in cascade if s[node])
        except ValueError:
            return None
        x = m[:-1, :]
        y = m[1:, node]
        return x, y

    pairs = (aux(cascade, node) for cascade in cascades)
    xs, ys = zip(*(pair for pair in pairs if pair))
    x = np.vstack(xs)
    y = np.concatenate(ys)
    return x, y


def simulate_cascade(x, graph):
    susc = x ^ np.ones(graph.shape[0]).astype(bool)
    yield x, susc
    while np.any(x):
        x = 1 - np.exp(-np.dot(graph.T, x))
        y = nr.random(x.shape[0])
        x = (x >= y) & susc
        yield x, susc
        susc ^= x


def simulate_cascades(n, graph):
    for _ in xrange(n):
        x = np.zeros(g.shape[0], dtype=bool)
        x[nr.randint(0, g.shape[0])] = True
        yield simulate_cascade(x, graph)


if __name__ == "__main__":
    g = np.array([[0, 1, 1, 0], [1, 0, 0, 1], [1, 0, 0, 1], [0, 1, 1, 0]])
    p = 0.5
    g = np.log(1. / (1 - p * g))
    cascades = simulate_cascades(10, g)
    print likelihood(p * np.ones(g.shape[0]), cascades, 3)