```#!/usr/bin/env python
#
#
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#
# Unless required by applicable law or agreed to in writing, software
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and

import matplotlib.pyplot as plt
import numpy as np
import statsmodels.api as sm
from datetime import datetime
import pytz

from alephnull.transforms import batch_transform

@batch_transform
def ols_transform(data, sid1, sid2):
"""Computes regression coefficient (slope and intercept)
via Ordinary Least Squares between two SIDs.
"""
p0 = data.price[sid1]
slope, intercept = sm.OLS(p0, p1).fit().params

return slope, intercept

"""Pairtrading relies on cointegration of two stocks.

The expectation is that once the two stocks drifted apart
(i.e. there is spread), they will eventually revert again. Thus,
if we short the upward drifting stock and long the downward
widened we can sell the spread with profit once they converged
again. A nice property of this algorithm is that we enter the
market in a neutral position.

This specific algorithm tries to exploit the cointegration of
Pepsi and Coca Cola by estimating the correlation between the
two. Divergence of the spread is evaluated by z-scoring.
"""

def initialize(self, window_length=100):
self.invested = 0
self.window_length = window_length
self.ols_transform = ols_transform(refresh_period=self.window_length,
window_length=self.window_length)

def handle_data(self, data):
######################################################
# 1. Compute regression coefficients between PEP and KO
params = self.ols_transform.handle_data(data, 'PEP', 'KO')
if params is None:
return
intercept, slope = params

######################################################
# 2. Compute spread and zscore
zscore = self.compute_zscore(data, slope, intercept)
self.record(zscores=zscore)

######################################################
# 3. Place orders
self.place_orders(data, zscore)

def compute_zscore(self, data, slope, intercept):
"""1. Compute the spread given slope and intercept.
"""
spread = (data['PEP'].price - (slope * data['KO'].price + intercept))
return zscore

def place_orders(self, data, zscore):
"""Buy spread if zscore is > 2, sell if zscore < .5.
"""
if zscore >= 2.0 and not self.invested:
self.order('PEP', int(100 / data['PEP'].price))
self.order('KO', -int(100 / data['KO'].price))
self.invested = True
elif zscore <= -2.0 and not self.invested:
self.order('PEP', -int(100 / data['PEP'].price))
self.order('KO', int(100 / data['KO'].price))
self.invested = True
elif abs(zscore) < .5 and self.invested:
self.invested = False

"""
decrease exposure, regardless of position long/short.
buy for a short position, sell for a long.
"""
ko_amount = self.portfolio.positions['KO'].amount
self.order('KO', -1 * ko_amount)
pep_amount = self.portfolio.positions['PEP'].amount
self.order('PEP', -1 * pep_amount)

if __name__ == '__main__':
start = datetime(2013, 1, 1, 0, 0, 0, 0, pytz.utc)
end = datetime(2014, 1, 1, 0, 0, 0, 0, pytz.utc)
start=start, end=end)