# Simple Trend Following with Concentration
# Fixes the 2021-2025 performance issues

from AlgorithmImports import *

class TrendMomentumStrategy(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2004, 11, 18)  
        self.SetEndDate(2025, 8, 31)    
        self.SetCash(1000000)
        
        # SIMPLIFIED PARAMETERS
        self.momentum_period = 90        # 3-month momentum (faster)
        self.max_positions = 10          # Maximum 10 stocks
        self.min_positions = 3           # Minimum 3 stocks (concentration)
        
        # SINGLE CLEAR TREND FILTER
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.qqq = self.AddEquity("QQQ", Resolution.Daily).Symbol
        self.trend_period = 150          # 150-day trend (between 100-200 MA)
        
        # Safe havens
        self.ief = self.AddEquity("IEF", Resolution.Daily).Symbol
        self.gld = self.AddEquity("GLD", Resolution.Daily).Symbol
        
        # Position tracking
        self.selected_stocks = []
        self.month = -1
        
        # Monthly rebalancing
        self.Schedule.On(
            self.DateRules.MonthStart("SPY"),
            self.TimeRules.AfterMarketOpen("SPY", 30),
            self.Rebalance
        )
        
        # Universe settings - BIGGER universe for more opportunities
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)
        
        self.SetBenchmark("SPY")
        self.SetWarmUp(200)
        
    def CoarseSelectionFunction(self, coarse):
        """Get liquid stocks only - no fundamental requirements"""
        if self.Time.month == self.month:
            return Universe.Unchanged
        self.month = self.Time.month
        
        if self.IsWarmingUp:
            return []
        
        # Very simple filter - just liquid stocks
        filtered = [x for x in coarse 
                   if x.Price > 10           # Above $10
                   and x.Price < 2000         # Below $2000 (avoid BRK.A)
                   and x.DollarVolume > 20000000  # $20M daily volume
                   and x.HasFundamentalData]
        
        # Sort by dollar volume and take top 500
        sorted_by_volume = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)
        
        return [x.Symbol for x in sorted_by_volume[:500]]
    
    def Rebalance(self):
        """Simple monthly rebalancing"""
        if self.IsWarmingUp:
            return
        
        # Step 1: Check if we should be in stocks at all
        market_trend = self.CheckMarketTrend()
        
        if market_trend == "BEAR":
            self.Debug(f"{self.Time.date()}: BEAR MARKET - Going defensive")
            self.Liquidate()
            self.SetHoldings(self.gld, 0.5)  # 50% gold
            self.SetHoldings(self.ief, 0.3)  # 30% bonds
            # 20% cash
            self.selected_stocks = []
            return
        
        # Step 2: Get momentum scores for all stocks
        momentum_scores = {}
        
        for kvp in self.ActiveSecurities:
            symbol = kvp.Key
            if symbol in [self.spy, self.qqq, self.ief, self.gld]:
                continue
                
            # Calculate simple momentum
            try:
                history = self.History(symbol, self.momentum_period + 20, Resolution.Daily)
                if history.empty or len(history) < self.momentum_period:
                    continue
                    
                closes = history['close']
                
                # Simple 3-month return
                momentum = (closes.iloc[-1] / closes.iloc[-self.momentum_period]) - 1
                
                # Only consider positive momentum
                if momentum > 0:
                    # Check if it's not in a downtrend
                    recent_high = closes.tail(20).max()
                    current = closes.iloc[-1]
                    
                    # Skip if more than 10% below recent high (failing stocks)
                    if current > recent_high * 0.90:
                        momentum_scores[symbol] = momentum
                        
            except:
                continue
        
        if len(momentum_scores) < 3:
            self.Debug(f"{self.Time.date()}: Too few momentum stocks")
            return
        
        # Step 3: Determine position count based on market strength
        if market_trend == "STRONG":
            # Concentrate in strong markets
            num_positions = min(self.min_positions + 2, len(momentum_scores))
            self.Debug(f"{self.Time.date()}: STRONG TREND - {num_positions} concentrated positions")
        else:
            # More diversified in normal markets
            num_positions = min(self.max_positions, len(momentum_scores))
            self.Debug(f"{self.Time.date()}: NORMAL TREND - {num_positions} positions")
        
        # Step 4: Select top momentum stocks
        sorted_momentum = sorted(momentum_scores.items(), key=lambda x: x[1], reverse=True)
        new_portfolio = [symbol for symbol, _ in sorted_momentum[:num_positions]]
        
        # Step 5: Execute trades
        self.ExecuteTrades(new_portfolio, market_trend)
    
    def CheckMarketTrend(self):
        """Determine market trend - BEAR, NORMAL, or STRONG"""
        try:
            spy_hist = self.History(self.spy, self.trend_period + 50, Resolution.Daily)
            qqq_hist = self.History(self.qqq, self.trend_period + 50, Resolution.Daily)
            
            if spy_hist.empty or qqq_hist.empty:
                return "NORMAL"
            
            spy_closes = spy_hist['close']
            qqq_closes = qqq_hist['close']
            
            # Current prices
            spy_current = spy_closes.iloc[-1]
            qqq_current = qqq_closes.iloc[-1]
            
            # Moving averages
            spy_ma = spy_closes.tail(self.trend_period).mean()
            qqq_ma = qqq_closes.tail(self.trend_period).mean()
            
            # Recent performance
            spy_1m = (spy_closes.iloc[-1] / spy_closes.iloc[-21]) - 1 if len(spy_closes) > 21 else 0
            spy_3m = (spy_closes.iloc[-1] / spy_closes.iloc[-63]) - 1 if len(spy_closes) > 63 else 0
            
            # BEAR: Both indices below MA or recent crash
            if (spy_current < spy_ma and qqq_current < qqq_ma) or spy_1m < -0.08:
                return "BEAR"
            
            # STRONG: Both above MA and strong recent performance
            if spy_current > spy_ma * 1.05 and qqq_current > qqq_ma * 1.05 and spy_3m > 0.05:
                return "STRONG"
            
            # NORMAL: Everything else
            return "NORMAL"
            
        except:
            return "NORMAL"
    
    def ExecuteTrades(self, new_portfolio, market_trend):
        """Execute trades with position sizing based on trend"""
        # Get current holdings
        current_holdings = [s for s in self.selected_stocks 
                           if self.Portfolio[s].Invested]
        
        # Identify winners to keep
        winners = []
        for symbol in current_holdings:
            if self.Portfolio[symbol].UnrealizedProfitPercent > 0.30:  # 30% gainers
                winners.append(symbol)
                if symbol not in new_portfolio:
                    new_portfolio.append(symbol)  # Keep winners
                    if len(new_portfolio) > self.max_positions:
                        new_portfolio = new_portfolio[:self.max_positions]
        
        # Sells
        sells = [s for s in current_holdings if s not in new_portfolio]
        for symbol in sells:
            self.Liquidate(symbol)
            self.Debug(f"  Selling {symbol.Value}")
        
        # Position sizing
        position_sizes = {}
        
        if market_trend == "STRONG" and len(new_portfolio) <= 5:
            # Concentrated sizing in strong trends
            for i, symbol in enumerate(new_portfolio):
                if i == 0:
                    position_sizes[symbol] = 0.25  # Top position 25%
                elif i < 3:
                    position_sizes[symbol] = 0.20  # Next 2: 20% each
                else:
                    position_sizes[symbol] = 0.15  # Others: 15%
        else:
            # Equal weight normally
            weight = 0.95 / len(new_portfolio) if new_portfolio else 0
            for symbol in new_portfolio:
                if symbol in winners:
                    position_sizes[symbol] = weight * 1.2  # Winners get 20% more
                else:
                    position_sizes[symbol] = weight
        
        # Normalize weights
        total = sum(position_sizes.values())
        if total > 0.95:
            for symbol in position_sizes:
                position_sizes[symbol] = position_sizes[symbol] * 0.95 / total
        
        # Set positions
        for symbol in position_sizes:
            self.SetHoldings(symbol, position_sizes[symbol])
            if symbol not in current_holdings:
                self.Debug(f"  Buying {symbol.Value} at {position_sizes[symbol]:.1%}")
        
        # Small cash/bond allocation
        remaining = 1.0 - sum(position_sizes.values())
        if remaining > 0.02:
            self.SetHoldings(self.ief, remaining)
        
        self.selected_stocks = list(position_sizes.keys())
        
        self.Debug(f"{self.Time.date()}: Holding {len(self.selected_stocks)} positions")
    
    def OnData(self, data):
        """Daily check for crashes only"""
        # Emergency exit on flash crashes
        if not self.IsWarmingUp and len(self.selected_stocks) > 0:
            spy_price = self.Securities[self.spy].Price
            spy_history = self.History(self.spy, 10, Resolution.Daily)
            
            if not spy_history.empty and len(spy_history) >= 5:
                five_day_return = (spy_price / spy_history['close'].iloc[-5]) - 1
                
                # Flash crash: Down 7% in 5 days
                if five_day_return < -0.07:
                    self.Debug(f"{self.Time.date()}: FLASH CRASH DETECTED - Emergency exit")
                    self.Liquidate()
                    self.SetHoldings(self.gld, 0.5)
                    self.SetHoldings(self.ief, 0.5)
                    self.selected_stocks = []