%pylab inline

import time, seaborn
import pandas as pd
import numpy as np
import pysal as ps
import matplotlib.pyplot as plt
from scipy.sparse import eye
from scipy.sparse.linalg import inv as spinv
from scipy.linalg import inv
from pysal.spreg.utils import spdot

%load_ext rpy2.ipython

%%R
library(spdep)
library(sphet)
print(sessionInfo())

%%R
# Utilities to run tests on the R sides
buildW <- function(n){
    t0 = Sys.time()
    w = cell2nb(sqrt(n), sqrt(n))
    we = nb2listw(w)
    t1 = Sys.time()
    out = list('we'=we, 'tim'=t1-t0)
    return(out)
}

runOLS <- function(y, x){
    db <- cbind(data.frame(y=y), data.frame(x))
    t0 = Sys.time()
    m <- lm('y ~ V1 + V2 + V3', db)
    t1 = Sys.time()
    return(t1-t0)
}

runGMERR <- function(y, x, w){
    db <- cbind(data.frame(y=y), data.frame(x))
    t0 = Sys.time()
    m <- spreg(y ~ V1 + V2 + V3, db, listw=w)
    t1 = Sys.time()
    return(t1-t0)
}

runMLERR <- function(y, x, w){
    db <- cbind(data.frame(y=y), data.frame(x))
    t0 = Sys.time()
    m <- errorsarlm(y ~ V1 + V2 + V3, db, listw=w)
    t1 = Sys.time()
    return(t1-t0)
}

runGMLAG <- function(y, x, w){
    db <- cbind(data.frame(y=y), data.frame(x))
    t0 = Sys.time()
    m <- spreg(y ~ V1 + V2 + V3, db, listw=w, model = "lag")
    t1 = Sys.time()
    return(t1-t0)
}

runMLLAG <- function(y, x, w){
    db <- cbind(data.frame(y=y), data.frame(x))
    t0 = Sys.time()
    m <- lagsarlm(y ~ V1 + V2 + V3, db, listw=w)
    t1 = Sys.time()
    return(t1-t0)
}

np.random.seed(1234)

def sim(sizes, rho, lamb, ml=False):
    timings = pd.DataFrame()
    %R timingsR = data.frame()
    for n in sizes:
            ### Data
        print "Working on Size %i"%n
        timing = {}
        # Python
        t0 = time.time()
        w = ps.lat2W(int(np.sqrt(n)), int(np.sqrt(n)))
        t1 = time.time()
        w.transform = 'R'
        timing['W'] = t1 - t0
        x = np.random.random((n, 3))
        X = np.hstack((np.ones((n, 1)), x))
        eps = np.random.normal(size=(n, 1))
        xb =np.dot(X, betas) 
        spxb_r = ps.spreg.utils.power_expansion(w, xb, rho)
        speps_r = ps.spreg.utils.power_expansion(w, eps, rho)
        speps_l = ps.spreg.utils.power_expansion(w, eps, lamb)
        y = xb + eps
        y_lag = spxb_r + speps_r
        y_err = xb + speps_l
        # R
        %Rpush n x y y_lag y_err
        %R x = as.data.frame(x)
        %R wtw = buildW(n)
        %R res = data.frame(W=wtw$tim)
        %R w = wtw$we
            ### OLS
        t0 = time.time()
        m = ps.spreg.OLS(y, x)
        t1 = time.time()
        timing['OLS'] = t1 - t0
        %R t = runOLS(y, x)
        %R res = cbind(res, data.frame(OLS=t))
            ### Spatial Error
        t0 = time.time()
        m = ps.spreg.GM_Error_Het(y_err, x, w=w)
        t1 = time.time()
        timing['Err_gm'] = t1 - t0
        %R t = runGMERR(y_err, x, w)
        %R res = cbind(res, data.frame(Err_gm=t))
        if ml:
            t0 = time.time()    
            m = ps.spreg.ML_Error(y, x, w=w)
            t1 = time.time()
            timing['Err_ml'] = t1 - t0
            %R t = runMLERR(y_err, x, w)
            %R res = cbind(res, data.frame(Err_ml=t))
        # Spatial lag
        t0 = time.time()
        m = ps.spreg.GM_Lag(y_lag, x, w=w)
        t1 = time.time()
        timing['Lag_gm'] = t1 - t0
        %R t = runGMLAG(y_lag, x, w)
        %R res = cbind(res, data.frame(Lag_gm=t))
        if ml:
            t0 = time.time()
            m = ps.spreg.ML_Lag(y_lag, x, w=w)
            t1 = time.time()
            timing['Lag_ml'] = t1 - t0
            %R t = runMLLAG(y_lag, x, w)
            %R res = cbind(res, data.frame(Lag_ml=t))
        timings['n%i'%n] = pd.Series(timing)
        %R row.names(res) = paste('n', n, sep='')
        %R timingsR = rbind(timingsR, res)
    timingsPY = timings.T
    timingsPY.index = pd.MultiIndex.from_arrays([['Py']*timingsPY.shape[0], \
                                               timingsPY.index])
    %Rpull timingsR
    ri = %R rownames(timingsR)
    timingsR.index = pd.MultiIndex.from_arrays([['R']*ri.shape[0], ri])
    timings = timingsPY.T.join(timingsR.T).T
    return timings

sizes = [100, 625, 900, 3025, 4900, 10000, 25600, 50625, 99225]
betas = np.array([[1, 1, 1, 1]]).T
rho = 0.5
lamb = 0.5

%%time
ssizes = sizes[:5]
stimings = sim(ssizes, rho, lamb, ml=True)

stimings

for col in stimings:
    f, ax = plt.subplots(1)
    stimings[col].groupby(level=0).plot(ax=ax)
    plt.xlabel('N. Observations')
    plt.ylabel('Seconds')
    plt.legend()
    plt.title(col)

%%time
timings = sim(sizes, rho, lamb)

timings

for col in timings:
    f, ax = plt.subplots(1)
    timings[col].groupby(level=0).plot(ax=ax)
    plt.xlabel('N. Observations')
    plt.ylabel('Seconds')
    plt.legend()
    plt.title(col)