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, 4)
        self.SetEndDate(2020, 7, 7)
        self.SetCash(30000)
        
        # Risk Management
        self.hwm = self.Portfolio.TotalPortfolioValue
        self.max_dd = 8000
        
        ### Treasury Strategy {
        self.TS_AR = .5 # 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.Close)
        self.Schedule.On(self.DateRules.EveryDay(self.tlt), self.TimeRules.AfterMarketOpen(self.tlt, 1), self.Rebalance)
        ### }
        
        ### 60:40 Strategy {
        self.SF_AR = .3
        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 = 0.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)
        ### }
        

    def Close(self):
        # 60:40
        self.sixty_forty_rebalance = True
        
        ### Treasury
        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
    ### }

    def Rebalance(self):
        ### Treasury 
        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)
        ### Treasury
        
    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 rebalancing
        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)
                    if quantity:
                        self.MarketOrder(ticker, quantity)
        self.sixty_forty_rebalance = False
        # 60:40
        
        
        # Turnaround Tuesday
        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:
                    self.MarketOrder(self.symbol, self.quantity)
            self.monday_open_price = 0
        
        elif self.tuesday_open:
            self.tuesday_open = False
            self.MarketOnCloseOrder(self.symbol, -self.quantity)
            self.quantity = 0