from AlgorithmImports import *
import numpy as np
from datetime import timedelta

class SharpeWeightedStrategy(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2010, 1, 1)  # Set start date
        self.SetCash(100000)  # Set strategy cash
        
        self.UniverseSettings.Resolution = Resolution.Daily
        
        # Define other settings
        self.minimum_market_cap = 20e9  # Minimum market cap
        self.number_of_stocks = 50  # Number of stocks to select
        self.risk_free_rate = 0.04/252  # Daily risk-free rate
        self.lookback = 252  # Trading days lookback
        self.next_rebalance = None
        self.symbols = []

        # Add universe
        self.AddUniverse(self.CoarseSelectionFunction)
        
    def CoarseSelectionFunction(self, coarse):
        if self.Time.day != 1:  # Only run on first day of month
            return Universe.Unchanged
            
        # Filter for stocks meeting our criteria
        selected = [x for x in coarse if x.HasFundamentalData and 
                   x.Market == "usa" and 
                   x.Price > 0 and
                   x.MarketCap > self.minimum_market_cap]
                   
        # Sort by dollar volume (for liquidity)
        selected = sorted(selected, key=lambda x: x.DollarVolume, reverse=True)
        
        return [x.Symbol for x in selected[:200]]  # Take top 200 by volume for detailed analysis
        
    def OnData(self, data):
        if not self.symbols:  # Skip if we haven't selected symbols yet
            return
            
        if self.next_rebalance is None:
            self.next_rebalance = self.Time + timedelta(days=1)
            return
            
        # Check if it's time to rebalance
        if self.Time >= self.next_rebalance:
            # Calculate next rebalance date (first day of next month)
            self.next_rebalance = self.Time.replace(day=1) + timedelta(days=32)
            self.next_rebalance = self.next_rebalance.replace(day=1)
            
            self.Rebalance()
    
    def Rebalance(self):
        sharpe_ratios = {}
        
        # Calculate Sharpe ratio for each symbol
        for symbol in self.symbols:
            history = self.History(symbol, self.lookback, Resolution.Daily)
            if len(history) < self.lookback:  # Skip if not enough data
                continue
                
            # Calculate daily returns
            returns = history['close'].pct_change().dropna()
            
            # Calculate Sharpe ratio
            excess_returns = returns - self.risk_free_rate
            expected_return = excess_returns.mean()
            volatility = returns.std()
            
            if volatility == 0:  # Skip if volatility is zero
                continue
                
            sharpe = expected_return / volatility * np.sqrt(252)  # Annualized Sharpe
            sharpe_ratios[symbol] = sharpe
        
        if not sharpe_ratios:  # If no valid Sharpe ratios, skip rebalance
            return
            
        # Select top stocks by Sharpe ratio
        sorted_symbols = sorted(sharpe_ratios.items(), key=lambda x: x[1], reverse=True)
        top_symbols = sorted_symbols[:self.number_of_stocks]
        
        # Calculate weights based on Sharpe ratios
        total_sharpe = sum(sharpe for _, sharpe in top_symbols)
        weights = {symbol: max(0, sharpe/total_sharpe) for symbol, sharpe in top_symbols}
        
        # Normalize weights to sum to 1
        weight_sum = sum(weights.values())
        weights = {symbol: weight/weight_sum for symbol, weight in weights.items()}
        
        # Liquidate positions we don't want anymore
        for holding in self.Portfolio.Values:
            if holding.Symbol not in weights and holding.Invested:
                self.Liquidate(holding.Symbol)
        
        # Rebalance to target weights
        for symbol, weight in weights.items():
            self.SetHoldings(symbol, weight)
            
    def OnSecuritiesChanged(self, changes):
        self.symbols = [x.Symbol for x in changes.AddedSecurities]
        # Clean up any removed symbols
        for removed in changes.RemovedSecurities:
            if removed.Symbol in self.symbols:
                self.symbols.remove(removed.Symbol)