#
#   QuantConnect Basic Template:
#	Fundamentals to using a QuantConnect algorithm.
#
#	You can view the QCAlgorithm base class on Github: 
#	https://github.com/QuantConnect/Lean/tree/master/Algorithm
#
import numpy as np

class DemoAlgorithm(QCAlgorithm):
	
	def __init__(self):
		self.symbols = ['MSFT','UNH','CTAS','JNS','COG']

	def Initialize(self):
		# Set the cash we'd like to use for our backtest
		# This is ignored in live trading 
		self.SetCash(20000)
		
		# Start and end dates for the backtest.
		# These are ignored in live trading.
		self.SetStartDate(2016,1,1)
		self.SetEndDate(2017,1,1)
		
		#convert the strings in our list into QuantConnect.Symbol object for convenience. 
		for i in range(len(self.symbols)):
			# Add those symbols to our universe
			symbol = self.AddEquity(self.symbols[i],Resolution.Daily).Symbol
			# replace the original strings with QuantConnect.Symbols
			self.symbols[i] = symbol
			# self.History returns the historical price data into a pandas DataFrame, 
			# with columns['close','open','high','low']
			history = self.History(self.symbols[i],30,Resolution.Daily)
			# assign the price list to the corresponding QuantConnect.symbol instance for convenience.
			self.symbols[i].prices = [x for x in history['close']]
			# Calculate the 30 days and 10 days sma.
			self.symbols[i].sma_10 = np.mean(self.symbols[i].prices[-10:])
			self.symbols[i].sma_30 = np.mean(self.symbols[i].prices)

        # the scheduled event will be on fire every Monday after the market open
		self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday, DayOfWeek.Monday), 
		                 self.TimeRules.AfterMarketOpen(self.symbols[0], 2), Action(self.Rebalance))



	def OnData(self, data):
		# the 'data' contains other information besides daipy price. e.g, if there is a stock split on weekend, it would
		# returns a data at that time. we check here to ensure it's a normal data slice contains all the price information. 
		if False in [data.ContainsKey(x) for x in self.symbols]:
			return
		
		# update each stocks' price list every day
		for asset in self.symbols:
			asset.prices.append(float(data[asset].Close))
			asset.prices = asset.prices[-30:]
	
	def Rebalance(self):
		# first calculate the raw_weight for each stock
		for asset in self.symbols:
			asset.sma_10 = np.mean(asset.prices[-10:])
			asset.sma_30 = np.mean(asset.prices)
			asset.raw_weight = (asset.sma_10 - asset.sma_30)/asset.sma_30
		
		# calculate the total weight
		total_weights = sum([abs(x.raw_weight) for x in self.symbols])
		
		# calculate normalized_weights and then long/short
		for asset in self.symbols:
			asset.normalized_weights = asset.raw_weight/total_weights
			self.SetHoldings(asset,asset.normalized_weights)
			if asset.normalized_weights > 0:
				self.Log('long %s at weight %f'%(str(asset),asset.normalized_weights))
			else:
				self.Log('short %s at weight %f'%(str(asset),asset.normalized_weights))