# region imports
from AlgorithmImports import *
# endregion

class OpeningRangeBreakoutUniverseAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2024, 1, 1)
        #self.set_end_date(2024, 2, 1)
        self.set_cash(10_000_000)
        self.settings.automatic_indicator_warm_up = True
        self._selected = []

        # Set the parameters.
        self._universe_size = 1000
        self._indicator_period = 14 # days
        self._stop_loss_atr_distance = 0.5 # 0.1 => 10% of ATR
        self._stop_loss_risk_size = 0.005 # 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 

        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 and f.symbol != self._spy], key=lambda f: f.dollar_volume)[-self._universe_size:]])

        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) # 1 minute late to allow consolidated bars time to update.
        self.schedule.on(self.date_rules.every_day(self._spy), self.time_rules.before_market_close(self._spy, 1), self._exit)

    def on_securities_changed(self, changes):
        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]
        history = self.history(symbols, 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)
        # Filter 1: Select assets with abnormally high volume for the day so far. (Relative Volume > 100%)
        equities = [equity for equity in equities if equity.relative_volume and equity.relative_volume > 1]
        if not equities:
            return
        # Filter 2: Select the top 20 assets with the greatest Relative Volume.
        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})
        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']))
            #self.log(f"target quantity for {equity.symbol}: {(quantity * order['entry_price'] / self.portfolio.total_portfolio_value)}")
            quantity_limit = self.calculate_order_quantity(equity.symbol, 1/self._max_positions)
            #if abs(quantity) <= quantity_limit:
            #    self.debug(f"{self.time} Within limit: {equity.symbol}")
            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 = []