from AlgorithmImports import *
import numpy as np

class EnhancedMomentumAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.SetCash(100000)
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2024, 12, 31)
        
        # Enhanced parameters
        self.portfolio_size = 15  # More concentrated for better alpha
        self.rebalance_weeks = 1  # Bi-weekly for faster trend capture
        
        # Risk management
        self.max_position_size = 0.10  # 10% max per stock
        self.momentum_threshold = 0.10  # Only buy stocks with >5% momentum
        self.volatility_lookback = 20
        
        # Enhanced universe - focus on growth and momentum sectors
        self.stock_symbols = [
            # Tech leaders
            "AAPL", "MSFT", "GOOGL", "AMZN", "META", "NVDA", "TSLA", "NFLX", "AMD", "ADBE",
            # Growth stocks
            "CRM", "SHOP", "SQ", "PYPL", "ROKU", "ZOOM", "SNOW", "PLTR", "RBLX", "U",
            # Momentum favorites
            "COIN", "ARKK", "QQQ", "TQQQ", "SPYG", "VUG", "IWF", "VGT", "XLK", "FTEC"
        ]
        
        # Add symbols with error handling
        self.symbols = []
        for ticker in self.stock_symbols:
            try:
                symbol = self.AddEquity(ticker, Resolution.Daily).Symbol
                self.symbols.append(symbol)
            except:
                continue
        
        # Advanced data storage
        self.price_history = {}
        self.volume_history = {}
        self.returns_history = {}
        
        # Performance tracking
        self.rebalance_count = 0
        self.winning_trades = 0
        self.total_trades = 0
        
        self.SetWarmUp(30, Resolution.Daily)  # More data for better calculations
        
        # Bi-weekly rebalancing
        self.Schedule.On(self.DateRules.Every([DayOfWeek.Monday, DayOfWeek.Thursday]), 
                        self.TimeRules.AfterMarketOpen("AAPL", 30), 
                        self.Rebalance)
    
    def OnData(self, data):
        if self.IsWarmingUp:
            return
        
        # Store enhanced data
        for symbol in self.symbols:
            if symbol in data and data[symbol] is not None:
                bar = data[symbol]
                
                # Initialize data structures
                if symbol not in self.price_history:
                    self.price_history[symbol] = RollingWindow[float](120)
                    self.volume_history[symbol] = RollingWindow[float](60)
                    self.returns_history[symbol] = RollingWindow[float](60)
                
                # Store price and volume
                if hasattr(bar, 'Close') and bar.Close > 0:
                    current_price = float(bar.Close)
                    self.price_history[symbol].Add(current_price)
                    
                    # Calculate daily returns
                    if self.price_history[symbol].Count > 1:
                        prev_price = self.price_history[symbol][1]
                        daily_return = (current_price / prev_price) - 1
                        self.returns_history[symbol].Add(daily_return)
                
                if hasattr(bar, 'Volume') and bar.Volume > 0:
                    self.volume_history[symbol].Add(float(bar.Volume))
    
    def Rebalance(self):
        if self.IsWarmingUp:
            return
        
        # Calculate enhanced momentum scores
        momentum_data = {}
        
        for symbol in self.symbols:
            if (symbol in self.price_history and 
                self.price_history[symbol].Count >= 60 and
                symbol in self.returns_history and
                self.returns_history[symbol].Count >= 20):
                
                try:
                    score_data = self.CalculateEnhancedMomentum(symbol)
                    if score_data['composite_score'] > self.momentum_threshold:
                        momentum_data[symbol] = score_data
                except Exception as e:
                    continue
        
        if len(momentum_data) < 5:
            self.Log("Insufficient momentum data for rebalancing")
            return
        
        # Rank by composite momentum score
        ranked_stocks = sorted(momentum_data.items(), 
                              key=lambda x: x[1]['composite_score'], 
                              reverse=True)
        
        # Select top performers with risk adjustment
        selected_stocks = []
        for symbol, data in ranked_stocks[:self.portfolio_size]:
            # Additional quality filters
            if (data['volatility'] < 0.5 and  # Not too volatile
                data['volume_trend'] > 0.8):   # Good liquidity
                selected_stocks.append((symbol, data))
        
        # Ensure minimum selection
        if len(selected_stocks) < 8:
            selected_stocks = [(s, d) for s, d in ranked_stocks[:max(8, len(ranked_stocks)//2)]]
        
        # Dynamic position sizing based on momentum strength
        self.ExecuteTrades(selected_stocks)
        
        # Logging
        self.rebalance_count += 1
        avg_momentum = np.mean([data['composite_score'] for _, data in selected_stocks])
        selected_tickers = [str(s).split()[0] for s, _ in selected_stocks]
        
        self.Log(f"Rebalance #{self.rebalance_count}: {len(selected_stocks)} stocks")
        self.Log(f"Avg momentum: {avg_momentum:.2%}")
        self.Log(f"Top picks: {', '.join(selected_tickers[:5])}")
    
    def CalculateEnhancedMomentum(self, symbol):
        """Calculate multi-factor momentum score"""
        prices = [self.price_history[symbol][i] for i in range(60)][::-1]  # Reverse for chronological order
        returns = [self.returns_history[symbol][i] for i in range(20)]
        
        # Multiple momentum timeframes
        mom_1w = (prices[-1] / prices[-5] - 1) if len(prices) >= 5 else 0
        mom_1m = (prices[-1] / prices[-20] - 1) if len(prices) >= 20 else 0
        mom_3m = (prices[-1] / prices[-60] - 1) if len(prices) >= 60 else 0
        
        # Risk-adjusted momentum
        volatility = np.std(returns) if len(returns) > 5 else 0.5
        sharpe_momentum = mom_1m / (volatility + 0.01) if volatility > 0 else 0
        
        # Acceleration (momentum of momentum)
        recent_mom = (prices[-1] / prices[-10] - 1) if len(prices) >= 10 else 0
        older_mom = (prices[-10] / prices[-20] - 1) if len(prices) >= 20 else 0
        acceleration = recent_mom - older_mom
        
        # Volume trend
        volumes = [self.volume_history[symbol][i] for i in range(min(10, self.volume_history[symbol].Count))]
        volume_trend = 1.0
        if len(volumes) >= 10:
            recent_vol = np.mean(volumes[:5])
            older_vol = np.mean(volumes[5:10])
            volume_trend = recent_vol / (older_vol + 1)
        
        # Composite score with weights optimized for 2023
        composite_score = (
            0.40 * mom_1m +        # Primary momentum
            0.25 * sharpe_momentum + # Risk-adjusted
            0.20 * mom_1w +         # Short-term trend
            0.10 * acceleration +   # Momentum acceleration
            0.05 * (mom_3m * 0.5)   # Long-term context (reduced weight)
        )
        
        return {
            'composite_score': composite_score,
            'momentum_1m': mom_1m,
            'volatility': volatility,
            'acceleration': acceleration,
            'volume_trend': volume_trend
        }
    
    def ExecuteTrades(self, selected_stocks):
        """Execute trades with dynamic position sizing"""
        
        # Calculate position sizes based on momentum strength
        total_momentum = sum([data['composite_score'] for _, data in selected_stocks])
        target_positions = {}
        
        for symbol, data in selected_stocks:
            # Base equal weight
            base_weight = 1.0 / len(selected_stocks)
            
            # Momentum tilt (up to 50% adjustment)
            if total_momentum > 0:
                momentum_weight = (data['composite_score'] / total_momentum) * 0.5
            else:
                momentum_weight = 0
            
            # Final weight with maximum position limit
            final_weight = min(base_weight + momentum_weight, self.max_position_size)
            target_positions[symbol] = final_weight
        
        # Normalize weights to sum to ~1.0
        total_weight = sum(target_positions.values())
        if total_weight > 0:
            target_positions = {k: v/total_weight for k, v in target_positions.items()}
        
        # Execute liquidations
        for symbol in list(self.Portfolio.Keys):
            if symbol not in target_positions and self.Portfolio[symbol].Invested:
                holding = self.Portfolio[symbol]
                if holding.UnrealizedProfitPercent > 0:
                    self.winning_trades += 1
                self.total_trades += 1
                self.Liquidate(symbol)
        
        # Execute new positions
        for symbol, weight in target_positions.items():
            if weight > 0.01:  # Minimum position size
                self.SetHoldings(symbol, weight)
    
    def OnEndOfAlgorithm(self):
        """Performance summary"""
        final_value = self.Portfolio.TotalPortfolioValue
        total_return = (final_value / 100000 - 1) * 100
        win_rate = (self.winning_trades / self.total_trades * 100) if self.total_trades > 0 else 0
        
        self.Log("=" * 50)
        self.Log("ENHANCED MOMENTUM STRATEGY RESULTS")
        self.Log("=" * 50)
        self.Log(f"Initial Capital: $100,000")
        self.Log(f"Final Portfolio Value: ${final_value:,.0f}")
        self.Log(f"Total Return: {total_return:.1f}%")
        self.Log(f"Win Rate: {win_rate:.1f}%")
        self.Log(f"Total Rebalances: {self.rebalance_count}")
        self.Log(f"Total Trades: {self.total_trades}")
        
        # Performance vs SPY (approximate)
        spy_2023_return = 24.2  # Actual SPY return for 2023
        alpha = total_return - spy_2023_return
        self.Log(f"Alpha vs SPY: {alpha:.1f}%")
        self.Log("=" * 50)