from AlgorithmImports import *

class Nasdaq100MeanReversion(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 1, 1)    # Set Start Date
        self.SetEndDate(2025, 3, 1)      # Set End Date
        self.SetCash(10000)              # Set Strategy Cash
        
        # Universe selection - using Nasdaq 100 ETF to get constituents
        self.nasdaq100 = self.AddEquity("QQQ", Resolution.Daily).Symbol
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.Universe.ETF(self.nasdaq100, self.UniverseSettings, self.EtfConstituentsFilter))
        

        # Strategy parameters
        self.spread_percentage = 0.0007  # 0.07% trading cost
        self.lookback = 14               # ROC lookback period
        self.investment_per_trade = 3000 # Amount to invest per trade
        self.max_positions = 10          # Maximum number of concurrent positions
        self.max_days_held = 20          # Maximum holding period
        
        # Data dictionaries for tracking
        self.assets = {}             # Dictionary to store securities data
        self.trades = {}                 # Dictionary to store active trades
        self.roc_data = {}               # Store ROC calculations
        self.atr_data = {}               # Store ATR calculations
        self.holdings = {}               # Store position holdings
        
        # Schedule the daily update function
        self.schedule.on(self.date_rules.every_day(), 
                         self.time_rules.before_market_close("QQQ", 10), 
                         self.DailyUpdate)
        
        # Portfolio tracking metrics
        self.net_profit = 0
        self.realized_trades = []
        self.Log(f"Strategy initialized with {self.investment_per_trade} investment per trade")

    def EtfConstituentsFilter(self, constituents: List[ETFConstituentUniverse]) -> List[Symbol]:
        # Filter only valid constituents with weight > 0
        selected = [c for c in constituents if c.Weight > 0 and c.Symbol is not None]
        
        if not selected:  # If no valid symbols, return an empty list safely
            self.Debug("No valid ETF constituents found!")
            return []
        
        # Sort by weight and take top 100
        selected = sorted(selected, key=lambda c: c.Weight, reverse=True)[:100]

        # Debugging: Show selected symbols
        # extracted_symbols = [c.Symbol.Value for c in selected]
        # self.Debug(f"Selected Constituents: {extracted_symbols}")

        return [c.Symbol for c in selected]
    
    
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol.Value
            
            if symbol not in self.assets:
                # Initialize data for the security
                self.assets[symbol] = security
                self.roc_data[symbol] = RateOfChangePercent(self.lookback)
                self.atr_data[symbol] = AverageTrueRange(self.lookback)
                self.holdings[symbol] = 0
                
                self.add_security(security.Type, symbol, Resolution.Daily)  # Register the security

                # Add the required data
                self.register_indicator(symbol, self.roc_data[symbol], Resolution.Daily)
                self.register_indicator(symbol, self.atr_data[symbol], Resolution.Daily)
                
                self.Debug(f"Added {symbol} to universe")
        
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            
            if symbol in self.assets:
                if self.Portfolio[symbol].Invested:
                    self.Liquidate(symbol)
                    self.Debug(f"Liquidated {symbol} due to removal from universe")
                
                self.assets.pop(symbol, None)
                self.roc_data.pop(symbol, None)
                self.atr_data.pop(symbol, None)
                self.holdings.pop(symbol, None)
                
                if symbol in self.trades:
                    self.trades.pop(symbol, None)
                
                self.Debug(f"Removed {symbol} from universe")
            
        self.Debug(f"Total securities in assets: {len(self.assets)}")
        self.Log(f"Total securities in assets: {len(self.assets)}")
    
    def DailyUpdate(self):
        # Check for exit conditions first
        self.CheckExitSignals()
        
        # Check for entry conditions if we have capacity
        if len(self.trades) < self.max_positions:
            self.CheckEntrySignals()
    
    def CheckEntrySignals(self):
        # Calculate how many new positions we can open
        available_slots = self.max_positions - len(self.trades)
        available_capital = self.Portfolio.Cash
        max_new_positions = min(int(available_capital / self.investment_per_trade), available_slots)
        
        if max_new_positions <= 0:
            return
        
        # Find potential entry signals
        potential_entries = []
        
        for symbol, security in self.assets.items():
            # Skip if already in a trade
            if symbol in self.trades:
                continue
                
            # Skip if data not ready
            if not self.roc_data[symbol].IsReady or not self.atr_data[symbol].IsReady:
                continue
            
            # Get ROC value
            roc = self.roc_data[symbol].Current.Value
            # if roc < -20:
            #     self.debug(f"Current ROC for {symbol}: {roc}")
            
            prev_roc = self.roc_data[symbol].Previous.Value
            # if prev_roc < -20:
            #     self.debug(f"Prev ROC for {symbol}: {roc}")
            
            
            # Entry condition: ROC < -20% and ROC increasing
            if roc is not None and prev_roc is not None:
                if roc < -20 and roc > prev_roc:
                    # Calculate momentum for ranking
                    momentum = roc - prev_roc
                    potential_entries.append((symbol, momentum))
        
        # Sort by momentum (highest first) and take top entries
        if potential_entries:
            sorted_entries = sorted(potential_entries, key=lambda x: x[1], reverse=True)
            entries_to_take = sorted_entries[:max_new_positions]
            
            for symbol, _ in entries_to_take:
                self.EnterTrade(symbol)

    
    def EnterTrade(self, symbol):
        if symbol in self.trades or symbol not in self.assets:
            self.debug(f"Already in trade, or no symbol for {symbol}")
            return
        
        security = self.assets[symbol]
        current_price = security.Price
        
        # Adjust entry price with spread
        entry_price = current_price * (1 + self.spread_percentage)
        
        # Calculate position size
        position_size = self.investment_per_trade / entry_price
        
        # Calculate stop loss and take profit based on ATR
        atr = self.atr_data[symbol].Current.Value
        stop_loss = entry_price - (2.5 * atr)
        take_profit = entry_price + (1 * atr)
        
        # Store trade info
        self.trades[symbol] = {
            'entry_date': self.Time,
            'entry_price': entry_price,
            'position_size': position_size,
            'stop_loss': stop_loss,
            'take_profit': take_profit,
            'days_held': 1
        }
        
        # Execute the trade
        self.MarketOrder(symbol, position_size)
        self.debug(f"Entered trade for {symbol} at {entry_price}, Stop: {stop_loss}, TP: {take_profit}")
    
    def CheckExitSignals(self):
        symbols_to_exit = []
        
        for symbol, trade in self.trades.items():
            security = self.assets[symbol]
            current_price = security.Price
            
            # Adjust exit price with spread
            exit_price = current_price * (1 - self.spread_percentage)
            
            # Check exit conditions
            exit_signal = None
            
            # 1. Take profit
            if exit_price >= trade['take_profit']:
                exit_signal = "Take Profit"
            
            # 2. Stop loss
            elif exit_price <= trade['stop_loss']:
                exit_signal = "Stop Loss"
            
            # 3. Max days held
            elif trade['days_held'] >= self.max_days_held:
                exit_signal = "Max Days Held"
            
            # If we have an exit signal, add to exit list
            if exit_signal:
                symbols_to_exit.append((symbol, exit_signal))
            else:
                # Increment days held
                trade['days_held'] += 1
        
        # Process exits
        for symbol, exit_signal in symbols_to_exit:
            self.ExitTrade(symbol, exit_signal)
    
    def ExitTrade(self, symbol, exit_signal):
        if symbol not in self.trades:
            return
        
        trade = self.trades[symbol]
        
        # Calculate profit
        security = self.assets[symbol]
        current_price = security.Price * (1 - self.spread_percentage)  # Adjust for spread
        profit_percentage = (current_price / trade['entry_price'] - 1) * 100
        profit_amount = trade['position_size'] * (current_price - trade['entry_price'])
        
        # Close the position
        self.Liquidate(symbol)
        
        # Record trade
        self.realized_trades.append({
            'symbol': symbol,
            'entry_date': trade['entry_date'],
            'exit_date': self.Time,
            'days_held': trade['days_held'],
            'entry_price': trade['entry_price'],
            'exit_price': current_price,
            'exit_signal': exit_signal,
            'profit_percentage': profit_percentage,
            'profit_amount': profit_amount
        })
        
        # Update net profit
        self.net_profit += profit_amount
        
        # Remove from active trades
        self.trades.pop(symbol)
        
        self.Debug(f"Exited {symbol} with {exit_signal}. Profit: {profit_amount:.2f} ({profit_percentage:.2f}%)")
    
    def OnData(self, data):
        # Main strategy logic runs in scheduled events (DailyUpdate)
        pass
    
    def on_end_of_algorithm(self) -> None:
        # Final statistics
        self.Debug(f"Strategy completed with net profit: ${self.net_profit:.2f}")
        self.Debug(f"Total trades: {len(self.realized_trades)}")
        
        if len(self.realized_trades) > 0:
            # Calculate win rate
            winning_trades = [t for t in self.realized_trades if t['profit_amount'] > 0]
            win_rate = len(winning_trades) / len(self.realized_trades) * 100
            
            # Calculate average profit
            avg_profit = sum(t['profit_amount'] for t in self.realized_trades) / len(self.realized_trades)
            
            # Calculate average holding period
            avg_days = sum(t['days_held'] for t in self.realized_trades) / len(self.realized_trades)
            
            self.Debug(f"Win rate: {win_rate:.2f}%")
            self.Debug(f"Average profit per trade: ${avg_profit:.2f}")
            self.Debug(f"Average holding period: {avg_days:.2f} days")