# 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(2024, 12, 31)
        self.SetCash(1000000)
        
        # CONCENTRATION PARAMETERS
        self.max_positions = 5  # Fewer positions for concentration
        self.min_positions = 2   # Can concentrate down to 2
        self.max_single_position = 0.50  # Allow 50% in a single winner
        self.max_leverage = 1.5  # Maximum 1.5x leverage in strong trends
        
        # MOMENTUM DETECTION
        self.momentum_fast = 10  # Faster: 10 day momentum
        self.momentum_slow = 30  # Faster: 30 day momentum  
        self.acceleration_period = 10  # Detect accelerating moves
        self.rebalance_days = 5   # More frequent rebalancing
        
        # 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
                
                # ERA-SPECIFIC BOOSTS for known winners
                year = self.Time.year
                
                # 2010-2014: Mobile/Cloud era
                if 2010 <= year <= 2014:
                    if name in ['AAPL', 'AMZN', 'NFLX']:
                        if acceleration > 0.03:
                            base_score *= 1.5
                            self.Debug(f"{name} showing mobile/cloud momentum: {acceleration:.2%}")
                
                # 2015-2019: FAANG dominance
                elif 2015 <= year <= 2019:
                    if name in ['META', 'AMZN', 'NFLX', 'GOOGL']:
                        if acceleration > 0.03:
                            base_score *= 1.5
                            self.Debug(f"{name} showing FAANG momentum: {acceleration:.2%}")
                
                # 2020-2022: COVID tech boom + TSLA
                elif 2020 <= year <= 2022:
                    if name in ['TSLA', 'NVDA', 'AMD', 'SNOW', 'NET']:
                        if acceleration > 0.03:
                            base_score *= 1.5
                            self.Debug(f"{name} showing COVID/EV momentum: {acceleration:.2%}")
                
                # 2023-2024: AI revolution
                elif year >= 2023:
                    if name in ['NVDA', 'META', 'AMD', 'SMCI']:
                        if acceleration > 0.05:
                            base_score *= 1.5
                            self.Debug(f"{name} showing AI momentum: {acceleration:.2%}")
                
                # 2024-2025: PLTR and AI implementation phase
                if year >= 2024:
                    if name == 'PLTR':
                        # PLTR special detection
                        if mom_20 > 0.15 and acceleration > 0.03:
                            base_score *= 2.0  # Double boost for PLTR
                            self.Debug(f"PLTR DETECTED: Momentum {mom_20:.1%}, Accel {acceleration:.1%}")
                    elif name in ['PLTR', 'SNOW', 'NET', 'DDOG', 'CRWD']:
                        # AI implementation/data plays
                        if acceleration > 0.04:
                            base_score *= 1.7
                            self.Debug(f"{name} showing AI implementation momentum")
            
            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
        
        # TORQUE MODE: More aggressive concentration
        leverage_multiplier = 1.0
        
        # Check market momentum for leverage
        spy_momentum = self.GetMarketMomentum()
        if spy_momentum > 0.10:  # Market up 10%+
            leverage_multiplier = 1.5
            self.Debug(f"TORQUE MODE: Market momentum {spy_momentum:.1%}, using leverage")
        
        # Categorize positions
        mega_winners = []
        winners = []
        normal = []
        
        for symbol, position in self.positions.items():
            ret = position.get('return', 0)
            name = position.get('name', '')
            
            # Lower thresholds for faster concentration
            if ret >= 0.30:  # Was 0.50
                mega_winners.append((symbol, ret, position))
                self.Debug(f"MEGA WINNER {name}: {ret:.1%}")
            elif ret >= 0.15:  # Was 0.20
                winners.append((symbol, ret, position))
            else:
                normal.append((symbol, ret, position))
        
        # EXTREME CONCENTRATION for torque
        new_sizes = {}
        
        # If we have mega winners, go ALL IN
        if mega_winners:
            remaining = 1.0 * leverage_multiplier
            
            # Give mega winners everything
            for i, (symbol, ret, position) in enumerate(mega_winners[:2]):  # Top 2 only
                if i == 0:
                    size = min(0.60 * leverage_multiplier, self.max_single_position * 1.5)
                else:
                    size = min(0.40 * leverage_multiplier, self.max_single_position)
                new_sizes[symbol] = size
                remaining -= size
            
            # Small positions for others
            for symbol, ret, position in winners[:3]:
                if remaining > 0:
                    new_sizes[symbol] = min(0.10, remaining)
                    remaining -= new_sizes[symbol]
        
        # Standard allocation if no mega winners
        else:
            for symbol, ret, position in winners:
                new_sizes[symbol] = 0.20 * leverage_multiplier
            
            for symbol, ret, position in normal:
                if ret > -0.03:
                    new_sizes[symbol] = 0.10
                else:
                    self.ExitPosition(symbol, f"Cutting at {ret:.1%}")
        
        # Execute concentration
        self.ExecuteConcentration(new_sizes)
    
    def GetMarketMomentum(self):
        """Get SPY momentum for leverage decisions"""
        try:
            spy_hist = self.History(self.spy, 30, Resolution.Daily)
            if len(spy_hist) >= 30:
                return (spy_hist['close'][-1] / spy_hist['close'][-30]) - 1
        except:
            pass
        return 0
    
    def ExecuteConcentration(self, new_sizes):
        """Execute the concentration with optional leverage"""
        
        # Normalize if over leverage limit
        total = sum(new_sizes.values())
        if total > self.max_leverage:
            scale = self.max_leverage / total
            for symbol in new_sizes:
                new_sizes[symbol] *= scale
        
        # Execute trades
        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.25:  # Big position
                            position = self.positions[symbol]
                            self.Debug(f"TORQUE POSITION: {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
            
            # Check if it's a short position
            is_short = position.get('is_short', False)
            
            if is_short:
                # SHORT POSITION MANAGEMENT
                entry_price = position['entry_price']
                if entry_price > 0:
                    # For shorts, profit when price goes down
                    ret = (entry_price / current_price) - 1
                    
                    # Cover short if losing (price went up)
                    if ret < -self.short_cover:
                        self.Liquidate(symbol)
                        self.Debug(f"Covered short {position.get('name', symbol)} at {ret:.1%} loss")
                        del self.positions[symbol]
                    # Take profit on shorts if down enough
                    elif ret > 0.20:  # 20% profit on short
                        self.Liquidate(symbol)
                        self.Debug(f"Covered short {position.get('name', symbol)} at {ret:.1%} profit")
                        del self.positions[symbol]
            else:
                # LONG POSITION MANAGEMENT (original logic)
                # 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)