# region imports
from AlgorithmImports import *
# endregion

# Strategy for 30m timeframe
    # RSI 14 (threshold on 25)
    # BB  30, stdv 2
    # BBwidth = (upper - lower) / middle
    # ATR 14

    # Conditions:
    # Previous candle = Close < Lower BB
    # Previous candle = RSI < 25
    # Current candle is green and closes above high of previous candle
    # BBwidth > 0.15

    # Exit:
    # Timeout after 2 bars OR TP: 1.2atr SL: 1.6atr
    
class SeriousBBRSI(QCAlgorithm): 
    def Initialize(self):
        # Set the start and end dates for backtesting
        self.SetStartDate(2024, 4, 1) 
        self.SetEndDate(2024, 5, 1)

        # Set the initial cash for the algorithm
        self.SetCash(10000) 

        # Add the SPY asset with minute resolution
        self.security = self.AddEquity("SPY", Resolution.Minute) 

        # Consolidate the SPY data into 30-minute bars and call OnDataConsolidated method for each consolidated bar
        self.Consolidate("SPY", timedelta(minutes=30), self.OnDataConsolidated) 

        # Set the brokerage model to Interactive Brokers with a margin account
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin) 

        # Create a rolling window of 2 TradeBars for security
        self.securityWin = RollingWindow[TradeBar](2)

        # Set SPY as the benchmark
        self.SetBenchmark("SPY") 

        # Create a 14-period RSI indicator with minute resolution for security
        self.rsi30 = self.RSI(self.security.Symbol, 14, Resolution.Minute) 
        
        # Create a 30-period 2 stdev BB indicator with minute resolution for security
        self.bb30 = self.bb(self.security.Symbol, 30, 2, Resolution.Minute)

        # Create a rolling window of 2 TradeBars for RSI
        self.rsiWin = RollingWindow[IndicatorDataPoint](2)

        # Create a rolling window of 2 TradeBars for BB
        self.bbLowerWin = RollingWindow[float](30)

    def OnDataConsolidated(self, consolidated):
        # Create plot of 30m Close price
        self.Plot("TradePlot", "close", consolidated.Close)

        # Add the Close price of the consolidated data to securityWin
        self.securityWin.Add(consolidated)

        # Add the rsi new consolidated bar to the rolling window
        self.rsiWin.Add(self.rsi30.Current)

        # Add the bb new consolidated bar to the rolling window
        # self.bbWin.Add(self.bb30)
        self.bbLowerWin.Add(self.bb30.LowerBand.Current.Value)

        # If the RSI and the rolling window are not ready, return
        if not (self.rsi30.IsReady and self.bb30.IsReady and self.securityWin.IsReady and self.rsiWin.IsReady and self.bbLowerWin.IsReady): 
            return 

        # Get the right data for conditions
        Close = self.securityWin[0].Close
        previousClose = self.securityWin[1].Close
        current_rsi = self.rsi30.Current.Value
        previous_rsi = self.rsiWin[1].Value
        previous_bb = self.bbLowerWin[1]

        # Plot the RSI value
        self.Plot("30minRSI", "rsi30", current_rsi) 

        # Plot the BB value
        self.Plot("TradePlot", "MiddleBand", self.bb30.MiddleBand.Current.Value)
        self.Plot("TradePlot", "UpperBand", self.bb30.UpperBand.Current.Value)
        self.Plot("TradePlot", "LowerBand", self.bb30.LowerBand.Current.Value)

        # If the SPY is in the portfolio
        if self.Portfolio[self.security.Symbol].Invested: 
            # If the RSI is above 50, liquidate the SPY
            if current_rsi > 50: 
                self.Liquidate(self.security.Symbol) 
                self.Debug("Sell") 
        # If the RSI is below 25 AND close is below the bblower, buy the SPY
        elif (previous_rsi < 25 and previousClose < previous_bb):
            self.SetHoldings(self.security.Symbol, 1.0) 
            self.Debug("Buy")