from AlgorithmImports import *

class MonthlyLowRsiSelection(QCAlgorithm):
    def initialize(self):
        self.set_start_date(2016, 1, 1)
        self.set_end_date(2025, 9, 7)
        self.set_cash(100000)
        self.universe_settings.resolution = Resolution.DAILY
        self.add_equity("SPY")  # To warm up the data
        self.spy = self.add_equity("SPY", Resolution.DAILY).symbol

        self.data = {}  # Store indicators
        self.selection_flag = False

        # Universe Selection for S&P 500 (simulated)
        self.add_universe(self.coarse_selection_function)
        
        self.last_month = -1

        self.rsi_period = 14
        self.sma_period = 14

        self.selected_symbols = []

        self.set_warm_up(self.rsi_period + 1)
        self.selection_flag = False
        self.changes = None

        self.schedule.on(
            self.date_rules.month_start("SPY"),
            self.time_rules.after_market_open("SPY", 30),
            self.monthly_rebalance
        )
        
        # No commissions
        self.set_brokerage_model(BrokerageName.DEFAULT, AccountType.MARGIN)
        self.settings.free_portfolio_value_percentage = 0  # Use all capital

    def coarse_selection_function(self, coarse):
        selected = [x.symbol for x in coarse if x.has_fundamental_data and x.price > 5 and x.dollar_volume > 1e7]
        return selected

    def universe_selection(self, symbols):
        for symbol in symbols:
            if symbol not in self.data:
                rsi = self.rsi(symbol, self.rsi_period, MovingAverageType.WILDERS, Resolution.DAILY)
                sma = self.sma(symbol, self.sma_period, Resolution.DAILY)
                self.data[symbol] = {'rsi': rsi, 'sma': sma}
        return symbols

    def on_securities_changed(self, changes):
        self.changes = changes
        for security in changes.removed_securities:
            symbol = security.symbol
            if symbol in self.data:
                self.data.pop(symbol)

    def monthly_rebalance(self):
        if self.is_warming_up:
            return

        # Filter to symbols with indicators ready
        filtered = [sym for sym in self.data.keys() 
                    if self.data[sym]['rsi'].is_ready and self.data[sym]['sma'].is_ready]
        
        # Compute RSI sorting lowest to highest
        stocks = []
        for sym in filtered:
            price = self.securities[sym].price
            sma = self.data[sym]['sma'].current.value
            rsi = self.data[sym]['rsi'].current.value
            # Filter for price above SMA
            if price > sma:
                stocks.append((sym, rsi))
        stocks.sort(key=lambda x: x[1])  # Sort by RSI

        # Pick the 10 lowest RSI stocks
        selected = [x[0] for x in stocks[:10]]
        self.selected_symbols = selected

        # Liquidate everything first
        self.liquidate()

        # Evenly allocate to 10 positions
        count = len(self.selected_symbols)
        if count == 0:
            return

        weight = 1.0 / count
        for symbol in self.selected_symbols:
            if self.securities[symbol].is_tradable:
                self.set_holdings(symbol, weight)

    def on_data(self, data):
        # No per-bar trading, all handled in scheduled monthly_rebalance
        pass