# region imports
from AlgorithmImports import *
# endregion

class ORB_v1(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2024, 1, 1)
        self.set_end_date(2024,12, 31)
        self.set_cash(1_000_000)
        self.settings.automatic_indicator_warm_up = True # pre-loading 28 days indicator before 
        self._selected = []

        # Set the parameters.
        self._universe_size = 1000   #question on this limit?
        self._indicator_period = 14 # days  
        self._stop_loss_atr_distance = 0.5 # 0.1 => 10% of ATR
        self._stop_loss_risk_size = 0.01 # 0.01 => Lose 1% of the portfolio if stop loss is hit
        self._max_positions = 20
        self._opening_range_minutes = 5  
        self._leverage = 4

        # Add SPY so there is at least 1 asset at minute resolution to step the algorithm along.
        self._spy = self.add_equity('SPY').symbol   # use as the time benchmark, will not participate in actual trading

        self.universe_settings.resolution = Resolution.DAILY
        self.universe_settings.schedule.on(self.date_rules.month_start(self._spy))

    
        self._universe = self.add_universe(
            lambda fundamentals: [
                f.symbol 
                for f in sorted(
                    [
                        f for f in fundamentals 
                        if f.price > 5  # price > 5
                        and f.symbol != self._spy
                        and f.volume > 1_000_000  # Volume greater than 1 million
                    ], 
                    key=lambda f: f.dollar_volume
                )[-self._universe_size:]
            ]
        )

        # run 5 min after the market open
        self.schedule.on(
            self.date_rules.every_day(self._spy), 
            self.time_rules.after_market_open(self._spy, self._opening_range_minutes), 
            self._scan_for_entries
        )  

        # exit 1 min before market close
        self.schedule.on(
            self.date_rules.every_day(self._spy), 
            self.time_rules.before_market_close(self._spy, 1), 
            self._exit
        )

        self.set_warm_up(timedelta(2 * self._indicator_period))

        # Set Benchmark
        self.SetBenchmark("SPY")

        # Variable to hold the last calculated benchmark value
        self.lastBenchmarkValue = None
        

        # Our inital benchmark value scaled to match our portfolio
        # dont know why the inital value of SP 500 is 1047215.721 at 2024/1/1, I scale back by just deducting a constant, still have some minor gap but should be negligible in the long term
        # need to adjust for the constant depends on the initial gap
        self.BenchmarkPerformance = self.Portfolio.TotalPortfolioValue - 47000





    def on_securities_changed(self, changes):  # updating ATR parameter for securuity have changes (we set daily change so everyday this function should run one time only)
        for security in changes.added_securities:
            security.atr = self.atr(security.symbol, self._indicator_period, resolution=Resolution.DAILY)
            security.volume_sma = SimpleMovingAverage(self._indicator_period)

        
    def _scan_for_entries(self):
        symbols = list(self._universe.selected)
        equities = [self.securities[symbol] for symbol in symbols]  


        # Filter 1: ATR> 0.5
        equities = [equity for equity in equities if equity.atr.is_ready and equity.atr.current.value > 0.5]
        if not equities:
            return

        
        # Filter 2: Relative Volume > 100%
        history = self.history([equity.symbol for equity in equities], 5, Resolution.MINUTE)
        volume_sum = history.volume.unstack(0).sum()
        equities = [equity for equity in equities if equity.symbol in volume_sum]

        for equity in equities:
            volume = volume_sum.loc[equity.symbol]
            equity.relative_volume = volume / equity.volume_sma.current.value if equity.volume_sma.is_ready else None
            equity.volume_sma.update(self.time, volume)
        if self.is_warming_up:
            return

        equities = [equity for equity in equities if equity.relative_volume and equity.relative_volume > 1]
        if not equities:
            return


        # Filter 3: top 20 assets with largest RV
        equities = sorted(equities, key=lambda equity: equity.relative_volume)[-self._max_positions:]
        
        history = history.loc[[equity.symbol for equity in equities]]
        open_by_symbol = history.open.unstack(0).iloc[0]
        close_by_symbol = history.close.unstack(0).iloc[-1]
        high_by_symbol = history.high.unstack(0).max()
        low_by_symbol = history.low.unstack(0).min()

        # Create orders for the target assets.
        # Calculate position sizes so that if you fill an order at the high (low) of the first 5-minute bar 
        # and hit a stop loss based on 10% of the ATR, you only lose x% of portfolio value.    
        orders = []

        for symbol in close_by_symbol[close_by_symbol > open_by_symbol].index:
            equity = self.securities[symbol]
            orders.append({
                'equity': equity, 
                'entry_price': high_by_symbol.loc[equity.symbol], 
                'stop_price': high_by_symbol.loc[equity.symbol] - self._stop_loss_atr_distance * equity.atr.current.value  # stop price
            })


        for symbol in close_by_symbol[close_by_symbol < open_by_symbol].index:
            equity = self.securities[symbol]
            orders.append({
                'equity': equity, 
                'entry_price': low_by_symbol.loc[equity.symbol], 
                'stop_price': low_by_symbol.loc[equity.symbol] + self._stop_loss_atr_distance * equity.atr.current.value
        })


        for order in orders:
            equity = order['equity']
            self._selected.append(equity)
            self._create_empty_order_tickets(equity)
            self.add_security(equity.symbol, leverage=self._leverage)
            quantity = int((self._stop_loss_risk_size * self.portfolio.total_portfolio_value / self._max_positions) / (order['entry_price'] - order['stop_price']))
            quantity_limit = self.calculate_order_quantity(equity.symbol, 1/self._max_positions)
            quantity = min(abs(quantity), quantity_limit) * np.sign(quantity)
            if quantity:
                equity.stop_loss_price = order['stop_price']
                equity.entry_ticket = self.stop_market_order(equity.symbol, quantity, order['entry_price'], tag='Entry')

    def on_order_event(self, order_event: OrderEvent) -> None:
        if order_event.status != OrderStatus.FILLED:
            return
        security = self.securities[order_event.symbol]
        # When the entry order is hit, place the exit order: Stop loss based on ATR.
        if order_event.ticket == security.entry_ticket:
            security.stop_loss_ticket = self.stop_market_order(order_event.symbol, -security.entry_ticket.quantity, security.stop_loss_price, tag='ATR Stop')
        # When the stop loss order is hit, cancel the MOC order.
        elif order_event.ticket == security.stop_loss_ticket:
            self._create_empty_order_tickets(security)

    # Create some members on the Equity object to store each order ticket.
    def _create_empty_order_tickets(self, equity):
        equity.entry_ticket = None
        equity.stop_loss_ticket = None

    # Liquidate the portfolio, remove order tickets, remove the minute-resolution data subscriptions.
    def _exit(self):
        self.liquidate()
        for equity in self._selected:
            self._create_empty_order_tickets(equity)
            self.remove_security(equity.symbol)
        self._selected = []

            
    def OnData(self, data):
       
        # Get the current close price for the benchmark (SPY)
        benchmark = self.Securities["SPY"].Close
  
        # If we had a previous close, update our benchmark performance.
        if self.lastBenchmarkValue is not None and self.lastBenchmarkValue != 0:
            self.BenchmarkPerformance = self.BenchmarkPerformance * (benchmark / self.lastBenchmarkValue)
  
        # Save the current close price for the next update
        self.lastBenchmarkValue = benchmark

        self.Plot("Strategy vs Benchmark", "Portfolio Value", self.Portfolio.TotalPortfolioValue)
        self.Plot("Strategy vs Benchmark", "Benchmark", self.BenchmarkPerformance)