from AlgorithmImports import *

class QQQORBTest(QCAlgorithm):
    """
    Opening Range Breakout (ORB) Day Trading Strategy for QQQ.
    
    Reference:
        - https://papers.ssrn.com/sol3/papers.cfm?abstract_id=4729284

    TLDR
        - This strategy identifies the breakout from the opening range during 
          the first 5 minutes of the trading session and takes positions accordingly.

    Strategy:
        - If the market moves up in the first 5 minutes, a bullish position is taken.
        - If the market moves down in the first 5 minutes, a bearish position is taken.
        - No positions are taken if the first 5-minute candle is a doji (open = close).
        - Stop loss is set at the low of the first 5-minute candle for long trades, and at the high for short trades.
        - The profit target is 10 times the risk (distance between entry price and stop).
        - If the profit target is not reached by the end of the day, the position is liquidated at market closure.
        - Trading size is calibrated to risk 1% of the capital per trade.

    Parameters:
        - Starting capital: $25,000
        - Maximum leverage: 4x
        - Commission: $0.0005 per share
        - Risk per trade: 1% of capital (ie: if stop loss hit)
        
    Contact
        - u/shock_and_awful 
    """
    
    ## System method, algo entry point
    ## -------------------------------
    def Initialize(self):
        self.InitParams()
        self.InitBacktest()
        self.InitData()
        self.InitIndicators()
        self.ScheduleRoutines()

    ## Initialize Parameters
    ## ----------------------
    def InitParams(self):
        self.risk_per_trade     = 0.01      # 1% risk per trade
        self.max_leverage       = 4         # in line with US FINRA regulations
        self.fixed_commission   = 0.0005    # fixed per trade
        self.minsAfterOpen      = 6         # Daily trade time: 6 mins after open (after first 5 min bar has closed)
        self.rrMultiplier       = float(self.get_parameter("rrMultiplier"))    # Risk reward multiplier for take profit
        self.rVolPeriod         = int(self.get_parameter("rVolPeriod"))
        self.rVolMagnitude      = float(self.get_parameter("rVolMagnitude"))
    ## Backtest properties
    ## -------------------
    def InitBacktest(self):
        self.SetStartDate(2018,1,1)  # Set Start Date
        # self.SetEndDate(2023,2,17)    # Set End Date
        self.SetCash(25000)
        
    ## Subscribe to asset data feed on right timeframe(s)
    ## --------------------------------------------------
    def InitData(self):
        self.ticker     = "QQQ"
        self.security   = self.AddEquity(self.ticker, Resolution.Minute)
        self.symbol     = self.security.Symbol
        self.SetBenchmark(self.ticker)

        # Set Brokerage model to IBkr
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # Set the feed model to fixed commission
        self.security.SetFeeModel(ConstantFeeModel(self.fixed_commission))

        # Set the margin model to PDT
        self.security.margin_model = PatternDayTradingMarginModel()

        # Initialize time frame (Five minute consolidator)
        self.fiveMinConsolidator = TradeBarConsolidator(timedelta(minutes = 5)) 
        self.fiveMinConsolidator.DataConsolidated += self.onFiveMinBarFormed
        self.SubscriptionManager.AddConsolidator(self.symbol, self.fiveMinConsolidator) 

        # State Tracking
        self.direction      = None
        self.stop_price     = 0
        self.entry_price    = 0
        self.highVolSpike   = False

    ## Handle 5 min bar consolidation
    ## -------------------------------
    def onFiveMinBarFormed(self, sender, bar):
        self.lastFiveMinBar = bar

    ## Initialize necessary indicators
    ## -------------------------------
    def InitIndicators(self):
        self.openingVolMA = SimpleMovingAverage(self.rVolPeriod)

    ## Schedule our 'chron jobs'
    ## -------------------------
    def ScheduleRoutines(self):
        ## Intraday selection
        self.Schedule.On(self.DateRules.EveryDay(), 
                         self.TimeRules.AfterMarketOpen(self.ticker, self.minsAfterOpen), 
                         self.RunAfterMarketOpen)                

        ## End of Day Liquidation
        self.Schedule.On(self.DateRules.EveryDay(), 
                         self.TimeRules.BeforeMarketClose(self.ticker, 2), 
                         self.LiquidateAtEoD)         


    ## A convenient 'Chron job', called on the 6th minute of every trading day.
    ## ------------------------------------------------------------------------
    def RunAfterMarketOpen(self):
        
        
        

        # Before updating with current bar volume, 
        # check if its a volume spike
        if(self.openingVolMA.is_ready):
            if (self.lastFiveMinBar.volume > (self.rVolMagnitude * self.openingVolMA.current.value)):
                self.highVolSpike = True
            else:
                self.highVolSpike = False
        
        self.openingVolMA.update(self.lastFiveMinBar.Time, self.lastFiveMinBar.volume)

        if(not self.openingVolMA.is_ready):
            return

        if(not self.highVolSpike):
            return

        openingBar     = self.lastFiveMinBar
        entry_price    = self.current_slice[self.symbol].open
        coefficient    = 1     
        
        # Doji candle. Don't trade.
        if( openingBar.close == openingBar.open ):
            self.Log("Doji Candle")
            return

        # Set direction based on bullish / bearish candle
        self.direction  = PortfolioBias.LONG if (openingBar.close > openingBar.open) else PortfolioBias.SHORT
        
        # Go Long if price hasnt already moved past stop
        if (self.direction == PortfolioBias.LONG) and \
             (entry_price > openingBar.low) :
            self.stop_price = openingBar.low 
            self.tp_price   = entry_price + (self.rrMultiplier * abs(entry_price - self.stop_price))
                                                
        # Go Short if price hasnt already moved past stop
        elif (self.direction == PortfolioBias.SHORT) and \
             (entry_price < openingBar.high) :
            self.stop_price = openingBar.high
            self.tp_price   = entry_price - (self.rrMultiplier * abs(entry_price - self.stop_price))
            coefficient = -1
                    
        num_shares = coefficient * self.CalcOrderSize(entry_price, self.stop_price)

        order_note = f"Opening position. going { 'long' if (self.direction == PortfolioBias.LONG) else 'short'}"  
        self.MarketOrder(self.symbol, num_shares, tag=order_note)
        self.Log("Order placed")


    ## System method, called as every new bar of data arrives. Check for exits (stop loss)
    ## -----------------------------------------------------------------------------------
    def OnData(self, data):        

        # If data is available and we are invested
        if ((self.ticker in data ) and (data[self.ticker] is not None)) and \
           (self.Portfolio.Invested):
            
            self.CheckForExits()
    

    ## Check if take profit or stop loss hit 
    ## -------------------------------------
    def CheckForExits(self):
        current_price = self.Securities[self.symbol].Price

        if (self.direction == PortfolioBias.LONG):
            if(current_price >= self.tp_price):
                self.LiquidateWithMsg("Take Profit")
            elif(current_price <= self.stop_price):
                self.LiquidateWithMsg("Stop Loss")

        elif (self.direction == PortfolioBias.SHORT):
            if(current_price <= self.tp_price):
                self.LiquidateWithMsg("Take Profit")
            elif(current_price >= self.stop_price):
                self.LiquidateWithMsg("Stop Loss")


    ## Calculate position size
    ## ---–---–---–---–---–---
    def CalcOrderSize(self, entry_price, stop_price):
        
        # Calculate the risk per share 
        risk_per_share = abs(entry_price - stop_price)
        
        # Get account size.
        # Currently using available cash but might consider using self.Portfolio.TotalPortfolioValue
        acct_size = self.Portfolio.Cash / 2

        # Calculate the total risk for the trade
        total_risk = acct_size * self.risk_per_trade
        
        # Calculate the number of shares to buy/sell
        shares = int(total_risk / risk_per_share)
        
        # Alt position size, adjust for leverage
        # with modifciation: multiply by risk per trade
        max_shares = int(acct_size * self.max_leverage / entry_price)

        return min(shares, max_shares)

    ## 'Chron Job' to Liquidate all holdings at the end of the day
    ## -----------------------------------------------------------
    def LiquidateAtEoD(self):
        self.LiquidateWithMsg("End-of-Day")
        

    ## Convenience method to liquidate with a message
    ## ----------------------------------------------
    def LiquidateWithMsg(self, exitReason):
        
        pnl         = round(100* self.Portfolio[self.symbol].UnrealizedProfitPercent,4)
        biasText    = 'Long' if (self.direction == PortfolioBias.LONG) else 'short'
        winlossText = 'win' if pnl > 0 else 'loss'        
        
        self.Liquidate(tag=f"{exitReason} Exiting {biasText} position with {pnl}% {winlossText}")