from AlgorithmImports import *

class SP500GrowthStrategy(QCAlgorithm):
    """
    S&P 500 Growth Strategy
    
    Growth Score (Higher = Better):
    - YoY Revenue Growth: 50% weight
    - RSI Level: 30% weight (higher RSI = momentum)
    - Price vs 20-day MA: 20% weight
    
    Portfolio: Top 200 Growth stocks
    Tier Weighting:
    - Ranks 1-50: 50% of portfolio weight
    - Ranks 51-100: 25% of portfolio weight
    - Ranks 101-200: 25% of portfolio weight
    """
    
    def initialize(self):
        self.set_start_date(2023, 1, 1)
        self.set_end_date(2023, 12, 31)
        self.set_cash(1000000)
        self.set_benchmark("SPY")
        
        # S&P 500 universe via SPY ETF constituents
        self.spy_symbol = self.add_equity("SPY", Resolution.DAILY).symbol
        self.add_universe(self.universe.etf(self.spy_symbol, self.universe_settings))
        
        # Strategy parameters
        self.rebalance_frequency = timedelta(hours=1)  # Hourly rebalancing as requested
        
        # Data tracking
        self.universe_symbols = []
        self.technical_indicators = {}
        self.last_rebalance = datetime.min
        
        # Schedule rebalancing every hour during market hours
        self.schedule.on(
            self.date_rules.every_day("SPY"),
            self.time_rules.every(timedelta(hours=1)),
            self.rebalance_strategy
        )
        
        self.debug("S&P 500 Growth strategy initialized")
    
    def on_securities_changed(self, changes):
        """Handle universe changes"""
        # Remove old securities and indicators
        for removed in changes.removed_securities:
            symbol = removed.symbol
            if symbol in self.universe_symbols:
                self.universe_symbols.remove(symbol)
            if symbol in self.technical_indicators:
                del self.technical_indicators[symbol]
        
        # Add new securities and create indicators
        for added in changes.added_securities:
            symbol = added.symbol
            if symbol != self.spy_symbol:  # Don't include SPY itself
                self.universe_symbols.append(symbol)
                
                # Create technical indicators using string ticker (to avoid RSI syntax error)
                ticker = str(symbol).split(' ')[0]
                try:
                    self.technical_indicators[symbol] = {
                        'rsi': self.rsi(ticker, 14),
                        'sma': self.sma(ticker, 20),
                        'ticker': ticker
                    }
                except Exception as e:
                    self.debug(f"Failed to create indicators for {ticker}: {e}")
        
        self.debug(f"Universe updated: {len(self.universe_symbols)} stocks")
    
    def rebalance_strategy(self):
        """Main rebalancing logic"""
        
        if (self.time - self.last_rebalance) < self.rebalance_frequency:
            return
            
        self.debug(f"Rebalancing at {self.time}")
        
        # Calculate scores for all stocks
        growth_scores = {}
        
        for symbol in self.universe_symbols:
            try:
                growth_score = self.calculate_growth_score(symbol)
                
                if growth_score is not None:
                    growth_scores[symbol] = growth_score
                    
            except Exception as e:
                continue  # Skip problematic stocks
        
        if len(growth_scores) < 50:
            self.debug(f"Insufficient data: {len(growth_scores)} growth stocks")
            return
        
        # Rank stocks (higher score = better)
        growth_rankings = self.rank_stocks(growth_scores, ascending=False)
        
        # Select top 200 from growth strategy
        top_200_growth = growth_rankings[:200]
        
        # Calculate position weights with tier structure (100% of portfolio)
        growth_weights = self.calculate_tier_weights(top_200_growth, 1.0)
        
        # Execute rebalancing
        self.rebalance_to_targets(growth_weights)
        
        # Log summary
        self.debug(f"Rebalanced: {len(growth_weights)} growth positions")
        
        # Log top picks for verification
        if len(growth_rankings) >= 10:
            top_growth_tickers = [self.get_ticker(s) for s in growth_rankings[:5]]
            self.debug(f"Top 5 Growth: {top_growth_tickers}")
        
        self.last_rebalance = self.time
    
    def calculate_growth_score(self, symbol):
        """
        Calculate growth score (higher = better)
        50% YoY Revenue Growth + 30% RSI + 20% Price vs 20-day MA
        """
        try:
            fund = self.securities[symbol].fundamentals
            if not fund.has_fundamental_data:
                return None
            
            # Component 1: YoY Revenue Growth (50% weight)
            revenue_growth = fund.operation_ratios.revenue_growth.one_year
            if revenue_growth <= 0:
                growth_component = 0
            else:
                growth_component = min(revenue_growth * 100, 100)  # Cap at 100% for scoring
            
            # Component 2: RSI Level (30% weight) - Higher RSI = momentum
            rsi_component = 0
            if symbol in self.technical_indicators:
                rsi_indicator = self.technical_indicators[symbol]['rsi']
                if rsi_indicator.is_ready:
                    rsi_component = rsi_indicator.current.value  # Already 0-100 scale
            
            # Component 3: Price vs 20-day MA (20% weight)
            ma_component = 0
            if symbol in self.technical_indicators:
                sma_indicator = self.technical_indicators[symbol]['sma']
                if sma_indicator.is_ready:
                    current_price = self.securities[symbol].price
                    ma_value = sma_indicator.current.value
                    if ma_value > 0:
                        price_ma_ratio = ((current_price / ma_value) - 1) * 100
                        ma_component = max(0, price_ma_ratio + 10)  # Shift to make positive
            
            # Calculate weighted growth score
            growth_score = (
                growth_component * 0.50 +
                rsi_component * 0.30 +
                ma_component * 0.20
            )
            
            return growth_score
            
        except Exception:
            return None
    
    def rank_stocks(self, scores_dict, ascending=True):
        """Rank stocks by scores and return ordered list"""
        sorted_items = sorted(scores_dict.items(), key=lambda x: x[1], reverse=not ascending)
        return [symbol for symbol, score in sorted_items]
    
    def calculate_tier_weights(self, ranked_symbols, total_weight):
        """
        Calculate tier-based weights
        Ranks 1-50: 50% of portfolio weight (1% each)
        Ranks 51-100: 25% of portfolio weight (0.5% each)
        Ranks 101-200: 25% of portfolio weight (0.25% each)
        """
        weights = {}
        
        # Tier 1: Ranks 1-50 get 50% of total weight
        tier1_individual_weight = (total_weight * 0.50) / 50  # 1% each
        for i in range(min(50, len(ranked_symbols))):
            weights[ranked_symbols[i]] = tier1_individual_weight
        
        # Tier 2: Ranks 51-100 get 25% of total weight
        tier2_individual_weight = (total_weight * 0.25) / 50  # 0.5% each
        for i in range(50, min(100, len(ranked_symbols))):
            weights[ranked_symbols[i]] = tier2_individual_weight
        
        # Tier 3: Ranks 101-200 get 25% of total weight
        tier3_individual_weight = (total_weight * 0.25) / 100  # 0.25% each
        for i in range(100, min(200, len(ranked_symbols))):
            weights[ranked_symbols[i]] = tier3_individual_weight
        
        return weights
    
    def rebalance_to_targets(self, target_weights):
        """Execute rebalancing to target weights"""
        
        # Get current holdings
        current_holdings = {kvp.key: kvp.value.holdings_value / self.portfolio.total_portfolio_value 
                           for kvp in self.portfolio if kvp.value.invested}
        
        # Liquidate positions not in target
        for symbol in current_holdings:
            if symbol not in target_weights:
                self.liquidate(symbol)
        
        # Set target positions
        successful_trades = 0
        for symbol, target_weight in target_weights.items():
            if target_weight > 0.0001:  # Minimum 0.01% position
                try:
                    self.set_holdings(symbol, target_weight)
                    successful_trades += 1
                except Exception:
                    continue  # Skip if order fails
        
        self.debug(f"Successfully placed {successful_trades} orders out of {len(target_weights)} targets")
    
    def get_ticker(self, symbol):
        """Extract ticker from symbol"""
        return str(symbol).split(' ')[0]
    
    def on_data(self, data):
        """Handle incoming data - indicators auto-update"""
        pass