import numpy as np
import numpy.random as nr


def likelihood(p, cascade, node):
    m = build_matrix(cascade, node)
    x = m[:-1, :]
    y = m[1:, node].reshape(-1)
    a = 1 - np.exp(-np.dot(x, p))
    return y * np.log(a) + (1 - y) * np.log(1-a)


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


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


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))
    x = np.zeros(g.shape[0]).astype(bool)
    x[0] = 1
    cascade = simulate_cascade(g, x)
    print likelihood(np.ones(g.shape[0]) * 0.5, cascade, 1)