import pandas as pd 
from pandas.tseries.offsets import BDay
from pandas.tseries.offsets import BMonthEnd

class InternFundAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 18)
        self.SetCash(2000)
        
        # Risk Management
        self.hwm = self.Portfolio.TotalPortfolioValue
        self.max_dd = 500
        
        ### Treasury Strategy {
        self.TS_AR = .2 # allocation ratio
        self.tlt = self.AddEquity('TLT', Resolution.Minute).Symbol
        self.Schedule.On(self.DateRules.MonthEnd(self.tlt), self.TimeRules.BeforeMarketClose(self.tlt, 1), self.MonthlyClose)
        self.Schedule.On(self.DateRules.EveryDay(self.tlt), self.TimeRules.AfterMarketOpen(self.tlt, 1), self.EveryDayAfterMarketOpen)
        ### }
        
        ### 60:40 Strategy {
        self.SF_AR = .5
        self.weight_by_ticker = {'SPY': 0.6, 'AGG': 0.4, 'VXX': 0.1}
        self.sixty_forty_tickers = list(self.weight_by_ticker.keys())
        for ticker in self.sixty_forty_tickers:
            self.AddEquity(ticker, Resolution.Minute)
        self.sixty_forty_rebalance = True
        ### }
        
        
        ### Turnaround Tuesday Strategy {
        self.TT_AR = .2
        self.spy = self.AddEquity("SPY", Resolution.Minute)
        self.symbol = self.spy.Symbol
        self.quantity = 0
        self.monday_open_price = 0
        self.monday_open = False
        self.monday_close = False
        self.tuesday_open = False
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.AfterMarketOpen("SPY", 0), self.SignalMondayOpen)
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.BeforeMarketClose("SPY", 0), self.SignalMondayClose)
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Tuesday), self.TimeRules.AfterMarketOpen("SPY", 0), self.SignalTuesdayOpen)
        ### }
        
        
        ### Reverse George Douglas Taylor Strategy {
        self.GDT_AR = .2
        tickers = ['IWM', 'QQQ']
        self.symbol_data_by_symbol = {}
        for ticker in tickers:
            symbol = self.AddEquity(ticker, Resolution.Minute).Symbol
            self.symbol_data_by_symbol[symbol] = SymbolData(symbol, self)
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 1), self.EveryDayBeforeMarketClose)
        ### }
        
        
        ### Leveraged ETFS {
        self.LE_AR = 0.05
        self.AddEquity('TQQQ', Resolution.Minute)    # 3x QQQ
        self.AddEquity('UBT', Resolution.Minute)     # 3x 20yr Treasury
        self.AddEquity('UST', Resolution.Minute)     # 3x 10yr Treasury
        self.LE_tickers = ['TQQQ', 'UBT', 'UST']
        self.weeks = 0
        self.Schedule.On(self.DateRules.WeekStart('UST'), self.TimeRules.AfterMarketOpen('UST', 150), self.LE_Rebalance)  
        ###}
        
        
        ### Bull Market Gap Down Rebound {
        self.GDR_AR = 0.3
        self.sma_short = SimpleMovingAverage(50)
        self.sma_long = SimpleMovingAverage(200)
        self.max = Maximum(50)
        closes = self.History(self.symbol, 200, Resolution.Daily).loc[self.symbol].close
        for time, close in closes.iteritems():
            self.sma_short.Update(time, close)
            self.sma_long.Update(time, close)
            self.max.Update(time, close)
        self.last_close = None
        self.rebound_quantity = 0
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), self.TimeRules.BeforeMarketClose(self.symbol, 15), self.FifteenBeforeClose)
        ### }
        
    
    def OnData(self, data):
        # Risk Management
        value = self.Portfolio.TotalPortfolioValue
        if value > self.hwm:
            self.hwm = value
        if self.hwm - value > self.max_dd:
            self.Debug("Max DD reached")
            self.Quit()
        
        
        ### 60:40 Strategy {
        if self.sixty_forty_rebalance:
            for ticker in self.sixty_forty_tickers:
                if data.ContainsKey(ticker): 
                    weight = self.weight_by_ticker[ticker]
                    quantity = self.CalculateOrderQuantity(ticker, weight * self.SF_AR / sum(self.weight_by_ticker.values()))
                    if quantity >= 1:
                        self.MarketOrder(ticker, quantity)
        self.sixty_forty_rebalance = False
        ### }
        
        
        ### Turnaround Tuesday Strategy {
        if not data.ContainsKey(self.symbol) or data[self.symbol] is None:
            return
        
        if self.monday_open:
            self.monday_open = False
            self.monday_open_price = data[self.symbol].Open
            
        elif self.monday_close:
            self.monday_close = False
            if data[self.symbol].Close < self.monday_open_price: # Monday is a down day
                self.quantity = self.CalculateOrderQuantity(self.symbol, self.TT_AR)
                if self.quantity >= 1:
                    self.MarketOrder(self.symbol, self.quantity)
            self.monday_open_price = 0
        
        elif self.tuesday_open:
            self.tuesday_open = False
            if self.quantity >= 1:
                self.MarketOnCloseOrder(self.symbol, -self.quantity)
            self.quantity = 0
        ### }
        

    def EveryDayAfterMarketOpen(self):
        ### Treasury Strategy {
        offset = BMonthEnd()
        last_day = offset.rollforward(self.Time)
        trigger_day = last_day - BDay(4)

        if self.Time == trigger_day:
            self.SetHoldings(self.tlt, self.TS_AR)
        ### }
        
        ### Reverse George Douglas Taylor Strategy {
        # Update days_held tracker for each symbol
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if self.Securities[symbol].Invested:
                symbol_data.days_held += 1
        self.make_orders('open')
        ### }
        
        ### Bull Market Gap Down Rebound {
        if self.last_close is None:
            return
        # Enter if:
        # - Gapped down this morning
        # - Yesterday closed at the highest close in the last 50 bars
        # - Short MA > Long MA
        if self.Securities[self.symbol].Open < self.last_close and \
            self.max.Current.Value == self.last_close and \
            self.sma_short > self.sma_long:
            self.rebound_quantity = self.CalculateOrderQuantity(self.symbol, self.GDR_AR)
            self.MarketOrder(self.symbol, self.rebound_quantity)
        ### }

    def EveryDayBeforeMarketClose(self):
        ### Reverse George Douglas Taylor Strategy {
        self.make_orders('close')
        
        # Exit orders with days held == 1
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if symbol_data.days_held == 1:
                symbol_data.days_held = 0
                self.Liquidate(symbol)
        ### }
    
    
    ### Bull Market Gap Down Rebound {
    def FifteenBeforeClose(self):
        if self.rebound_quantity > 0:
            self.MarketOrder(self.symbol, -self.rebound_quantity)
            self.rebound_quantity = 0
    ### }
        

    def MonthlyClose(self):
        ### 60:40 Strategy {
        self.sixty_forty_rebalance = True
        ### }
        
        ### Treasury Strategy {
        self.Liquidate(self.tlt)
        ### }
        
        
    ### Turnaround Tuesday {
    def SignalMondayOpen(self):
        if self.spy.IsTradable:
            self.monday_open = True
        
    def SignalMondayClose(self):
        if self.monday_open_price:
            self.monday_close = True
    
    def SignalTuesdayOpen(self):
        if self.quantity:
            self.tuesday_open = True
    ### }
    
    
    ### Bull Market Gap Down Rebound {
    def OnEndOfDay(self):
        self.last_close = self.Securities[self.symbol].Price
        
        self.sma_short.Update(self.Time, self.last_close)
        self.sma_long.Update(self.Time, self.last_close)
        self.max.Update(self.Time, self.last_close)
    ### }
    
    
    ### Reverse George Douglas Taylor Strategy {
    def make_orders(self, at):
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if at == 'open':
                price = self.Securities[symbol].Open
            else:
                price = self.Securities[symbol].Close
            signal = symbol_data.generate_signal(price, at)
            if signal:
                if self.Securities[symbol].Invested:
                    # Extend exit date
                    symbol_data.days_held -= 1
                else:
                    # Make order
                    quantity = self.CalculateOrderQuantity(symbol, self.GDT_AR / len(self.symbol_data_by_symbol))
                    if quantity >= 1:
                        self.MarketOrder(symbol, quantity)
    ### }
    
    ### Leveraged ETFs {
    def LE_Rebalance(self):
        if self.weeks % 2 == 0:
            alloc = self.LE_AR / len(self.LE_tickers)
            for ticker in self.LE_tickers:
                quantity = self.CalculateOrderQuantity(ticker, alloc)
                if quantity >= 1:
                    self.MarketOrder(ticker, quantity)
            
        self.weeks += 1
    ### }
    
### Reverse George Douglas Taylor Strategy {
class SymbolData:
    open_price = None
    days_held = 0

    def __init__(self, symbol, algorithm):
        self.symbol = symbol
        self.window = RollingWindow[TradeBar](3)
        
        # Warm up history
        history = algorithm.History(symbol, 3, Resolution.Daily)
        for idx, row in history.iterrows():
            tradebar = TradeBar(idx, symbol, row.open, row.high, row.low, row.close, row.volume)
            self.window.Add(tradebar)
        
        # Setup consolidator
        self.consolidator = TradeBarConsolidator(timedelta(1))
        self.consolidator.DataConsolidated += self.ConsolidationHandler
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        
    def ConsolidationHandler(self, sender, consolidated):
        self.window.Add(consolidated)
        
    def generate_signal(self, price, at):
        if at == 'open':
            self.open_price = price

        # If up two days in a row
        if self.window[2].Close < self.window[1].Close and self.window[1].Close < self.window[0].Close:
            # Gaps up the third day
            if at == 'open' and self.window[0].Close < self.open_price:
                return True # Short open
            
            # Doesn't gap up the third day, but the third day is an up day
            if at == 'close' and self.window[0].Close >= self.open_price and self.window[0].Close < price:
                return True # Short close
### }