from AlgorithmImports import *


class PivotBasedTradingAlgorithm(QCAlgorithm):

    

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        # self.SetEndDate(2020, 1, 1)
        self.SetCash(100000)

        self.Debug(f"The algorithm time zone is set to: {self.TimeZone}")
        self.SetTimeZone("America/New_York")

        # Set brokerage model
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # Add symbol with 5-minute resolution
        # Add symbol
        self.symbol = self.AddEquity("QQQ", Resolution.Minute).Symbol
        
        # Create our consolidator with a period of 30 mins
        consolidator = TradeBarConsolidator(timedelta(minutes=5))
        consolidator.DataConsolidated += self.OnDataConsolidated
        self.SubscriptionManager.AddConsolidator(self.symbol, consolidator)

        # Get the last calendar day worth of daa at the specified resolution
        self.tradeBarHistory = self.History([self.symbol], timedelta(1), Resolution.Daily)


        # Schedule CalculatePivotPoints to run every trading day before market open
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), 
                         self.TimeRules.AfterMarketOpen(self.symbol, -10), 
                         self.CalculatePivotPoints)

        # Schedule function to liquidate 15 minutes before market close
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), 
                         self.TimeRules.BeforeMarketClose(self.symbol, 15), 
                         self.LiquidatePosition)

        self.previousDayHigh = None
        self.previousDayLow = None
        self.previousDayClose = None

        # Initialize pivot point attributes
        self.r1 = None
        self.r2 = None
        self.s1 = None
        self.s2 = None

    def OnDataConsolidated(self, sender, bar):

        '''
        This is called for each 5-minute bar of data. If the pivot points haven't been calculated yet (they are None), it returns early. 
        If it's the first 5-minute bar of the day (9:30 AM), it calculates the pivot points. 
        The trading logic checks if the current bar's closing price is greater than r1 to enter a long position or if it's less than r1 to exit the position.
        '''

        # Calculate Pivot Points at the start of the trading day
        if bar.Time.hour == 9 and bar.Time.minute == 30:
            self.CalculatePivotPoints()

        self.Debug(f"Processing bar for {bar.EndTime}: Close={bar.Close}")
        # Skip if pivot points haven't been calculated yet
        if self.r1 is None or self.r2 is None or self.s1 is None or self.s2 is None:
            self.Debug(f"Pivot points not yet calculated {self.Time}")
            return
        
        self.Debug(f"Pivot points: R1={self.r1}, R2={self.r2}, S1={self.s1}, S2={self.s2}")

        # Trading logic
        if bar.Close > self.r1 and not self.Portfolio.Invested:
            self.Debug(f"Signal to go Long: 5m Close: {bar.Close} > R1: {self.r1}")
            self.SetHoldings(self.symbol, 1)  # Long entry
        elif self.Portfolio.Invested and bar.Close < self.r1:
            self.Debug(f"Signal to go Close: 5m Close: {bar.Close} < R1: {self.r1}")
            self.Liquidate()  # Exit position

    def CalculatePivotPoints(self):
        '''
        Calculates the pivot points using the previous day's high, low, and close prices.
        This method is called at the beginning of the trading day to set the r1, r2, s1, and s2 values.
        '''
        # Fetch historical data for the previous trading day
        tradeBarHistory = self.tradeBarHistory

        self.Debug(f"History Count: {len(tradeBarHistory)}")

        if not tradeBarHistory.empty:

            self.previousDayHigh = tradeBarHistory['high'][0]
            self.previousDayLow = tradeBarHistory['low'][0]
            self.previousDayClose = tradeBarHistory['close'][0]
            self.Debug(f"Stored previous day's OHLC: High={self.previousDayHigh}, Low={self.previousDayLow}, Close={self.previousDayClose}")
        else:
            self.Debug("History data is empty. Cannot store previous day's OHLC.")
            

        self.Debug(f"Previous High: {self.previousDayHigh}, Low: {self.previousDayLow}, Close: {self.previousDayClose}")
        if self.previousDayHigh is not None and self.previousDayLow is not None and self.previousDayClose is not None:
            p = (self.previousDayHigh + self.previousDayLow + self.previousDayClose) / 3
            self.r1 = 2 * p - self.previousDayLow
            self.r2 = p + (self.previousDayHigh - self.previousDayLow)
            self.s1 = 2 * p - self.previousDayHigh
            self.s2 = p - (self.previousDayHigh - self.previousDayLow)
            self.Debug(f"Calculated Pivot Points - R1: {self.r1}, R2: {self.r2}, S1: {self.s1}, S2: {self.s2}")
        else:
            self.Debug("Previous day OHLC values not set. Cannot calculate pivot points.")

    def LiquidatePosition(self):
        '''
        Liquidates any invested position 15 minutes before the market close.
        '''
        # Liquidate 15 minutes before market close
        if self.Portfolio.Invested:
            self.Liquidate(self.symbol)