from AlgorithmImports import *
class GapUpAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        # 1) Basic settings
        self.set_start_date(2025, 5, 25)
        self.set_end_date(2025, 6, 8)
        self.set_cash(100000)
        # Store dictionaries for tracking
        self._open_prices: Dict[Symbol, float] = {}
        self._entry_prices: Dict[Symbol, float] = {}
        self._watchlist: List[Symbol] = []
        # Set brokerage / resolution / timezone / etc.
        self.set_brokerage_model(BrokerageName.DEFAULT, AccountType.MARGIN)
        self.universe_settings.resolution = Resolution.DAILY
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.set_time_zone(TimeZones.NEW_YORK)
        # Add a daily SPY to anchor our scheduling rules
        self.add_equity("SPY", Resolution.DAILY)
        # Add Universe with both coarse and fine selection functions
        self.add_universe(self.coarse_selection_function, self.fine_selection_function)
        # Schedule the morning scan at 0 minutes after market open (can adjust to 10 seconds if preferred)
        self.schedule.on(
            self.date_rules.every_day("SPY"),
            self.time_rules.after_market_open("SPY", 0),
            self.morning_scan
        )
        # Schedule liquidation before market close (1 minute before close)
        self.schedule.on(
            self.date_rules.every_day("SPY"),
            self.time_rules.before_market_close("SPY", 1),
            self.end_of_day_liquidation
        )
        # We want second-resolution updates for the watchlist
        self.set_security_initializer(lambda s: s.set_data_normalization_mode(DataNormalizationMode.RAW))
    def coarse_selection_function(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        # Filter US equities priced >= 0.95
        return [
            c.symbol for c in coarse
            if c.has_fundamental_data and c.price >= 0.95 and c.symbol.id.market == Market.USA
        ]
    def fine_selection_function(self, fine: List[FineFundamental]) -> List[Symbol]:
        # For this strategy, just return all symbols passed from coarse
        return [f.symbol for f in fine]
    def morning_scan(self) -> None:
        # 2) Morning scan: compute gap for all active securities vs previous close
        if len(self.active_securities) == 0:
            self.debug("No active securities in universe for morning scan.")
            return
        symbols = [sec.symbol for sec in self.active_securities.values]
        self.debug(f"Morning scan with {len(symbols)} symbols")
        if not symbols:
            return
        # Get last 2 daily bars for each symbol to compute gap
        history = self.history(symbols, 2, Resolution.DAILY)
        if history.empty:
            self.debug("No historical data returned.")
            return
        gap_info = []
        # Handle multi-index data from pandas
        unique_syms = set(history.index.get_level_values(0))
        for sym in symbols:
            if sym not in unique_syms:
                continue
            sym_hist = history.xs(sym, level=0)
            if len(sym_hist) < 2:
                continue
            # Last two daily bars
            prev_bar = sym_hist.iloc[-2]
            curr_bar = sym_hist.iloc[-1]
            prev_close = float(prev_bar["close"])
            today_open = float(curr_bar["open"])
            if prev_close <= 0:
                continue
            gap_percent = (today_open - prev_close) / prev_close
            # Only consider if open >= $0.95
            if today_open >= 0.95:
                gap_info.append((sym, gap_percent, today_open))
        # Sort descending by gap
        gap_info.sort(key=lambda x: x[1], reverse=True)
        # Select top 100
        top_gap = gap_info[:10]
        # Store watchlist and open prices
        self._watchlist = [x[0] for x in top_gap]
        self._open_prices = {x[0]: x[2] for x in top_gap}
        self.debug(f"Selected top gap up stocks: {len(self._watchlist)}")
        # Subscribe for second data
        for symbol, _, _ in top_gap:
            self.add_equity(symbol.value, Resolution.SECOND)
    def on_data(self, slice: Slice) -> None:
        # 3) Per-second monitoring
        if not self._watchlist:
            return
        for symbol in self._watchlist:
            bar = slice.bars[symbol] if symbol in slice.bars else None
            if bar is None:
                continue
            current_price = bar.close
            open_price = self._open_prices.get(symbol, 0.0)
            if open_price <= 0:
                continue
            invested = self.portfolio[symbol].invested
            # If not invested, check if current price >= 1.4 * open
            if not invested and current_price >= 1.4 * open_price:
                quantity = self.calculate_order_quantity(symbol, 0.01)
                if quantity != 0:
                    self.market_order(symbol, quantity)
                    self._entry_prices[symbol] = current_price
                    self.debug(f"BUY {symbol} at {current_price}")
            # If invested, check for 30% up or down from entry
            if invested:
                entry_price = self._entry_prices.get(symbol, 0.0)
                if entry_price > 0:
                    # Up 30%?
                    if current_price >= 1.3 * entry_price:
                        self.liquidate(symbol)
                        self.debug(f"TAKE PROFIT {symbol} at {current_price}")
                    # Down 30%?
                    elif current_price <= 0.7 * entry_price:
                        self.liquidate(symbol)
                        self.debug(f"STOP LOSS {symbol} at {current_price}")
    def end_of_day_liquidation(self) -> None:
        # 5) Liquidate positions and reset daily state
        self.debug("EndOfDayLiquidation: Liquidating all positions.")
        self.liquidate()
        self._watchlist.clear()
        self._open_prices.clear()
        self._entry_prices.clear()
    def on_securities_changed(self, changes: SecurityChanges) -> None:
        # 7) Event triggered when securities are added/removed
        pass