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(30000)
        
        # Risk Management
        self.hwm = self.Portfolio.TotalPortfolioValue
        self.max_dd = 1500
        
        ### Treasury Strategy {
        self.TS_AR = .6 # 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 = .2
        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)
        ### }
        
    
    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')
        ### }

    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)
        ### }

    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
    ### }
    
    
    ### 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)
    ### }
    
    
### 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
### }