from cPickle import load
from itertools import product
from math import exp, log, sin
import sys
from multiprocessing import Pool


def iter_events(events):
    for n, s in events.iteritems():
        for t in s:
            yield (n, t)


def inprod(a, b):
    return tuple(float(x * y) for x, y in zip(a, b))


def ll(lamb, alpha, mu):
    r1 = sum(log(lamb * (1 + 0.43 * sin(0.0172 * t1 + 4.36))
                 + sum(alpha / d * mu * exp(-mu * (t1 - t2))
                       for (n2, t2, d) in s))
             for ((n1, t1), s) in event_edges.iteritems())
    r2 = sum(sum(alpha / d * (1 - exp(-mu * (nodes[n2] - t1)))
                 for n2, d in edges[n1].iteritems() if nodes[n2] > t1)
             for (n1, t1) in iter_events(events))
    r3 = lamb * sum(nodes.itervalues())
    return -(r1 - r2 - r3)


def get_values():
    d = {}
    for line in open(sys.argv[1]):
        v = map(float, line.strip().split())
        d[tuple(v[:3])] = v[3]
    l = d.items()
    l.sort(key=lambda x: x[1])
    for line in open("refine.txt"):
        v = map(float, line.strip().split())
        d[tuple(v[:3])] = v[3]
    for a, _ in l[:100]:
        t = [1. / i for i in range(2, 4)] + [float(i) for i in range(1, 4)]
        for b in product(t, repeat=3):
            l, al, m = inprod(a, b)
            if (l, al, m) in d:
                continue
            yield (l, al, m)


def refine():
    p = Pool(5)

    def aux(x):
        l, a, m = x
        print l, a, m, ll(l, a, m)
        sys.stdout.flush()

    p.map(aux, get_values())


if __name__ == "__main__":
    nodes, edges, events, event_edges = load(open("data-dist1.pickle", "rb"))
    refine()