# region imports
from AlgorithmImports import *  # Import necessary QuantConnect algorithm classes and methods.
import numpy as np  # Import numpy for mathematical operations.
# endregion

class SelectionData(object):
    """
    This class stores data related to each symbol in the universe.
    It tracks the symbol, calculates its SMA, and stores volume and price-related metrics.
    """

    def __init__(self, symbol, period):
        # Initialize with the symbol and the period for the SMA calculation.
        self.symbol = symbol  # Ticker symbol (e.g., 'AAPL').
        self.sma = SimpleMovingAverage(period)  # Create an SMA indicator with the given period.
        self.is_above_sma = False  # Boolean indicating if the current price is above the SMA.
        self.volume = 0  # Stores the current volume of the symbol.
        self.price_to_sma_ratio = 0  # Stores the ratio of price to SMA for ranking purposes.

    def update(self, time, price, volume):
        """
        Update the SMA and related metrics with new price and volume data.
        """
        self.volume = volume  # Update the volume with the latest data.
        
        # If the SMA value updates with the new price, recalculate metrics.
        if self.sma.Update(time, price):
            # Check if the current price is above the SMA.
            self.is_above_sma = price > self.sma.Current.Value  
            # Calculate the price-to-SMA ratio, handling division by zero.
            self.price_to_sma_ratio = price / self.sma.Current.Value if self.sma.Current.Value != 0 else 0

class MyAlgorithm(QCAlgorithm):
    """
    Main trading algorithm class that selects and manages a portfolio of stocks based on
    fundamental and technical indicators.
    """

    def __init__(self):
        # Initialize key variables for the algorithm.
        self.num_fine = 30  # Number of fine-selected symbols to keep.
        self.smaDict = {}  # Dictionary to store SelectionData objects by symbol.
        self.lastWeek = None  # Track the last week of trading to rebalance monthly.
        self.currently_held_symbols = set()  # Track symbols currently held in the portfolio.
        self.fineDict = {}  # Store fine selection data by symbol.
        self.initial_market_cap = 1e9  # Initial minimum market cap threshold.
        self.inflation_rate = 0.03  # Annual inflation rate for market cap adjustment.

    def Initialize(self):
        """
        Set up the initial algorithm parameters and universe settings.
        """
        self.SetStartDate(1999, 1, 1)  # Start backtest from January 1, 1999.
        self.SetCash(100000)  # Set initial cash to $100,000.
        
        # Add a universe using coarse and fine selection filters.
        self.AddUniverse(self.SelectCoarse, self.SelectFine)
        
        # Use raw data normalization and daily resolution for analysis.
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.symbols = []  # List of selected symbols after filtering.
        self.SetBenchmark('SPY')  # Use SPY as the benchmark index.
        
        # Warm up the algorithm with 300 days of data for SMA calculations.
        self.SetWarmUp(TimeSpan.FromDays(300))

    def SelectCoarse(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        """
        Coarse selection step: Filter stocks based on whether their current price is above the SMA.
        """
        for c in coarse:
            # If the symbol is not already in the SMA dictionary, add it.
            if c.Symbol not in self.smaDict:
                self.smaDict[c.Symbol] = SelectionData(c.Symbol, 200)  # Create SMA with 200-day period.

            # Update the symbol's SMA and volume data.
            avg = self.smaDict[c.Symbol]
            avg.update(c.EndTime, c.AdjustedPrice, c.DollarVolume)

        # Select symbols where the current price is above the SMA.
        return [c.Symbol for c in coarse if self.smaDict[c.Symbol].is_above_sma]

    def SelectFine(self, fine: List[FineFundamental]) -> List[Symbol]:
        """
        Fine selection step: Filter stocks based on fundamental metrics and select the top symbols.
        """
        # Store fine selection data for quick access during rebalancing.
        self.fineDict = {x.Symbol: x for x in fine}

        # Calculate the number of years since the start date to adjust market cap for inflation.
        years_since_start = (self.Time - self.StartDate).days / 365.25
        adjusted_market_cap = self.initial_market_cap * (1 + self.inflation_rate) ** years_since_start

        # Filter symbols based on fundamental metrics.
        filtered_fine = [
            x for x in fine if
            x.Symbol in self.smaDict  # Ensure the symbol exists in the SMA dictionary.
            and x.MarketCap > adjusted_market_cap  # Market cap must exceed the inflation-adjusted threshold.
            and x.ValuationRatios.PSRatio < 1  # Price-to-sales ratio must be less than 1.
            and x.ValuationRatios.PriceChange1M > 0.01  # Price must have increased over the last month.
            and x.OperationRatios.RevenueGrowth.ThreeMonths > 0.01  # Revenue growth over three months must be positive.
            and x.OperationRatios.OperationIncomeGrowth.ThreeMonths > 0.01  # Operating income growth must be positive.
            and x.OperationRatios.NetIncomeGrowth.ThreeMonths > 0.01  # Net income growth must be positive.
            and x.OperationRatios.NetIncomeContOpsGrowth.ThreeMonths > 0.01  # Growth in income from operations must be positive.
            and x.OperationRatios.CFOGrowth.OneYear > 0.01  # CFO growth over one year must be positive.
            and x.OperationRatios.FCFGrowth.OneYear > 0.01  # FCF growth over one year must be positive.
        ]

        # Sort the filtered symbols by price-to-SMA ratio in descending order and select the top symbols.
        top = sorted(filtered_fine, key=lambda x: self.smaDict[x.Symbol].price_to_sma_ratio, reverse=True)[:self.num_fine]
        
        # Store the selected symbols and update the currently held symbols.
        self.symbols = [x.Symbol for x in top]
        self.currently_held_symbols = set(self.Portfolio.Keys)
        
        return self.symbols

    def OnData(self, data):
        """
        Rebalance the portfolio monthly by selling underperforming symbols and setting new holdings.
        """
        # Only rebalance on Wednesdays (weekday 2).
        if self.Time.weekday() != 2:
            return

        # Get the current week number.
        week_number = int(self.Time.strftime('%V'))

        # Rebalance every 4th week (monthly) and ensure it's not the same week as the last rebalance.
        if week_number % 4 == 0 and week_number != self.lastWeek:
            # Liquidate symbols that are no longer selected.
            symbols_to_liquidate = self.currently_held_symbols - set(self.symbols)
            for symbol in symbols_to_liquidate:
                self.Liquidate(symbol)

            # Calculate the total log market cap for weight distribution.
            total_log_market_cap = sum(np.log(self.fineDict[symbol].MarketCap) for symbol in self.symbols if symbol in self.fineDict)

            # Set portfolio holdings based on the log of market cap as a weight.
            for symbol in self.symbols:
                if symbol in data.Keys:
                    log_market_cap_weight = np.log(self.fineDict[symbol].MarketCap) / total_log_market_cap if total_log_market_cap != 0 else 0
                    self.SetHoldings(symbol, log_market_cap_weight)

            # Update the last week of rebalancing.
            self.lastWeek = week_number