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):
        '''Initialise  the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''

        self.SetStartDate(2015, 05, 1)  # Set Start Date
        self.SetEndDate(2015, 07, 19)
        self.SetCash(100000)  # Set Strategy Cash
        self.AddEquity("SPY", Resolution.Second)

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

        consolidator_minute = TradeBarConsolidator(60)
        consolidator_minute.DataConsolidated += self.OnMinuteData
        self.SubscriptionManager.AddConsolidator("SPY", consolidator_minute)

        self.daily_rw = RollingWindow[TradeBar](2)
        self.minute_rw = RollingWindow[TradeBar](2)
        self.window = RollingWindow[TradeBar](2)

        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.At(9, 31),
                         Action(self.one_minute_after_open_market))

        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.At(15, 59, 56),
                         Action(self.before_close_market))

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

    def OnMinuteData(self, sender, bar):
        self.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.
        """
        if not (self.window.IsReady and self.daily_rw.IsReady and self.minute_rw.IsReady): return
        last_close = self.window[0].Close
        yesterday_daily_close = self.daily_rw[1].Close
        first_minute_close = self.minute_rw[1].Close
        first_minute_open = self.minute_rw[1].Open
        
        gap = last_close - yesterday_daily_close
        first_minute_bar = first_minute_close - first_minute_open

        if not self.Portfolio["SPY"].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("SPY", 1)
                self.Debug('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("SPY", -1)
                self.Debug('GOING SHORT')

    def before_close_market(self):
        if self.Portfolio["SPY"].IsShort:
            self.Liquidate("SPY")
            self.Debug('LIQUIDATE SHORT End of Day')

    # Add second bar to window rolling window
    def OnData(self, data):
        if data["SPY"] is None:
            return
        self.window.Add(data["SPY"])
        if not (self.window.IsReady):
            return
        # self.Debug("haha")
        factor = d.Decimal(1.01)

        currBar = self.window[0].Close

        if self.Portfolio["SPY"].AveragePrice * factor < currBar:
            self.Liquidate("SPY")
            self.Debug('LIQUIDATE AT THRESHOLD REACHED.')