import numpy as np
import pandas as pd
import statsmodels.api as sm

class TrendFollowingAlgorithm(QCAlgorithm):
	def __init__(self):
		self.symbols = ['DIA','QQQ','LQD','HYG','USO','GLD','VNQ','RWX','UNG','DBA']
		self.num_history = 252/2
		self.slope_min = 0.15
		self.exit = 1
		self.multiple = 5

	def Initialize(self):
		self.SetCash(50000)
		self.SetStartDate(2013,03,01)
		self.SetEndDate(2014,04,01)
		
		for i in range(len(self.symbols)):
			self.symbols[i] = self.AddEquity(self.symbols[i],Resolution.Minute).Symbol

		
		history = self.History(self.num_history, Resolution.Daily)
		for i in self.symbols:
			bars = map(lambda x: x[i],history)
			i.time = [str(x.Time) for x in bars]
			i.prices = [float(x.Open) for x in bars]
			i.stopprice = None
			i.slope = None
			
			
		self.Schedule.On(self.DateRules.EveryDay(self.symbols[0]), 
		     self.TimeRules.AfterMarketOpen(self.symbols[0], 2), Action(self.do_regression))
		
		# self.Schedule.On(self.DateRules.EveryDay(self.symbols[0]), 
		#      self.TimeRules.AfterMarketOpen(self.symbols[0], 4), Action(self.trade))

	def OnData(self, data):
		if False in [data.ContainsKey(x) for x in self.symbols]:
			return
		
		time = str(self.Time).split(' ')[1]
		hour = time.split(':')[0]
		m = time.split(':')[1]
		#update price for regression
		if hour == '09' and m == '31':
			for i in self.symbols:
				i.prices.append(float(data[i].Open))
				i.prices.pop(0)
		
		#stop loss	
		if float(m)%5 == 0:
			for i in self.symbols:
				price = float(data[i].Close)
				if self.Portfolio[i].Quantity > 0:
					if price/i.order_price - 1 < -0.02:
						self.Liquidate(i)
					elif price/i.order_price - 1 > 0.03:
						self.Liquidate(i)
				elif self.Portfolio[i].Quantity < 0:
					if price/i.order_price - 1 > 0.02:
						self.Liquidate(i)
					elif price/i.order_price - 1 < -0.03:
						self.Liquidate(i)
				else:
					pass
				
				
				# i.stoploss = abs(i.slope*self.num_history / 252)+1
				
				# if self.Portfolio[i].Quantity > 0:
				# 	if i.stopprice < 0:
				# 		i.stopprice = price/i.stoploss
				# 	else:
				# 		i.stopprice = max(price/i.stoploss,i.stopprice)
				# 		if price < i.stopprice:
				# 			self.Liquidate(i)
				# elif self.Portfolio[i].Quantity < 0:
				# 	if i.stopprice < 0:
				# 		i.stopprice = price * i.stoploss
				# 	else:
				# 		i.stopprice = min(price*i.stoploss,i.stopprice)
				# 		if price > i.stopprice:
				# 			self.Liquidate(i)
				# else:
				# 	i.stopprice = None

	def regression(self,asset):
		x = sm.add_constant(range(len(asset.prices)))
		asset.model = sm.OLS(asset.prices,x).fit()
	
	def do_regression(self):
		for i in self.symbols:
			self.regression(i)
			i.beta = i.model.params[1]
			i.intercept = i.model.params[0]
			i.slope = i.beta/i.intercept * 252.0
			# i.delta = i.prices - (np.dot(i.beta,range(len(i.prices))) + i.intercept)
			i.delta = i.model.resid
			i.sd = np.std(i.prices)
			
			
			#exit long postion
			if self.Portfolio[i].Quantity > 0 and i.slope < 0:
				self.Liquidate(i)
			#exit short position
			if self.Portfolio[i].Quantity < 0 and i.slope > 0:
				self.Liquidate(i)
			
			#trend is up:
			if i.slope > self.slope_min:
				if i.delta[-1] > 0 and i.delta[-2] < 0 and self.Portfolio[i].Quantity == 0:
					i.stopprice = None
					self.SetHoldings(i, i.slope*self.multiple)
					i.order_price = float(self.Portfolio[i].Price)
					
				if i.delta[-1] > self.exit*i.sd and self.Portfolio[i].Quantity > 0:
					self.Liquidate(i)
					self.Log('profit take')
			#trend is down
			if i.slope < -self.slope_min:
				if i.delta[-1] < 0 and i.delta[-2] > 0 and self.Portfolio[i].Quantity == 0:
					i.stopprice = None
					self.SetHoldings(i,i.slope*self.multiple)
					i.order_price = float(self.Portfolio[i].Price)
					
				if i.delta[-1] < -self.exit * i.sd and self.Portfolio[i].Quantity < 0:
					self.Liquidate(i)
					self.Log('profit take')