from AlgorithmImports import *
import numpy as np

class ShortTermCompetitionAlgo(QCAlgorithm):
    def Initialize(self):
        self.set_start_date(2025, 1, 1)
        self.set_cash(100000)
        
        self.spy = self.add_equity("SPY", Resolution.MINUTE)   
        self.set_benchmark("SPY")
        
        self.universe_settings.resolution = Resolution.MINUTE
        self.universe_settings.leverage = 2.0 
        
        # --- Parameters ---
        self.rsi_period = 2
        self.sma_period = 5
        
        # Increased to 25 to ensure we get enough signals to rotate capital
        self.buy_threshold = 25   
        self.max_positions = 10 
        
        self.active_symbols = set()
        self.indicators = {} 
        
        self.add_universe(self.coarse_selection)

        self.schedule.on(self.date_rules.every_day("SPY"), 
                 self.time_rules.before_market_close("SPY", 15),
                 self.run_strategy)

    def coarse_selection(self, coarse):
        # 1. Filter
        filtered = [x for x in coarse if x.price > 10 and x.has_fundamental_data]
        # 2. Sort by Dollar Volume
        sorted_by_volume = sorted(filtered, key=lambda x: x.dollar_volume, reverse=True)
        # 3. Take Top 100
        return [x.symbol for x in sorted_by_volume[:100]]

    def OnSecuritiesChanged(self, changes):
        # Handle Removals
        for security in changes.removed_securities:
            symbol = security.symbol
            
            # --- BUG FIX: Zombie Prevention ---
            # If we currently own this stock, DO NOT remove it from our tracker.
            # We need to keep updating its indicators so we know when to sell it.
            if self.portfolio[symbol].invested:
                continue

            if symbol in self.active_symbols:
                self.active_symbols.remove(symbol)
                if symbol in self.indicators:
                    self.unregister_indicator(self.indicators[symbol]['rsi'])
                    self.unregister_indicator(self.indicators[symbol]['sma'])
                    del self.indicators[symbol]

        # Handle Additions
        for security in changes.added_securities:
            symbol = security.symbol
            if symbol not in self.active_symbols and symbol.value != "SPY":
                self.active_symbols.add(symbol)
                
                # Indicators
                rsi = self.rsi(symbol, self.rsi_period, MovingAverageType.WILDERS, Resolution.HOUR)
                sma = self.sma(symbol, self.sma_period, Resolution.HOUR)
                
                # Warmup Loop
                try:
                    history = self.history(symbol, 40, Resolution.HOUR)
                    if not history.empty:
                        # Fix for MultiIndex issues
                        if 'symbol' in history.index.names:
                            history = history.droplevel('symbol')

                        for time, row in history.iterrows():
                            # Fix for Series vs Scalar issues
                            close_price = row.get('close') if isinstance(row, pd.Series) else row['close']
                            if isinstance(close_price, pd.Series):
                                close_price = close_price.iloc[0]

                            if pd.notna(close_price):
                                rsi.update(time, float(close_price))
                                sma.update(time, float(close_price))
                except Exception as e:
                    self.debug(f"Warmup Failed for {symbol.value}: {str(e)}")
                
                self.indicators[symbol] = {'rsi': rsi, 'sma': sma}

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

        targets = {}
        candidates = []
        
        # 1. Identify Candidates & Manage Positions
        # We iterate a copy (list(...)) so we can modify the set safely if needed
        for symbol in list(self.active_symbols):
            if symbol not in self.indicators: continue
            
            inds = self.indicators[symbol]
            
            if not inds['rsi'].is_ready or not inds['sma'].is_ready: 
                continue
            
            rsi_val = inds['rsi'].current.value
            sma_val = inds['sma'].current.value
            
            if not self.securities[symbol].has_data: continue
            price = self.securities[symbol].price
            
            # --- EXIT LOGIC ---
            if self.portfolio[symbol].invested:
                # 1. Profit Take / Standard Exit
                if price > sma_val:
                    targets[symbol] = 0.0 
                    self.log(f"EXIT PROFIT {symbol.value}: Price {price:.2f} > SMA {sma_val:.2f}")
                
                # 2. Stop Loss (New Feature to free up slots)
                elif self.portfolio[symbol].unrealized_profit_percent < -0.10:
                    targets[symbol] = 0.0
                    self.log(f"EXIT STOP LOSS {symbol.value}: Loss < -10%")
                
                else:
                    # Hold position
                    continue
            
            # --- ENTRY LOGIC ---
            else:
                if rsi_val < self.buy_threshold:
                    candidates.append((symbol, rsi_val))

        # 2. Select Top Candidates
        candidates.sort(key=lambda x: x[1]) 
        
        # Calculate available slots
        # Note: We re-check portfolio keys because 'active_symbols' might include non-invested stuff
        current_invested = [s for s in self.portfolio.keys() if self.portfolio[s].invested]
        
        # Count pending exits to free up slots immediately
        pending_exits = 0
        for sym, weight in targets.items():
            if weight == 0.0 and self.portfolio[sym].invested:
                pending_exits += 1

        slots_available = self.max_positions - len(current_invested) + pending_exits
        
        for i in range(min(len(candidates), slots_available)):
            symbol = candidates[i][0]
            targets[symbol] = 2.0 / self.max_positions # Also controls leverage
            self.log(f"ENTRY {symbol.value}: RSI {candidates[i][1]:.2f} < {self.buy_threshold}")
        
        # 3. Execution
        for symbol, weight in targets.items():
            if weight == 0:
                self.liquidate(symbol)
                
                # Clean up: If we sold it, and it's NOT in the top 100 anymore, 
                # we can finally remove it from memory.
                # (Simple check: if we just liquidated it, we can try to remove it 
                # if the universe doesn't want it. For simplicity, we leave it until next universe refresh).
            else:
                self.set_holdings(symbol, weight)