from AlgorithmImports import *

class SP500PEFilteredStrategy(QCAlgorithm):
    """
    S&P 500 P/E Filtered Market-Cap Weighted Strategy
    
    Strategy:
    - Start with S&P 500 universe
    - Exclude 50 stocks with lowest P/E ratios (value traps)
    - Market-cap weight all remaining stocks (~450 stocks)
    - Rebalance weekly
    """
    
    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(days=7)
        self.exclude_lowest_pe = 100  # Exclude 100 lowest P/E stocks
        
        # 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 Market-Cap Weighted Strategy initialized")
        self.debug(f"Excluding: {self.exclude_lowest_pe} lowest P/E stocks")
    
    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 and market caps for all stocks
        pe_ratios = {}
        market_caps = {}
        
        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
                    
                    # 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
                        
            except Exception:
                continue
        
        self.debug(f"Found {len(pe_ratios)} stocks with valid P/E ratios and market caps")
        
        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 for market-cap weighting
        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 for market-cap weighting")
        
        if len(included_stocks) < 50:
            self.debug("Too few stocks remaining after filters")
            return
        
        # Calculate market-cap weights for included stocks
        total_market_cap = sum(market_caps[symbol] for symbol in included_stocks)
        
        target_weights = {}
        for symbol in included_stocks:
            # Calculate proportional market cap weight
            market_cap_weight = market_caps[symbol] / total_market_cap
            target_weights[symbol] = market_cap_weight
        
        # Execute rebalancing
        successful_trades = self.rebalance_to_targets(target_weights)
        
        self.debug(f"Market-cap weighted {len(target_weights)} positions")
        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 some examples of excluded stocks
        self.debug(f"Sample low P/E excluded: {', '.join([self.get_ticker(s) for s, _ in sorted_by_pe[:5]])}")
        
        # 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:
            self.debug(f"  {self.get_ticker(symbol)}: {weight:.2%}")
        
        self.last_rebalance = self.time
    
    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