from AlgorithmImports import *

class SP500PEFilteredBetaStrategy(QCAlgorithm):
    """
    S&P 500 P/E Filtered Beta-Weighted Strategy
    
    Strategy:
    - Start with S&P 500 universe
    - Exclude 100 stocks with lowest P/E ratios (value traps)
    - Market-cap weight all remaining stocks (~400 stocks)
    - Double the weight of any stock with beta > 1.0
    - Rebalance weekly
    """
    
    def initialize(self):
        self.set_start_date(2024, 1, 1)
        self.set_end_date(2025, 6, 1)
        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(days=7)
        self.exclude_lowest_pe = 100  # Exclude 100 lowest P/E stocks
        self.beta_multiplier_threshold = 1.0  # Double weight if beta > 1.0
        
        # Data tracking
        self.universe_symbols = []
        self.last_rebalance = datetime.min
        self.initial_rebalance_done = False
        
        # Schedule weekly rebalancing
        self.schedule.on(
            self.date_rules.week_start("SPY"),
            self.time_rules.at(10, 0),
            self.rebalance_strategy
        )
        
        self.debug("S&P 500 P/E Filtered Beta-Weighted Strategy initialized")
        self.debug(f"Excluding: {self.exclude_lowest_pe} lowest P/E stocks")
        self.debug(f"Beta > {self.beta_multiplier_threshold}: 2x weight")
    
    def on_securities_changed(self, changes):
        """Handle universe changes"""
        for removed in changes.removed_securities:
            symbol = removed.symbol
            if symbol in self.universe_symbols:
                self.universe_symbols.remove(symbol)
        
        for added in changes.added_securities:
            symbol = added.symbol
            if symbol != self.spy_symbol:
                self.universe_symbols.append(symbol)
        
        self.debug(f"Universe updated: {len(self.universe_symbols)} stocks")
    
    def rebalance_strategy(self):
        """Main rebalancing logic"""
        
        # Allow initial rebalance immediately, then enforce weekly frequency
        if not self.initial_rebalance_done:
            self.initial_rebalance_done = True
        elif (self.time - self.last_rebalance) < self.rebalance_frequency:
            return
            
        self.debug(f"=== Rebalancing at {self.time} ===")
        
        # Collect P/E ratios, market caps, and betas for all stocks
        pe_ratios = {}
        market_caps = {}
        betas = {}
        
        for symbol in self.universe_symbols:
            try:
                fund = self.securities[symbol].fundamentals
                
                if fund.has_fundamental_data:
                    pe_ratio = fund.valuation_ratios.pe_ratio
                    market_cap = fund.market_cap
                    beta = fund.asset_classification.morningstar_sector_code  # Placeholder
                    
                    # Get actual beta from security
                    if symbol in self.securities:
                        security = self.securities[symbol]
                        # Calculate beta using 252-day returns vs SPY
                        beta = self.calculate_beta(symbol)
                    
                    # Only include stocks with valid P/E ratios and market cap
                    if pe_ratio > 0 and pe_ratio < 1000 and market_cap > 0:
                        pe_ratios[symbol] = pe_ratio
                        market_caps[symbol] = market_cap
                        if beta is not None:
                            betas[symbol] = beta
                        
            except Exception:
                continue
        
        self.debug(f"Found {len(pe_ratios)} stocks with valid data")
        
        if len(pe_ratios) < 100:
            self.debug("Insufficient stocks with valid P/E ratios")
            return
        
        # Sort stocks by P/E ratio
        sorted_by_pe = sorted(pe_ratios.items(), key=lambda x: x[1])
        
        # Identify stocks to exclude (only lowest P/E)
        lowest_pe_stocks = [symbol for symbol, pe in sorted_by_pe[:self.exclude_lowest_pe]]
        excluded_stocks = set(lowest_pe_stocks)
        
        # Get remaining stocks
        included_stocks = [symbol for symbol in pe_ratios.keys() if symbol not in excluded_stocks]
        
        self.debug(f"Excluded {len(excluded_stocks)} lowest P/E stocks")
        self.debug(f"Included {len(included_stocks)} stocks")
        
        if len(included_stocks) < 50:
            self.debug("Too few stocks remaining after filters")
            return
        
        # Calculate base market-cap weights
        base_weights = {}
        total_market_cap = sum(market_caps[symbol] for symbol in included_stocks)
        
        for symbol in included_stocks:
            base_weight = market_caps[symbol] / total_market_cap
            base_weights[symbol] = base_weight
        
        # Apply beta multiplier
        adjusted_weights = {}
        high_beta_count = 0
        
        for symbol in included_stocks:
            weight = base_weights[symbol]
            
            # Double weight if beta > 1.0
            if symbol in betas and betas[symbol] > self.beta_multiplier_threshold:
                weight *= 2.0
                high_beta_count += 1
            
            adjusted_weights[symbol] = weight
        
        # Normalize weights to sum to 1.0
        total_adjusted_weight = sum(adjusted_weights.values())
        target_weights = {symbol: weight / total_adjusted_weight 
                         for symbol, weight in adjusted_weights.items()}
        
        self.debug(f"High beta stocks (>1.0) with 2x weight: {high_beta_count}")
        
        # Execute rebalancing
        successful_trades = self.rebalance_to_targets(target_weights)
        
        self.debug(f"Successfully executed {successful_trades} trades")
        
        # Log P/E range that was excluded
        self.debug(f"Lowest P/E excluded: {sorted_by_pe[0][1]:.1f} to {sorted_by_pe[self.exclude_lowest_pe-1][1]:.1f}")
        
        # Log top holdings by weight
        top_holdings = sorted(target_weights.items(), key=lambda x: x[1], reverse=True)[:5]
        self.debug(f"Top 5 holdings by weight:")
        for symbol, weight in top_holdings:
            beta_str = f"(beta: {betas[symbol]:.2f})" if symbol in betas else ""
            self.debug(f"  {self.get_ticker(symbol)}: {weight:.2%} {beta_str}")
        
        self.last_rebalance = self.time
    
    def calculate_beta(self, symbol):
        """Calculate beta vs SPY using historical returns"""
        try:
            # Get 252 trading days of history
            history = self.history([symbol, self.spy_symbol], 252, Resolution.DAILY)
            
            if history.empty or len(history) < 100:
                return None
            
            # Calculate returns
            symbol_returns = history.loc[symbol]['close'].pct_change().dropna()
            spy_returns = history.loc[self.spy_symbol]['close'].pct_change().dropna()
            
            # Align the data
            aligned_data = symbol_returns.align(spy_returns, join='inner')
            symbol_returns_aligned = aligned_data[0]
            spy_returns_aligned = aligned_data[1]
            
            if len(symbol_returns_aligned) < 50:
                return None
            
            # Calculate beta using covariance method
            covariance = symbol_returns_aligned.cov(spy_returns_aligned)
            variance = spy_returns_aligned.var()
            
            if variance > 0:
                beta = covariance / variance
                return beta
            
            return None
            
        except Exception:
            return None
    
    def rebalance_to_targets(self, target_weights):
        """Execute rebalancing to target weights"""
        
        # Liquidate positions not in target
        current_holdings = {kvp.key for kvp in self.portfolio if kvp.value.invested}
        for symbol in current_holdings:
            if symbol not in target_weights:
                self.liquidate(symbol)
        
        # Set target positions
        successful_trades = 0
        failed_trades = 0
        
        for symbol, target_weight in target_weights.items():
            if target_weight > 0.0001:
                try:
                    if symbol in self.securities and self.securities[symbol].price > 0:
                        self.set_holdings(symbol, target_weight)
                        successful_trades += 1
                    else:
                        failed_trades += 1
                except Exception:
                    failed_trades += 1
                    continue
        
        if failed_trades > 0:
            self.debug(f"Failed to place {failed_trades} orders")
        
        return successful_trades
    
    def get_ticker(self, symbol):
        """Extract ticker from symbol"""
        return str(symbol).split(' ')[0]
    
    def on_data(self, data):
        """Handle incoming data"""
        pass