from AlgorithmImports import *
import numpy as np
import pandas as pd
from System import DayOfWeek

class DynamicTop20LiquidStocks(QCAlgorithm):

    def Initialize(self):
        # Set backtest period and cash
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2024, 8, 31)
        self.SetCash(100000)
        
        # Warm-up for 252 days so that historical data is available for return calculations.
        self.SetWarmUp(252)
        
        # Use dynamic coarse universe selection to pick the top 20 stocks by volume
        self.AddUniverse(self.CoarseSelectionFunction)
        
        # Initialize variable to hold current universe symbols
        self.selectedSymbols = []
        
        # Schedule the rebalancing function for every Wednesday at 11:30 AM
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Wednesday),
                 self.TimeRules.At(11, 30),
                 self.RebalancePortfolio)

    def CoarseSelectionFunction(self, coarse):
        # Filter to stocks that have fundamental data.
        filtered = [x for x in coarse if x.HasFundamentalData]
        # Sort by volume in descending order and select the top 20
        sorted_coarse = sorted(filtered, key=lambda x: x.Volume, reverse=True)
        top20 = [x.Symbol for x in sorted_coarse[:20]]
        self.selectedSymbols = top20
        return top20

    def RebalancePortfolio(self):
        # Liquidate all existing holdings to free up cash before rebalancing.
        self.Liquidate()

        # Ensure we have symbols to work with.
        if not self.selectedSymbols:
            self.Debug("No symbols in universe selection; skipping rebalancing.")
            return

        # Retrieve historical daily price data for the selected symbols over the past 253 days.
        history = self.History(self.selectedSymbols, 253, Resolution.Daily)
        if history.empty:
            self.Debug("Historical data is empty; skipping rebalancing.")
            return

        # Prepare a dictionary to store daily returns for each symbol.
        returns_dict = {}

        # Group historical data by symbol and compute daily returns.
        for symbol in self.selectedSymbols:
            # Check if the symbol is in the historical data index.
            if symbol not in history.index.get_level_values(0):
                self.Debug(f"Symbol {symbol} not found in historical data; skipping.")
                continue

            try:
                # Get the data for the symbol and sort by time.
                df = history.loc[symbol].sort_index()
                # Ensure we have enough data.
                if len(df) < 253:
                    self.Debug(f"Not enough data for {symbol}; skipping.")
                    continue
                # Calculate daily returns: (close_today / close_yesterday) - 1.
                df['return'] = df['close'].pct_change()
                returns = df['return'].dropna().values
                # Take only the last 252 returns.
                if len(returns) >= 252:
                    returns_dict[symbol] = returns[-252:]
            except Exception as e:
                self.Debug(f"Error processing {symbol}: {e}")

        # If no symbol has sufficient data, exit the rebalancing routine.
        if not returns_dict:
            self.Debug("Not enough historical data for any symbol; skipping rebalancing.")
            return

        # List of symbols with valid data.
        symbols_list = list(returns_dict.keys())
        n = len(symbols_list)

        # Build a matrix of shape (n, 252) where each row corresponds to a symbol's returns.
        returns_matrix = np.array([returns_dict[s] for s in symbols_list])

        # Compute the expected daily return (average return) for each symbol.
        expected_daily_returns = returns_matrix.mean(axis=1)

        # Compute the covariance matrix.
        cov_matrix = np.cov(returns_matrix)

        # Invert the covariance matrix. If the matrix is singular, log and exit.
        try:
            inv_cov = np.linalg.inv(cov_matrix)
        except np.linalg.LinAlgError:
            self.Debug("Covariance matrix not invertible; skipping rebalancing.")
            return

        # Define the daily risk free rate: 2.5%/252.
        risk_free_rate_daily = 0.025 / 252

        # Calculate expected daily excess return by subtracting the risk free rate.
        expected_excess_returns = expected_daily_returns - risk_free_rate_daily

        # Compute raw weights using matrix multiplication.
        raw_weights = inv_cov.dot(expected_excess_returns)

        # Set any negative weights to zero.
        raw_weights = np.where(raw_weights < 0, 0, raw_weights)

        # Normalize the weights so that they sum to 1.
        weight_sum = np.sum(raw_weights)
        if weight_sum > 0:
            weights = raw_weights / weight_sum
        else:
            weights = raw_weights  # all zero weights

        # Log the computed weights for debugging.
        self.Debug("Rebalancing on {}:".format(self.Time))
        for i, symbol in enumerate(symbols_list):
            self.Debug(f"{symbol.Value}: weight = {weights[i]:.4f}")

        # Set portfolio target holdings based on the calculated weights.
        for i, symbol in enumerate(symbols_list):
            self.SetHoldings(symbol, weights[i])