from AlgorithmImports import *

class BondTradingAlgorithm(QCAlgorithm):
    def Initialize(self):
        # Set the start and end date for the backtest
        self.SetStartDate(2020, 1, 1)
        #self.SetEndDate(2024, 1, 1)
        
        # Set the initial cash balance
        self.SetCash(100000)
        
        # Add Treasury bond data
        self.bond = self.AddEquity("TLT", Resolution.Daily).Symbol
        
        # Schedule the trading events
        self.Schedule.On(self.DateRules.MonthEnd(self.bond), self.TimeRules.BeforeMarketClose(self.bond, 0), self.SellAtMonthEnd)
        self.Schedule.On(self.DateRules.MonthEnd(self.bond), self.TimeRules.BeforeMarketClose(self.bond, 0), self.SellShortAtMonthEnd)
        
        # Schedule the Buy and Cover events dynamically
        self.Schedule.On(self.DateRules.EveryDay(self.bond), self.TimeRules.BeforeMarketClose(self.bond, 0), self.CheckSeventhLastTradingDay)
        
    def CheckSeventhLastTradingDay(self):
        # Get the current date
        current_date = self.Time.date()
        
        # Get the last trading day of the current month
        last_trading_day = self.GetLastTradingDayOfMonth(current_date)
        
        # Calculate the seventh last trading day of the month
        seventh_last_trading_day = self.GetNthLastTradingDayOfMonth(current_date, 7)
        
        # Check if today is the seventh last trading day of the month
        if current_date == seventh_last_trading_day:
            self.BuyAtSeventhLastTradingDay()
        elif current_date == self.GetNthTradingDayOfMonth(current_date, 7):
            self.CoverAtSeventhLastTradingDay()
        
    def GetLastTradingDayOfMonth(self, date):
        # Get the last day of the month
        last_day = date.replace(day=1) + timedelta(days=32)
        last_day = last_day.replace(day=1) - timedelta(days=1)
        
        # Find the last trading day of the month
        while not self.Securities[self.bond].Exchange.DateIsOpen(last_day):
            last_day -= timedelta(days=1)
        
        return last_day
    
    def GetNthLastTradingDayOfMonth(self, date, n):
        # Get the last trading day of the month
        last_trading_day = self.GetLastTradingDayOfMonth(date)
        
        # Find the nth last trading day of the month
        count = 0
        while count < n:
            last_trading_day -= timedelta(days=1)
            if self.Securities[self.bond].Exchange.DateIsOpen(last_trading_day):
                count += 1
        
        return last_trading_day
    
    def GetNthTradingDayOfMonth(self, date, n):
        # Get the first day of the month
        first_day = date.replace(day=1)
        
        # Find the nth trading day of the month
        count = 0
        while count < n:
            if self.Securities[self.bond].Exchange.DateIsOpen(first_day):
                count += 1
            first_day += timedelta(days=1)
        
        return first_day - timedelta(days=1)
    
    def BuyAtSeventhLastTradingDay(self):
        # Buy at the close of the seventh last trading day of the month
        self.SetHoldings(self.bond, 1)
        self.Debug(f"Bought {self.bond} at {self.Time}")
        
    def SellAtMonthEnd(self):
        # Sell at the close of the last trading day of the month
        self.Liquidate(self.bond)
        self.Debug(f"Sold {self.bond} at {self.Time}")
        
    def SellShortAtMonthEnd(self):
        # Sell short at the close of the month
        self.SetHoldings(self.bond, -1)
        self.Debug(f"Sold short {self.bond} at {self.Time}")
        
    def CoverAtSeventhLastTradingDay(self):
        # Cover at the close of the seventh trading day of the month
        self.Liquidate(self.bond)
        self.Debug(f"Covered short {self.bond} at {self.Time}")