from QuantConnect.Data.Market import TradeBar
from datetime import timedelta
from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
import decimal as d


class MyAlgorithm(QCAlgorithm):
	
    def Initialize(self):
        self.SetStartDate(2013, 05, 1)  # Set Start Date
        self.SetEndDate(2013, 06, 01)
        self.SetCash(100000)  # Set Strategy Cash
        
        self.symbolData = dict()
        
        for ticker in ["SPY", "FB", "TWTR"]:
			symbol = self.AddEquity(ticker, Resolution.Second).Symbol

			consolidator_daily = TradeBarConsolidator(timedelta(1))
			consolidator_daily.DataConsolidated += self.OnDailyData
			self.SubscriptionManager.AddConsolidator(symbol, consolidator_daily)

			consolidator_minute = TradeBarConsolidator(60)
			consolidator_minute.DataConsolidated += self.OnMinuteData
			self.SubscriptionManager.AddConsolidator(symbol, consolidator_minute)
			
			self.symbolData[symbol] = SymbolData()

        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.AfterMarketOpen('SPY', 2),
                         Action(self.one_minute_after_open_market))
                         
        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.BeforeMarketClose('SPY', 1),
                         Action(self.before_close_market))

    # Add daily bar to daily rolling window
    def OnDailyData(self, sender, bar):
        self.symbolData[bar.Symbol].daily_rw.Add(bar)

    def OnMinuteData(self, sender, bar):
        self.symbolData[bar.Symbol].minute_rw.Add(bar)

    def one_minute_after_open_market(self):
        """
        At 9:31 check if there has been a gap at the market open from the previous day.
        If so and the stock is gapping up and the first minute bar is negative, create a short selling signal.
        If the stock is gapping down and the first minute bar is positive, create a buying signal.
        """

        
        
        for k in self.symbolData: 
         if not (self.symbolData[k].window.IsReady and self.symbolData[k].daily_rw.IsReady and self.symbolData[k].minute_rw.IsReady): return
         last_close = self.symbolData[k].window[0].Close
         yesterday_daily_close = self.symbolData[k].daily_rw[1].Close
         first_minute_close = self.symbolData[k].minute_rw[1].Close
         first_minute_open = self.symbolData[k].minute_rw[1].Open
        
         gap = last_close - yesterday_daily_close
         first_minute_bar = first_minute_close - first_minute_open

         if not self.Portfolio[k].Invested:
            # If the stock is gapping down and the first minute bar is positive, create a buying signal.
            if gap < 0 and first_minute_bar > 0:
                self.SetHoldings(k, 0.333)
                self.Log('GOING LONG')
            # If the stock is gapping up and the first minute bar is negative, create a short selling signal
            elif gap > 0 and first_minute_bar < 0:
                self.SetHoldings(k, -0.333)
                self.Log('GOING SHORT')

    def before_close_market(self):
    	"""
    	At the end of the day, if there is a short position, close it.
    	"""
    	for k in self.symbolData:
    	 if self.Portfolio[k].Invested:
             self.Liquidate(k)
             self.Log('LIQUIDATE SHORT End of Day')


    def OnData(self, data):
        if data["SPY"] is None:
            return
        for k in self.symbolData:
         if not self.symbolData[k].window.IsReady:
            return
     	 if self.Portfolio[k].Invested:
          factor = d.Decimal(1.01)
          currBar = self.symbolData[stock].window[0].Close
        # Every second, check the price and if it's higher than the price the stock was bought for times 1.01, close the position.
          if self.Portfolio[k].AveragePrice * factor < currBar:
            self.Liquidate(k)
            self.Log('LIQUIDATE AT THRESHOLD REACHED.')
    
    def OnEndOfDay(self):
    	self.Plot("Portfolio", "MarginRemaining", self.Portfolio.MarginRemaining)
    
    def OnEndOfAlgorithm(self):
    	self.Liquidate()
    	self.Log('LIQUIDATE AT End Of Algorithm.')
    	
class SymbolData(object):
    		
    def __init__(self):
		self.daily_rw = RollingWindow[TradeBar](2)
		self.minute_rw = RollingWindow[TradeBar](2)
		self.window = RollingWindow[TradeBar](2)