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
from scipy.optimize import minimize
import numpy as np

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 approx(x):
    if x > 1e-10:
        return 1 - exp(-x)
    else:
        return x


def ll_gradient(lamb, alpha, mu):
    l1 = [(1 + 0.43 * sin(0.0172 * t1 + 4.36),
          sum(1. / d ** 2 * mu * exp(-mu * (t1 - t2)) for (n2, t2, d) in s))
          for ((n1, t1), s) in event_edges.iteritems()]
    r1 = sum(log(lamb * a + alpha * b) for a, b in l1)
    g1lambda = sum(a / (lamb * a + alpha * b) for a, b in l1)
    g1alpha = sum(b / (lamb * a + alpha * b) for a, b in l1)

    t2 = sum(sum(1. / d ** 2 * approx(mu * (nodes[n2][0] - t1))
                 for n2, d in edges[n1].iteritems()
                 if nodes[n2][0] > t1)
             for (n1, t1) in iter_events(events))
    r2 = alpha * t2
    g2lambda = 0
    g2alpha = t2

    t3 = sum(node[1] for node in nodes.itervalues())
    r3 = lamb * t3
    g3lambda = t3
    g3alpha = 0

    return [-(r1 - r2 - r3), np.array([-(g1lambda - g2lambda - g3lambda),
                                       -(g1alpha - g2alpha - g3alpha)])]


def ll(lamb, alpha, mu):
    r1 = sum(log(lamb * (1 + 0.43 * sin(0.0172 * t1 + 4.36))
                 + sum(alpha / d ** 2 * mu * exp(-mu * (t1 - t2))
                       for (n2, t2, d) in s))
             for ((n1, t1), s) in event_edges.iteritems())
    r2 = sum(sum(alpha / d ** 2 * approx(mu * (nodes[n2][0] - t1))
                 for n2, d in edges[n1].iteritems()
                 if nodes[n2][0] > t1)
             for (n1, t1) in iter_events(events))
    r3 = lamb * sum(node[1] for node in 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), 0)
        zn = max(z + gauss(0, sigma * z), 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 sa_bis(x, y, z, sigma=0.5, niter=30, 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()
        zn = max(z + gauss(0, sigma * z), 0)
        fn = ll(x, y, zn)
        if fn < fo or exp((fo - fn) / T) > random():
            z = zn
            sigma *= 1.5
            fo = fn
        else:
            sigma /= e
        if sigma < 1e-5:
            break
        T *= 0.99
    return z, fo


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

    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_bfgs(x, y, z, niter=4):

    def f(x0, zum):
        return ll_gradient(x0[0], x0[1], zum)

    z, fc = sa_bis(x, y, z)
    for _ in xrange(niter):
        result = minimize(f, np.array([x, y]), method='L-BFGS-B', jac=True,
                          bounds=[(0, None), (0, None)],
                          options={"maxiter": 50, "disp": True},
                          args=[z])
        x, y = result.x
        fc = result.fun
        z, fc = sa_bis(x, y, z, fc=fc)
        print x, y, z, fc
        sys.stdout.flush()


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

    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 = [5e-6, 1e-5, 1.5e-5, 2e-5, 2.5e-5]
    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("data-all.pickle", "rb"))
    x = 1.875e-5
    y = 6.1e14
    z = 3e-20
    # optimize_with_sa(1.25e-05, 2.048e+13, 9.1552734375e-20)
    with open(sys.argv[1]) as fh:
        l = [map(float, line.strip().split()[:3]) for line in fh]
    for e in l:
        optimize_with_bfgs(*e)

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