import csv
from StringIO import StringIO
from datetime import datetime
from decimal import *
from sets import Set

class BasicTemplateAlgorithm(QCAlgorithm):

	def Initialize(self):
		
		raw_trade_data = """
		20170201 08:37,BAC,bear
20170201 09:11,ADP,bull
20170201 09:34,AAPL,bull
20170201 09:36,AAPL,bull
20170201 09:38,GM,bull
20170201 09:38,F,bull
20170201 09:48,AAPL,bull
20170201 09:52,TSLA,bear
20170201 10:35,WMT,bull
20170201 10:37,WMT,bear
20170201 12:22,AAPL,bear
20170201 13:39,AAPL,bull
20170201 14:12,AAPL,bull
20170201 14:59,AAPL,bear
20170201 15:23,AMZN,bull
20170202 08:06,BAC,bear
20170202 08:21,RL,bear
20170202 12:31,AMZN,bear
20170202 13:27,GS,bear
20170202 15:11,GOOG,bear
20170202 15:11,GOOGL,bear
20170203 08:16,PFE,bull
20170203 08:24,AAPL,bear
20170203 10:01,GS,bull
20170203 10:03,GS,bull
20170203 10:03,V,bull
20170203 12:28,UA,bear
20170203 12:42,JWN,bull
20170203 15:08,AMZN,bear
20170206 08:14,C,bull
20170206 08:56,AMZN,bull
20170207 08:34,AAPL,bull
20170207 08:34,GOOG,bull
20170207 08:34,GOOGL,bull
20170207 10:32,MSFT,bear
20170207 10:32,INTC,bear
20170207 14:47,DIS,bear
20170207 15:13,MAT,bear
20170207 15:13,HAS,bear
20170208 08:24,INTU,bear
20170208 08:25,SBUX,bull
20170208 11:00,JWN,bull
20170208 11:05,JPM,bull
20170208 11:10,JWN,bear
20170208 11:26,NVDA,bull
20170208 12:24,TSLA,bull
20170208 13:30,NVDA,bull
20170208 14:37,JWN,bull
20170208 14:43,DIS,bull
20170209 09:20,AAPL,bull
20170209 09:59,GS,bull
20170209 11:02,KO,bull
20170209 12:48,MINT,bull
20170209 15:25,BRKA,bull
20170209 15:25,BRKB,bull
20170210 08:55,AAPL,bull
20170210 09:51,SHLD,bull
20170210 10:47,NVDA,bull
20170210 10:48,SKX,bull
20170213 08:33,AAPL,bull
20170213 09:24,AAPL,bull
20170213 10:32,GS,bull
20170213 10:32,AAPL,bull
20170213 11:32,SBUX,bull
20170214 08:18,MAT,bear
20170214 08:18,BABA,bear
20170214 08:54,GS,bear
20170214 09:26,TGT,bull
20170214 09:26,GPS,bull
20170214 09:26,BBY,bull
20170214 10:11,AAPL,bull
20170214 13:03,MS,bear
20170214 13:14,AAPL,bull
20170214 13:32,BAC,bear
20170214 13:34,TSLA,bull
20170215 09:58,AAPL,bull
20170215 10:24,BX,bear
20170215 11:00,GS,bull
20170216 13:02,MS,bull
20170217 08:21,UN,bear
20170217 08:59,UN,bear
20170217 09:51,BA,bull
20170217 10:16,BA,bull
20170217 13:59,CSCO,bull
20170217 13:59,CSCO,bull
20170217 14:05,BA,bull
20170217 14:56,UN,bear
20170217 15:41,CSCO,bear
20170220 14:25,MSFT,bear
20170221 09:33,AAPL,bull
20170221 09:46,AMZN,bull
20170221 11:28,AAPL,bull
20170221 11:28,MS,bull
20170221 12:22,SKX,bull
20170221 13:27,PEP,bull
20170221 13:42,WFC,bear
20170221 14:59,CVX,bear
20170221 14:59,C,bear
20170221 15:46,UN,bear
20170222 08:24,GOOG,bull
20170222 08:24,GOOGL,bull
20170222 09:19,AAPL,bull
20170222 13:56,BRKA,bear
20170222 13:56,BRKB,bear
20170222 15:25,TSLA,bull
20170222 15:25,SCTY,bull
20170222 15:36,Fit,bull
20170222 15:36,GRMN,bull
20170223 08:38,BSX,bear
20170223 08:39,F,bull
20170223 08:39,GM,bull


		"""
		
		self.trade_data = {}
		equities = Set()
		
		# store a list of trades if there are multiple per minute
		f = StringIO(raw_trade_data)
		for row in csv.reader(f, delimiter=','):
			if row: # blank lines
				key = row[0] # [:-3] # take last 3 chars off if we have seconds
				if key in self.trade_data:
					self.trade_data[key].append(row)
				else:
					self.trade_data[key] = [row]
				if len(row) > 1:
					equities.add(row[1])
		
		for symbol in equities:
			self.AddEquity(symbol)
		
		self.SetCash(100000)
		self.SetStartDate(2017,2,1)
		self.SetEndDate(2017,2,22)
		self.AddEquity("SPY", Resolution.Minute)
		self.Schedule.On(
			self.DateRules.EveryDay("SPY"),
			self.TimeRules.BeforeMarketClose("SPY", 5),
			Action(self.Liquidate)
		)

	def OnData(self, slice):
		
		timekey = datetime(slice.Time).strftime("%Y%m%d %H:%M")
		
		if timekey in self.trade_data:
			rows = self.trade_data[timekey]
			
			# TODO delete the key from self.trade_data, just in case OnData fires
			# more than once within same minute.
			
			for row in rows:
				symbol = row[1]
				strategy = row[2]
				
				
				self.Log(repr(self.Identity(symbol).ToString()))
				
				# cash = min(self.Portfolio.CashBook.Values[0].Amount, 100000)
				# stock_price = self.Identity(symbol).Current.Value
				# self.Log(str(stock_price))
				
				# qty = int(Decimal(cash) // Decimal(stock_price)) / 2
				

				if strategy == "bull":
					newTicket = self.MarketOrder(symbol, 100, asynchronous = False)

				elif strategy == "bear":
					newTicket = self.MarketOrder(symbol, -100, asynchronous = False)
					# self.SetHoldings(symbol, -0.5)