# MEGA-CAP MOMENTUM CONCENTRATION STRATEGY
# Designed to catch and concentrate into mega-winners like NVDA/META 2023-2024

from AlgorithmImports import *
import numpy as np
from datetime import timedelta
from collections import deque

class MegaCapConcentrationStrategy(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(1998, 1, 1)  # Start before the big run
        self.SetEndDate(2025, 9, 10)
        self.SetCash(1000000)
        
        # CONCENTRATION PARAMETERS
        self.max_positions = 10  # Start with up to 10
        self.min_positions = 2   # Can concentrate down to 2
        self.max_single_position = 0.40  # Allow 40% in a single winner
        
        # MOMENTUM DETECTION
        self.momentum_fast = 20  # 20 day momentum
        self.momentum_slow = 60  # 60 day momentum  
        self.acceleration_period = 10  # Detect accelerating moves
        
        # MEGA-CAP FOCUS
        self.min_market_cap = 100000000000  # $100B minimum
        self.mega_cap_bias = 2.0  # Double weight for $500B+ companies
        
        # WINNER RULES
        self.winner_threshold = 0.20  # 20% gain = winner
        self.mega_winner_threshold = 0.50  # 50% gain = mega winner
        self.max_winner_allocation = 0.40  # Can go 40% in mega winners
        
        # RISK MANAGEMENT
        self.stop_loss = -0.15  # 15% stop (wider for big names)
        self.trailing_stop_trigger = 0.30  # Trailing stop after 30%
        self.trailing_stop_distance = 0.10  # 10% trailing
        
        # State
        self.positions = {}
        self.momentum_scores = {}
        self.rebalance_days = 10
        self.next_rebalance = self.Time
        
        # DEFINE MEGA-CAPS TO WATCH
        # Tech mega-caps that could explode
        self.mega_caps = {
            'NVDA': self.AddEquity("NVDA", Resolution.Daily).Symbol,
            'META': self.AddEquity("META", Resolution.Daily).Symbol,
            'MSFT': self.AddEquity("MSFT", Resolution.Daily).Symbol,
            'GOOGL': self.AddEquity("GOOGL", Resolution.Daily).Symbol,
            'AAPL': self.AddEquity("AAPL", Resolution.Daily).Symbol,
            'AMZN': self.AddEquity("AMZN", Resolution.Daily).Symbol,
            'TSLA': self.AddEquity("TSLA", Resolution.Daily).Symbol,
        }
        
        # Other large caps to consider
        self.large_caps = {
            'AMD': self.AddEquity("AMD", Resolution.Daily).Symbol,
            'AVGO': self.AddEquity("AVGO", Resolution.Daily).Symbol,
            'ORCL': self.AddEquity("ORCL", Resolution.Daily).Symbol,
            'CRM': self.AddEquity("CRM", Resolution.Daily).Symbol,
            'ADBE': self.AddEquity("ADBE", Resolution.Daily).Symbol,
            'NFLX': self.AddEquity("NFLX", Resolution.Daily).Symbol,
            'QCOM': self.AddEquity("QCOM", Resolution.Daily).Symbol,
            'ARM': None,  # Will be added when available
            'SMCI': None,  # Super Micro Computer
        }
        
        # Try to add symbols that might not exist in earlier dates
        try:
            self.large_caps['ARM'] = self.AddEquity("ARM", Resolution.Daily).Symbol
        except:
            pass
        
        try:
            self.large_caps['SMCI'] = self.AddEquity("SMCI", Resolution.Daily).Symbol
        except:
            pass
        
        # Market indicators
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.qqq = self.AddEquity("QQQ", Resolution.Daily).Symbol
        self.vix = self.AddEquity("VIX", Resolution.Daily).Symbol
        
        # AI/Tech ETFs for sector momentum
        self.sector_etfs = {
            'SMH': self.AddEquity("SMH", Resolution.Daily).Symbol,  # Semiconductors
            'IGV': self.AddEquity("IGV", Resolution.Daily).Symbol,  # Software
            'SOXX': self.AddEquity("SOXX", Resolution.Daily).Symbol, # Semiconductors
        }
        
        # All symbols to track
        self.all_symbols = list(self.mega_caps.values()) + [s for s in self.large_caps.values() if s]
        self.all_symbols.extend([self.spy, self.qqq])
        self.all_symbols.extend(self.sector_etfs.values())
        
        # Universe
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.SelectStocks)
        
        # Scheduling
        self.Schedule.On(
            self.DateRules.EveryDay(self.spy),
            self.TimeRules.AfterMarketOpen(self.spy, 30),
            self.DailyCheck
        )
        
        self.Schedule.On(
            self.DateRules.EveryDay(self.spy),
            self.TimeRules.AfterMarketOpen(self.spy, 45),
            self.ConcentrateWinners
        )
        
        self.SetWarmUp(65, Resolution.Daily)
        
        self.SetBrokerageModel(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        
        self.Debug("Mega-Cap Concentration Strategy Initialized")
        self.Debug(f"Tracking: {[str(s) for s in self.mega_caps.keys()]}")
    
    def SelectStocks(self, coarse):
        """Focus on mega-caps and high momentum large caps"""
        
        if self.Time < self.next_rebalance:
            return Universe.Unchanged
        
        if self.IsWarmingUp:
            return []
        
        # Always include our defined mega/large caps
        return self.all_symbols
    
    def DailyCheck(self):
        """Daily momentum and risk check"""
        
        if self.IsWarmingUp:
            return
        
        # Update momentum scores
        self.UpdateMomentumScores()
        
        # Rebalance check
        if self.Time >= self.next_rebalance:
            self.next_rebalance = self.Time + timedelta(days=self.rebalance_days)
            self.Rebalance()
        
        # Risk management
        self.CheckStops()
    
    def UpdateMomentumScores(self):
        """Calculate momentum with bias toward mega-caps"""
        
        self.momentum_scores = {}
        
        # Check mega-caps first (they get priority)
        for name, symbol in self.mega_caps.items():
            if symbol and symbol in self.Securities:
                score = self.CalculateMomentum(symbol, name, is_mega_cap=True)
                if score:
                    self.momentum_scores[symbol] = score
        
        # Then large caps
        for name, symbol in self.large_caps.items():
            if symbol and symbol in self.Securities:
                score = self.CalculateMomentum(symbol, name, is_mega_cap=False)
                if score:
                    self.momentum_scores[symbol] = score
    
    def CalculateMomentum(self, symbol, name, is_mega_cap):
        """Calculate momentum score with acceleration detection"""
        
        history = self.History(symbol, 65, Resolution.Daily)
        
        if len(history) < 65:
            return None
        
        try:
            closes = history['close'].values
            
            # Basic momentum
            mom_20 = (closes[-1] / closes[-self.momentum_fast]) - 1
            mom_60 = (closes[-1] / closes[-self.momentum_slow]) - 1
            
            # Need positive momentum
            if mom_20 <= 0 or mom_60 <= 0:
                return None
            
            # Acceleration check (is momentum increasing?)
            mom_10_new = (closes[-1] / closes[-10]) - 1
            mom_10_old = (closes[-11] / closes[-21]) - 1
            acceleration = mom_10_new - mom_10_old
            
            # Volume check
            volumes = history['volume'].values
            avg_volume = np.mean(volumes[-20:])
            recent_volume = np.mean(volumes[-5:])
            volume_surge = recent_volume / avg_volume if avg_volume > 0 else 1
            
            # Volatility (prefer lower vol for position sizing)
            returns = np.diff(closes[-20:]) / closes[-20:-1]
            volatility = np.std(returns) * np.sqrt(252)
            
            # Score calculation
            base_score = mom_20 * 2 + mom_60  # Weight recent more
            
            # Boost for acceleration
            if acceleration > 0:
                base_score *= (1 + acceleration)
            
            # Boost for volume surge
            if volume_surge > 1.5:
                base_score *= 1.2
            
            # MEGA-CAP BIAS - This is KEY
            if is_mega_cap:
                base_score *= self.mega_cap_bias
                
                # Extra boost for specific AI plays during 2023-2024
                if self.Time.year >= 2023:
                    if name in ['NVDA', 'META', 'AMD', 'SMCI']:
                        # These were the AI winners
                        if acceleration > 0.05:  # If accelerating strongly
                            base_score *= 1.5
                            self.Debug(f"{name} showing AI momentum acceleration: {acceleration:.2%}")
            
            return {
                'symbol': symbol,
                'name': name,
                'score': base_score,
                'momentum_20': mom_20,
                'momentum_60': mom_60,
                'acceleration': acceleration,
                'volatility': volatility,
                'volume_surge': volume_surge,
                'is_mega_cap': is_mega_cap
            }
            
        except:
            return None
    
    def Rebalance(self):
        """Rebalance with heavy concentration in top performers"""
        
        if not self.momentum_scores:
            return
        
        # Sort by score
        sorted_scores = sorted(self.momentum_scores.items(), 
                             key=lambda x: x[1]['score'], 
                             reverse=True)
        
        # Take top N
        top_picks = sorted_scores[:self.max_positions]
        
        # Enter positions
        for symbol, score_data in top_picks:
            if symbol not in self.positions:
                self.EnterPosition(symbol, score_data)
        
        # Remove positions not in top picks (unless they're big winners)
        top_symbols = [s for s, _ in top_picks]
        for symbol in list(self.positions.keys()):
            if symbol not in top_symbols:
                position = self.positions[symbol]
                # Keep if it's a big winner
                if position.get('return', 0) < self.winner_threshold:
                    self.ExitPosition(symbol, "No longer in top momentum")
    
    def ConcentrateWinners(self):
        """Aggressively concentrate into mega winners like NVDA/META"""
        
        if self.IsWarmingUp:
            return
        
        # Update returns
        for symbol in list(self.positions.keys()):
            if symbol in self.Securities and symbol in self.Portfolio:
                if self.Portfolio[symbol].Invested:
                    position = self.positions[symbol]
                    current_price = self.Securities[symbol].Price
                    entry_price = position['entry_price']
                    
                    if current_price > 0 and entry_price > 0:
                        position['return'] = (current_price / entry_price) - 1
        
        # Categorize positions
        mega_winners = []
        winners = []
        normal = []
        
        for symbol, position in self.positions.items():
            ret = position.get('return', 0)
            name = position.get('name', '')
            
            # SPECIAL HANDLING for AI boom stocks in 2023-2024
            if self.Time.year >= 2023 and name in ['NVDA', 'META', 'AMD', 'SMCI']:
                if ret > 0.10:  # Lower threshold for known AI winners
                    mega_winners.append((symbol, ret, position))
                    self.Debug(f"AI Winner {name}: {ret:.1%} return")
            elif ret >= self.mega_winner_threshold:
                mega_winners.append((symbol, ret, position))
            elif ret >= self.winner_threshold:
                winners.append((symbol, ret, position))
            else:
                normal.append((symbol, ret, position))
        
        # Calculate new position sizes
        new_sizes = {}
        
        # Mega winners get huge allocation
        for symbol, ret, position in mega_winners:
            # Size based on return magnitude
            if ret > 1.0:  # 100%+ return
                size = self.max_single_position
            elif ret > 0.5:  # 50%+ return
                size = 0.30
            else:
                size = 0.20
            
            new_sizes[symbol] = size
            
        # Regular winners
        for symbol, ret, position in winners:
            new_sizes[symbol] = 0.10
        
        # Normal positions get minimum
        for symbol, ret, position in normal:
            if ret > -0.05:  # Not down too much
                new_sizes[symbol] = 0.05
            else:
                # Cut losers
                self.ExitPosition(symbol, f"Cutting loser at {ret:.1%}")
        
        # Normalize allocations
        total = sum(new_sizes.values())
        if total > 1.2:  # Allow some leverage
            scale = 1.2 / total
            for symbol in new_sizes:
                new_sizes[symbol] *= scale
        
        # Execute concentration
        for symbol, target_size in new_sizes.items():
            if symbol in self.Portfolio and symbol in self.Securities:
                current_size = self.Portfolio[symbol].HoldingsValue / self.Portfolio.TotalPortfolioValue
                
                if abs(target_size - current_size) > 0.02:
                    try:
                        self.SetHoldings(symbol, target_size)
                        
                        if target_size > 0.15:  # Big position
                            position = self.positions[symbol]
                            self.Debug(f"CONCENTRATED: {position.get('name', symbol)} to {target_size:.1%} (Return: {position.get('return', 0):.1%})")
                    except:
                        pass
    
    def EnterPosition(self, symbol, score_data):
        """Enter a new position"""
        
        if symbol not in self.Securities:
            return
        
        # Initial position sizing
        if score_data['is_mega_cap']:
            size = 0.10  # Start with 10% for mega-caps
        else:
            size = 0.05  # 5% for others
        
        try:
            self.SetHoldings(symbol, size)
            
            self.positions[symbol] = {
                'entry_price': self.Securities[symbol].Price,
                'entry_time': self.Time,
                'size': size,
                'name': score_data['name'],
                'return': 0,
                'peak_price': self.Securities[symbol].Price
            }
            
            self.Debug(f"Entered {score_data['name']} - Score: {score_data['score']:.2f}")
            
        except:
            pass
    
    def ExitPosition(self, symbol, reason=""):
        """Exit a position"""
        
        if symbol in self.Portfolio and self.Portfolio[symbol].Invested:
            self.Liquidate(symbol)
            
            if reason and symbol in self.positions:
                name = self.positions[symbol].get('name', symbol)
                self.Debug(f"Exited {name}: {reason}")
        
        if symbol in self.positions:
            del self.positions[symbol]
    
    def CheckStops(self):
        """Check stop losses and trailing stops"""
        
        for symbol in list(self.positions.keys()):
            if symbol not in self.Portfolio or not self.Portfolio[symbol].Invested:
                continue
            
            if symbol not in self.Securities:
                continue
            
            position = self.positions[symbol]
            current_price = self.Securities[symbol].Price
            
            if current_price <= 0:
                continue
            
            # Update peak
            if current_price > position.get('peak_price', 0):
                position['peak_price'] = current_price
            
            ret = position.get('return', 0)
            
            # Stop loss
            if ret < self.stop_loss:
                self.ExitPosition(symbol, f"Stop loss at {ret:.1%}")
                continue
            
            # Trailing stop for big winners
            if ret > self.trailing_stop_trigger:
                peak = position['peak_price']
                drawdown = (current_price / peak) - 1 if peak > 0 else 0
                
                if drawdown < -self.trailing_stop_distance:
                    self.ExitPosition(symbol, f"Trailing stop at {ret:.1%}")
    
    def OnData(self, data):
        """Process data"""
        pass
    
    def OnSecuritiesChanged(self, changes):
        """Handle security changes"""
        for security in changes.AddedSecurities:
            security.FeeModel = ConstantFeeModel(1)
    
    def OnEndOfAlgorithm(self):
        """Final report"""
        
        final = self.Portfolio.TotalPortfolioValue
        ret = (final / 1000000 - 1) * 100
        
        self.Debug("=" * 60)
        self.Debug("MEGA-CAP CONCENTRATION - FINAL REPORT")
        self.Debug("=" * 60)
        self.Debug(f"Final Value: ${final:,.0f}")
        self.Debug(f"Total Return: {ret:.1f}%")
        
        # Show final positions
        self.Debug("\nFinal Positions:")
        for symbol, position in self.positions.items():
            if symbol in self.Portfolio and self.Portfolio[symbol].Invested:
                name = position.get('name', symbol)
                size = self.Portfolio[symbol].HoldingsValue / self.Portfolio.TotalPortfolioValue
                position_return = position.get('return', 0)
                self.Debug(f"  {name}: {size:.1%} position, {position_return:.1%} return")
        
        self.Debug("=" * 60)