from numpy.random import poisson, uniform
import numpy as np
from math import cos, exp, log, pi
from itertools import chain
from networkx.generators.random_graphs import watts_strogatz_graph
from csv import DictWriter
import networkx as nx


def homogenous_poisson_times(tmax, mu):
    n = poisson(mu * tmax)
    u = np.sort(uniform(size=n))
    return tmax * u


def _ih_poisson_times(tmax, integral, inverse):
    itmax = integral(tmax)
    n = poisson(itmax)
    if n == 0:
        return np.array([])
    u = uniform(size=n)
    iv = np.vectorize(lambda x: inverse(x * itmax))
    u = np.sort(iv(u))
    return u


def sin_rate_times(tmax, mu, a, w, phi):
    def integral(t):
        return mu * t + mu * a / w * (cos(phi) - cos(w * t + phi))

    def inverse(z):
        a, b = 0, tmax
        while b - a >= 1e-10:
            mid = (a + b) / 2.
            if integral(mid) <= z:
                a = mid
            else:
                b = mid
        return mid

    return _ih_poisson_times(tmax, integral, inverse)


def exp_rate_times(tmax, alpha, beta):
    def integral(t):
        return alpha * (1 - exp(-beta * t))

    def inverse(z):
        return - 1 / beta * log(1 - z / alpha)

    return _ih_poisson_times(tmax, integral, inverse)


def bfs(g, node):
    d = {node: 0}
    layer = g[node]
    dist = 0
    while dist < 3:
        dist += 1
        for n in layer:
            d[n] = dist
        layer = chain.from_iterable([u for u in g[n] if u not in d]
                                    for n in layer)
    del d[node]
    return d


def simulate(g, tmax, mu, a, w, phi, alpha, beta):
    events = set()
    gen = set()
    gen_number = 1

    for node in xrange(len(g)):
        gen |= set((node, t) for t in sin_rate_times(tmax, mu, a, w, phi))
        events |= gen

    print gen_number, len(events)

    while gen:
        gen_number += 1
        next_gen = set()
        for node, time in gen:
            for neighbor, d in bfs(g, node).items():
                next_gen |= set((neighbor, time + t)
                                for t in exp_rate_times(tmax - time,
                                                        alpha / d ** 2, beta))
        gen = next_gen
        events |= gen
        print gen_number, len(events)
    return events


def write(g, events):
    with open('test_nodes.csv', 'w') as fh:
        fieldnames = ['name', 'fatal_day']
        writer = DictWriter(fh, fieldnames=fieldnames)
        writer.writeheader()
        for n in g:
            writer.writerow({'name': n, 'fatal_day': 'NA'})
    with open('test_edges.csv', 'w') as fh:
        fieldnames = ['from', 'to', 't1', 'dist']
        writer = DictWriter(fh, fieldnames=fieldnames)
        writer.writeheader()
        for n, t in events:
            for neighbor, d in bfs(g, n).items():
                writer.writerow({'from': n, 'to': neighbor,
                                 'dist': d, 't1': t})


if __name__ == "__main__":
    g = watts_strogatz_graph(10000, 30, 0.2)
    print nx.is_connected(g)
    events = simulate(g, 3012, 1.19e-5, 0.43, 2 * pi / 365.24, 4.36,
                      7.82e-3, 3.74e-3)
    write(g, events)