import pylab as pl
import scipy.stats as ss

# simulate two random walks using a series of random coinflips. the random walk is the cumulative sum of the number of heads (+1)
# or tails (-1) flipped. 
np.random.seed(1)
coinflip1_v = sign(1*(np.random.random(1000) > 0.5)-0.5)
coinflip2_v = sign(1*(np.random.random(1000) > 0.5)-0.5)
x1_v = cumsum(coinflip1_v)
x2_v = cumsum(coinflip2_v)

pl.figure(figsize=(24,2))
pl.plot(coinflip1_v)
pl.ylim([-1.5,1.5])
pl.xlabel('Coinflip')
pl.ylabel('Heads/Tails')
pl.title('Coinflips #1');

pl.figure(figsize=(24,2))
pl.plot(coinflip2_v,'r')
pl.ylim([-1.5,1.5])
pl.xlabel('Coinflip')
pl.ylabel('Heads/Tails')
pl.title('Coinflips #2');

pl.figure()
pl.plot(x1_v,label='Walk #1')
pl.plot(x2_v,'r',label='Walk #2')
pl.grid()
pl.title('Random walks')
pl.legend(loc=2)
pl.xlabel('Coinflip')
pl.ylabel('Position');

xcorr_t = ss.pearsonr(x1_v,x2_v)
print "Random walk cross-correlation (rho): %1.2f" % xcorr_t[0]
print "Random walk cross-correlation (pval): %1.4e" % xcorr_t[1]

pl.figure()
pl.plot(x1_v[:-1],x1_v[1:],'o')
pl.title('Walk #1')
pl.xlabel('Position at Coinflip i')
pl.ylabel('Position at Coinflip i+1')

pl.figure()
pl.plot(x2_v[:-1],x2_v[1:],'or')
pl.title('Walk #2')
pl.xlabel('Position at Coinflip i')
pl.ylabel('Position at Coinflip i+1');

print "Random walk 1 auto-correlaton (rho): %1.2f" % ss.pearsonr(x1_v[:-1],x1_v[1:])[0]
print "Random walk 2 auto-correlaton (rho): %1.2f" % ss.pearsonr(x2_v[:-1],x2_v[1:])[0]

# this example calls R code. to run it, you'll need to install R, the rpy2 Python module, and the 'forecast' R package.
import rpy2.robjects as R
import rpy2.robjects.numpy2ri
rpy2.robjects.numpy2ri.activate()
R.r('library(forecast)');

R.r.assign('x1_v',x1_v)
R.r.assign('x2_v',x2_v)

R.r('x1_arima<-auto.arima(x=x1_v,max.p=2,d=1,max.q=2,ic="bic",approximation=FALSE)');
R.r('x2_arima<-auto.arima(x=x2_v,max.p=2,d=1,max.q=2,ic="bic",approximation=FALSE)');

residual1_v = np.array(R.r('x1_arima$residuals'))
residual2_v = np.array(R.r('x2_arima$residuals'))

pl.figure(figsize=(24,2))
pl.plot(residual1_v)
pl.ylim([-1.5,1.5])
pl.xlabel('Coinflip')
pl.ylabel('Residual')
pl.title('Residuals #1');

pl.figure(figsize=(24,2))
pl.plot(residual2_v,'r')
pl.ylim([-1.5,1.5])
pl.xlabel('Coinflip')
pl.ylabel('Residual')
pl.title('Residuals #2');

print "Residual 1 auto-correlaton (rho): %1.3f" % ss.pearsonr(residual1_v[:-1],residual1_v[1:])[0]
print "Resdiual 2 auto-correlaton (rho): %1.3f" % ss.pearsonr(residual2_v[:-1],residual2_v[1:])[0]

xcorr_t = ss.pearsonr(residual1_v,residual2_v)
print "Random walk cross-correlation (rho): %1.2f" % xcorr_t[0]
print "Random walk cross-correlation (pval): %1.2f" % xcorr_t[1]