import numpy as np
from scipy import optimize
from IPython.kernel import client
from timeit import default_timer as timer

T = 100000

def obj(params, Y, X):
    """
    Normal log-likelihood
    """
    nobs2 = len(X)/2.
    resid = Y - np.dot(X,params)[:,None]
    return - nobs2*np.log(2*np.pi)-nobs2*np.log(1/(2*nobs2) *\
            np.dot(resid.T, resid)) - nobs2

np.random.seed(12345)
X = np.random.uniform(0,10,size=(T,2))
beta = np.array([3.5,-3.5])
Y = np.dot(X,beta[:,None]) + np.random.normal(size=(T,1))
mec = client.MultiEngineClient()

eta = 1e-8
print "T = "+str(T)
mec.push(dict(optimize=optimize, np=np, Y=Y,X=X,eta=eta))

mec.push_function(dict(obj=obj))
start_params = np.array([0,0])
start_params1 = start_params.copy()
start_params2 = start_params + 5
start_params3 = start_params - 5
start_params4 = np.random.randint(-10,10,size=2)

t = timer()
mec.push(dict(params=start_params1), targets=0)
mec.push(dict(params=start_params2), targets=1)
mec.push(dict(params=start_params3), targets=2)
mec.push(dict(params=start_params4), targets=3)
mec.execute("""ret = optimize.fmin_tnc(obj, params, approx_grad=True, eta=eta,
                maxfun=3000, args=(Y,X), messages=0)""")
ret1 = mec.pull('ret', targets=0)
ret2 = mec.pull('ret', targets=1)
ret3 = mec.pull('ret', targets=2)
ret4 = mec.pull('ret', targets=3)
print str(timer() - t)+" seconds for parallel optimization"

t = timer()
ret1 = optimize.fmin_tnc(obj, start_params1, approx_grad=True, eta=eta,
                maxfun=3000, args=(Y,X), messages=0)
ret2 = optimize.fmin_tnc(obj, start_params2, approx_grad=True, eta=eta,
                maxfun=3000, args=(Y,X), messages=0)
ret3 = optimize.fmin_tnc(obj, start_params3, approx_grad=True, eta=eta,
                maxfun=3000, args=(Y,X), messages=0)
ret4 = optimize.fmin_tnc(obj, start_params4, approx_grad=True, eta=eta,
                maxfun=3000, args=(Y,X), messages=0)
print str(timer() - t)+" seconds for serial optimization"