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

class InterFundAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 4)
        self.SetCash(30000)
        
        # Risk Management
        self.hwm = self.Portfolio.TotalPortfolioValue
        self.max_dd = 1000
        
        ### 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 = .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
        ### }
        

    def Close(self):
        # 60:40
        self.sixty_forty_rebalance = True
        
        ### Treasury
        self.Liquidate(self.tlt)
        

    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