# region imports
from AlgorithmImports import *
# endregion

class Mindthegapv2(QCAlgorithm):

    def initialize(self):
        # Set the backtest time frame and initial capital
        self.set_start_date(2019, 2, 12)
        self.set_end_date(2019, 4, 11)
        self.set_cash(100_000)

        # Set warm-up period for indicators (2 days of daily data)
        self.set_warm_up(2, Resolution.DAILY)
        self.is_warming_up = True

        # Dictionary to hold SelectionData objects for each symbol
        self._selection_data_by_symbol = {}

        # List of selected symbols that meet SMA criteria
        self.selected_symbols = []

        # Set universe using SPY constituents, filtered by SMA criteria
        spy = Symbol.create("SPY", SecurityType.EQUITY, Market.USA)
        self._universe = self.add_universe(self.universe.etf(spy), self._fundamental_selection)
        self.UniverseSettings.Resolution = Resolution.MINUTE

        # Schedule trades: open shortly after market open, close before close
        self.schedule.on(self.date_rules.every_day(spy), self.time_rules.after_market_open(spy, 2), self.open_positions)
        self.schedule.on(self.date_rules.every_day(spy), self.time_rules.after_market_open(spy, 15), self.close_positions)

    def _fundamental_selection(self, fundamental: List[Fundamental]) -> List[Symbol]:
        """
        Universe selection method that filters symbols trading above their 100-day SMA.
        """
        current_symbols = [f.symbol for f in fundamental]

        for f in fundamental:
            symbol = f.symbol
            if symbol not in self._selection_data_by_symbol:
                self._selection_data_by_symbol[symbol] = SelectionData(self, symbol)

            # Update the SMA for each symbol
            self._selection_data_by_symbol[symbol].update_sma(self.time, f.adjusted_price)

        # Remove outdated symbols
        for symbol in list(self._selection_data_by_symbol):
            if symbol not in current_symbols:
                del self._selection_data_by_symbol[symbol]

        # Select symbols currently trading above their SMA
        selected = [
            symbol for symbol, data in self._selection_data_by_symbol.items()
            if data.is_above_sma
        ]

        #self.plot("Universe", "Possible", len(fundamental))
        #self.Debug(f"Universe Time {self.time}. Selected Symbols: {selected}")
        self.selected_symbols = selected
        return selected

    def on_data(self, data: Slice):
        """
        Called every minute. Updates ATR values for selected symbols after market open.
        """
        if self.is_warming_up or not (self.time.hour == 9 and self.time.minute == 31):
            return

        for symbol in self.selected_symbols:
            trade_bar = data.bars.get(symbol)
            if trade_bar:
                self._selection_data_by_symbol[symbol]._atr.update(trade_bar)

    def open_positions(self):
        """
        Open positions in symbols that gapped down >3% and meet the ATR threshold.
        """
        if self.is_warming_up:
            return

        # Filter out symbols with valid data and price > $10
        symbols = [
            s for s in self.selected_symbols
            if self.securities.contains_key(s)
            and self.securities[s].has_data
            and self.securities[s].price > 10
        ]

        if not symbols:
            self.Debug("No valid symbols to trade.")
            return

        # Get yesterday’s close prices
        history = self.history(symbols, 1, Resolution.DAILY)
        symbols_to_trade = []

        for symbol in symbols:
            if symbol in history.index.levels[0] and len(history.loc[symbol]) >= 1:
                yesterday_close = history.loc[symbol].iloc[0]['close']
                current_price = self.securities[symbol].price
                atr = self._selection_data_by_symbol[symbol]._atr.current.value
                gap = yesterday_close - current_price

                #self.Debug(f"{symbol}: Close={yesterday_close}, Current={current_price}, ATR={atr}, Gap={gap}")

                # ATR-based entry threshold
                if atr > 0 and gap >= 1.2 * atr:
                    symbols_to_trade.append(symbol)
                    #self.Debug(f"{symbol} passed ATR-based entry condition.")

        # Sort symbols by ATR value descending, take top 20
        symbols_to_trade = sorted(
            symbols_to_trade,
            key=lambda x: self._selection_data_by_symbol[x]._atr.current.value,
            reverse=True
        )[:5]

        if not symbols_to_trade:
            self.Debug("No symbols passed gap down and ATR filter.")
            return

        # Equal weight allocation
        weight = 1 / len(symbols_to_trade)
        self.set_holdings([PortfolioTarget(s, weight) for s in symbols_to_trade], True)
        self.Debug(f"Set Holdings at {self.time}: {symbols_to_trade}")
        self.plot("Universe", "Symbols Traded", len(symbols_to_trade))

    def close_positions(self):
        """
        Liquidates all open positions before market close.
        """
        self.liquidate()
        self.Debug(f"Liquidated all positions at {self.time}")


# Supporting class for holding SMA and ATR indicators
class SelectionData:

    def __init__(self, algorithm: QCAlgorithm, symbol: Symbol, sma_period=100, atr_period=14):
        self.algorithm = algorithm
        self.symbol = symbol

        # Create indicators
        self._sma = SimpleMovingAverage(sma_period)
        self._atr = AverageTrueRange(atr_period, MovingAverageType.Simple)

        # Warm up indicators with historical data
        algorithm.warm_up_indicator(symbol, self._sma, Resolution.DAILY)
        algorithm.warm_up_indicator(symbol, self._atr, Resolution.DAILY)

        # Boolean flag for SMA filter
        self.is_above_sma = False

    def update_sma(self, time, value):
        """
        Updates the SMA and evaluates if the price is above SMA.
        """
        #self.algorithm.Debug(f"[{self.symbol}] SMA update at {time} with price {value}")
        if self._sma.update(time, value):
            self.is_above_sma = value > self._sma.current.value
            #self.algorithm.Debug(f"[{self.symbol}] SMA: {self._sma.current.value}, Above SMA: {self.is_above_sma}")