import numpy as np

class BasicTemplateAlgorithm(QCAlgorithm):

	def Initialize(self):
		self.SetCash(25000)
		self.SetStartDate(2006,4,26)
		self.SetEndDate(2012,4,9)
		
		self.ewa = self.AddEquity('EWA',Resolution.Minute).Symbol
		self.ewc = self.AddEquity('EWC',Resolution.Minute).Symbol
		
		#specify variables

		self.delta = 0.0001
		self.vw = self.delta/(1-self.delta)*np.eye(2)
		self.ve = 0.001
		
		self.beta = np.zeros(2)
		self.P = np.zeros((2,2))
		self.R = None
		
		self.pos = None
		self.day = None

	def OnData(self, data):
		time = str(data.Time).split(' ')[1]
		day = str(data.Time).split(' ')[0].split('-')[2]
		if time.split(':')[0] != '15' or time.split(':')[1] not in ['55','56','57','58','59','60']:
			return
		
		if self.day is not None and self.day == day:
			return
		
		
		x = np.asarray([float(data[self.ewa].Close),1.0]).reshape((1,2))
		y = float(data[self.ewc].Close)
		
		if self.R is not None:
			self.R = self.P + self.vw
		else:
			self.R = np.zeros((2,2))
			
		yhat = x.dot(self.beta)
		
		Q = x.dot(self.R).dot(x.T)+self.ve
		threshold = np.sqrt(Q)
		error = y - yhat
		Kalman = self.R.dot(x.T)/Q
		self.beta = self.beta + Kalman.flatten()*error
		self.P = self.R - Kalman*x.dot(self.R)
		
		if self.pos is None:
			if error < -threshold:
				quantity = float(self.Portfolio.Cash)/float(data[self.ewa].Close)
				self.Buy(self.ewc,1000)
				self.Sell(self.ewa,1000*self.beta[0])
				self.pos = 'long'
				self.day = day
				
			elif error > threshold:
				quantity = float(self.Portfolio.Cash)/float(data[self.ewa].Close)
				self.Sell(self.ewc,1000)
				self.Buy(self.ewa,1000*self.beta[0])
				self.pos = 'short'
				self.day = day
				
		else:
			if self.pos == 'long' and error > -threshold:
				self.Liquidate()
				self.day = day
				self.pos = None
			elif self.pos == 'short' and error < threshold:
				self.Liquidate()
				self.day = day
				self.pos = None