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

GR = (sqrt(5) - 1) / 2


def gss(f, a, b, tol=1e-20):
    """golden section search to minimize along one coordinate"""
    c = b - GR * (b - a)
    d = a + GR * (b - a)
    while abs(c - d) > tol:
        fc = f(c)
        fd = f(d)
        if fc < fd:
            b = d
            d = c
            c = b - GR * (b - a)
        else:
            a = c
            c = d
            d = a + GR * (b - a)
        sys.stderr.write("gss:" + " ".join(map(str, [a, b, fc, fd])) + "\n")
        sys.stderr.flush()
    return (b + a) / 2, fc


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


def ll_old(lamb, alpha, mu):
    r1 = sum(log(lamb + sum(alpha * w * mu * exp(-mu * (t1 - t2))
                            for (n2, t2, w) in s))
             for ((n1, t1), s) in event_edges.iteritems())
    r2 = sum(sum(alpha * w * (1 - exp(-mu * (nodes[n2] - t1)))
                 for n2, w 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 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 sa(x, y, z, sigma=0.5, niter=70, fc=None):
    T = 0.1
    e = 1.1
    if fc:
        fo = fc
    else:
        fo = ll(x, y, z)
    for _ in xrange(niter):
        sys.stderr.write("sa: " + " ".join(map(str, [T, sigma, x, y, z, fo]))
                         + "\n")
        sys.stderr.flush()
        yn = max(y + gauss(0, sigma * y + 1e-10), 0)
        zn = max(z + gauss(0, sigma * z + 1e-10), 0)
        fn = ll(x, yn, zn)
        if fn < fo or exp((fo - fn) / T) > random():
            y = yn
            z = zn
            sigma *= 2
            fo = fn
        else:
            sigma /= e
        if sigma < 1e-5:
            break
        T *= 0.99
    return y, z, fo


def optimize_with_sa(x, y, z, niter=200):

    def f(x):
        return ll(x, y, z)

    # y, z = sa(x, y, z)
    y, z, fc = sa(x, y, z)
    for _ in xrange(niter):
        x, fc = gss(f, 0, 1e-3, tol=1e-10)
        y, z, fc = sa(x, y, z, fc=fc)
        print x, y, z, fc
        sys.stdout.flush()


def optimize_with_gss(x, y, z, niter=100):

    def f(x):
        return ll(x, y, z)

    def g(y):
        return ll(x, y, z)

    def h(z):
        return ll(x, y, z)

    for _ in xrange(niter):
        y, fc = gss(g, 0, 1, tol=1e-10)
        z, fc = gss(h, 0, 1, tol=1e-10)
        x, fc = gss(f, 0, 1e-3, tol=1e-10)
        print x, y, z, fc
        sys.stdout.flush()


def coarse_search():
    d = {}
    for line in open("values.txt"):
        v = map(float, line.strip().split())
        d[tuple(v[:3])] = v[3]
    p = Pool(5)
    lamb = [20. ** i for i in range(-10, 0)]
    alpha = [20. ** i for i in range(-15, 15)]
    mu = [20. ** i for i in range(-15, 15)]

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

    p.map(aux, product(lamb, alpha, mu))


if __name__ == "__main__":
    nodes, edges, events, event_edges = load(open("data2.pickle", "rb"))

    x = 1.2e-5
    y = 0.002
    z = 0.004
    # sa(x, y, z)

    with open(sys.argv[1]) as fh:
        l = [map(float, line.strip().split()[:3]) for line in fh]
    for e in l:
        optimize_with_sa(*e)

    # optimize_with_gss(x, y, z)
    # print ll(x, y, z)
    # coarse_search()