from AlgorithmImports import *

class MarketSentimentAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 1, 1)
        self.SetEndDate(2024, 12, 31)
        self.SetCash(100000)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)
        
        # Sentiment variables
        self.ad_line = 0
        self.total_sentiment_score = 0
        
        # Dictionary to hold SMA indicators
        self.sma_indicators = {}
        
        # Log control
        self.max_logs_per_day = 50
        self.current_log_count = 0
        self.last_log_date = self.Time.date()
        
        # Initialize plots
        self.Plot("Sentiment", "Market Sentiment", 0)
        self.Plot("Indicators", "A/D Line", 0)
        self.Plot("Indicators", "A/D Ratio", 0)
        self.Plot("Indicators", "TRIN", 0)
        self.Plot("Indicators", "% Above MA", 0)
        self.Plot("Indicators", "Volume Breadth", 0)

        """
        Plot Descriptions:
        
        1. Sentiment - Market Sentiment:
        - **Description**: This plot represents the composite sentiment score derived from multiple market indicators, including the Advance/Decline Ratio (A/D Ratio), Percentage of Stocks Above 50-Day Moving Average (% Above MA), and the TRIN (Arms Index).
        
        - **Calculation**:
            - **A/D Ratio**: Calculated as the ratio of advancing stocks to declining stocks. A higher ratio indicates more stocks are advancing compared to declining.
            - **% Above MA**: Represents the percentage of stocks trading above their 50-day Simple Moving Average (SMA). A higher percentage suggests a majority of stocks are in an upward trend.
            - **TRIN**: The TRIN Index is calculated as the ratio of the A/D Ratio to the Volume Breadth (difference between advancing and declining volumes). It provides insights into market pressure and investor sentiment.
            - **Composite Sentiment Score**: Combines the normalized values of the above indicators using weighted factors (e.g., 0.4 for A/D Ratio, 0.4 for % Above MA, and 0.2 for TRIN) to produce a single sentiment score ranging between -1 and 1.
        
        - **Purpose**:
            - **Overall Market Sentiment**: Provides a unified measure of the market's bullish or bearish sentiment based on breadth indicators.
            - **Trend Identification**: Helps in identifying prevailing market trends and potential reversals by analyzing the combined signals from multiple indicators.
            - **Decision-Making**: Assists in making informed trading decisions by offering a consolidated view of market health and momentum.
        
        - **Interpretation**:
            - **Positive Sentiment Score**: Indicates a bullish market sentiment where the majority of stocks are advancing, and more volume is supporting upward movements.
            - **Negative Sentiment Score**: Reflects a bearish market sentiment with more declining stocks and higher volume on the downside.
            - **Neutral Range**: Scores around zero suggest a balanced market with no clear bullish or bearish dominance.
        
        - **Usage in Strategy**:
            - **Entry and Exit Signals**: Utilize the sentiment score to determine optimal points for entering or exiting positions based on prevailing market conditions.
            - **Risk Management**: Adjust position sizing or implement hedging strategies in response to significant shifts in market sentiment.
            - **Confirmation Tool**: Use the sentiment score in conjunction with other technical or fundamental indicators to validate trading signals.
        
        - **Visualization Enhancements**:
            - **Threshold Lines**: Incorporate horizontal lines at key sentiment score levels (e.g., -0.5, 0, +0.5) to quickly identify overbought or oversold conditions.
            - **Color Coding**: Apply color gradients to the sentiment score plot to visually distinguish between bullish (e.g., green), bearish (e.g., red), and neutral (e.g., gray) sentiments.
            - **Annotations**: Add markers or annotations for significant sentiment score changes to highlight potential trading opportunities or risks.
        
        - **Benefits**:
            - **Holistic View**: Combines multiple indicators into a single metric, simplifying the analysis of market conditions.
            - **Dynamic Adaptation**: Adjusts to changing market environments by reflecting real-time sentiment shifts.
            - **Enhanced Clarity**: Facilitates easier interpretation and quicker decision-making through a consolidated sentiment measure.

        2. Indicators - A/D Line:
           - **Description**: The Advance/Decline Line, a cumulative indicator that sums the difference between advancing and declining stocks.
           - **Purpose**: Measures the breadth of the market, showing the net number of advancing versus declining stocks over time.
           - **Usage**: An upward trend indicates increasing market strength, whereas a downward trend suggests weakening market conditions.
        
        3. Indicators - A/D Ratio:
           - **Description**: The ratio of advancing stocks to declining stocks.
           - **Purpose**: Provides a quick snapshot of market sentiment by comparing the number of advancing stocks against declining ones.
           - **Usage**: Ratios above 1 indicate more advancing stocks, signaling bullish sentiment. Ratios below 1 suggest bearish sentiment.
        
        4. Indicators - TRIN:
           - **Description**: The TRIN (Arms Index), calculated as the ratio of the A/D Ratio to the Volume Breadth.
           - **Purpose**: Combines market breadth and volume to assess overall market pressure and investor sentiment.
           - **Usage**: Values below 1 typically indicate bullish sentiment, while values above 1 suggest bearish sentiment.
        
        5. Indicators - % Above MA:
           - **Description**: The percentage of stocks trading above their 50-day Simple Moving Average (SMA).
           - **Purpose**: Measures the overall market trend by determining the proportion of stocks in an uptrend.
           - **Usage**: Higher percentages indicate a majority of stocks are in an uptrend, signaling bullish market conditions. Lower percentages suggest a bearish trend.
        
        6. Indicators - Volume Breadth:
           - **Description**: The difference between advancing and declining volumes.
           - **Purpose**: Assesses the strength behind market movements by comparing the volume of advancing stocks against declining ones.
           - **Usage**: Positive values indicate more volume in advancing stocks, reinforcing bullish sentiment. Negative values point to stronger volume in declining stocks, supporting bearish sentiment.
        """        
        
    def CoarseSelectionFunction(self, coarse):
        # Enhanced Filtering:
        # 1. Has Fundamental Data
        # 2. Security Type is Equity
        # 3. Dollar Volume > $1,000,000
        # 4. Symbol does not contain spaces
        filtered = [
            x for x in coarse 
            if x.HasFundamentalData 
            and x.Symbol.SecurityType == SecurityType.Equity
            and x.DollarVolume > 1e6
            and ' ' not in x.Symbol.Value
        ]
        
        # Select top 500 based on dollar volume
        selected = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)[:500]
        self.LogLimited(f"Coarse Selection: {len(selected)} equity symbols selected based on dollar volume.")
        
        # Initialize SMA for new symbols
        for coarse_fundamental in selected:
            symbol = coarse_fundamental.Symbol
            if symbol not in self.sma_indicators:
                try:
                    sma = self.SMA(symbol, 50, Resolution.Daily)
                    self.sma_indicators[symbol] = sma
                    self.RegisterIndicator(symbol, sma, Resolution.Daily)
                    self.LogLimited(f"SMA(50) initialized for {symbol}.")
                except Exception as e:
                    self.LogLimited(f"Error initializing SMA for {symbol}: {e}")
                    continue
        
        return [x.Symbol for x in selected]

    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.sma_indicators:
                try:
                    sma = self.SMA(symbol, 50, Resolution.Daily)
                    self.sma_indicators[symbol] = sma
                    self.RegisterIndicator(symbol, sma, Resolution.Daily)
                    self.LogLimited(f"SMA(50) initialized for newly added {symbol}.")
                except Exception as e:
                    self.LogLimited(f"Error initializing SMA for newly added {symbol}: {e}")
        
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.sma_indicators:
                self.sma_indicators.pop(symbol)
                self.LogLimited(f"SMA(50) removed for {symbol}.")

    def OnData(self, data):
        if not data.Bars:
            return  # Skip if no data available
        
        advancing_stocks = 0
        declining_stocks = 0
        advancing_volume = 0
        declining_volume = 0
        total_stocks = len(data.Bars)
        
        # Loop through each symbol in the current data
        for symbol, bar in data.Bars.items():
            try:
                if bar.Close > bar.Open:
                    advancing_stocks += 1
                    advancing_volume += bar.Volume
                elif bar.Close < bar.Open:
                    declining_stocks += 1
                    declining_volume += bar.Volume
                # If bar.Close == bar.Open, neither advancing nor declining
            except Exception as e:
                self.LogLimited(f"Error processing symbol {symbol}: {e}")
                continue  # Skip to the next symbol
        
        self.LogLimited(f"Advancing Stocks: {advancing_stocks}, Declining Stocks: {declining_stocks}")
        self.LogLimited(f"Advancing Volume: {advancing_volume}, Declining Volume: {declining_volume}")
        
        # Calculate Indicators
        ad_ratio = (advancing_stocks / declining_stocks) if declining_stocks > 0 else float('inf')
        volume_breadth = advancing_volume - declining_volume
        percent_above_ma = self.CalculatePercentAboveMovingAverage(data, total_stocks)
        trin = ((advancing_stocks / declining_stocks) / (advancing_volume / declining_volume)) if declining_volume > 0 else 0
        
        self.LogLimited(f"A/D Ratio: {ad_ratio:.2f}, Volume Breadth: {volume_breadth}")
        self.LogLimited(f"Percent Above MA: {percent_above_ma:.2f}%, TRIN: {trin:.2f}")
        
        # Update A/D Line
        net_advances = advancing_stocks - declining_stocks
        self.ad_line += net_advances
        
        # Sentiment Score
        self.total_sentiment_score = self.CalculateSentimentScore(ad_ratio, percent_above_ma, trin)
        
        # Log the sentiment metrics
        self.LogLimited(f"Market Sentiment: {self.total_sentiment_score:.2f} | A/D Line: {self.ad_line} | A/D Ratio: {ad_ratio:.2f} | TRIN: {trin:.2f}")
        
        # Plotting
        self.Plot("Sentiment", "Market Sentiment", self.total_sentiment_score)
        self.Plot("Indicators", "A/D Line", self.ad_line)
        self.Plot("Indicators", "A/D Ratio", ad_ratio)
        self.Plot("Indicators", "TRIN", trin)
        self.Plot("Indicators", "% Above MA", percent_above_ma)
        self.Plot("Indicators", "Volume Breadth", volume_breadth)

    def CalculatePercentAboveMovingAverage(self, data, total_stocks):
        """Calculate the percentage of stocks above their 50-day moving average."""
        count_above_ma = 0
        for symbol, bar in data.Bars.items():
            if symbol in self.sma_indicators:
                sma = self.sma_indicators[symbol].Current.Value
                if sma is not None and self.Securities[symbol].Price > sma:
                    count_above_ma += 1
        percent_above_ma = (count_above_ma / total_stocks) * 100 if total_stocks > 0 else 0
        self.LogLimited(f"Count Above MA: {count_above_ma}, Total Stocks: {total_stocks}, Percent Above MA: {percent_above_ma:.2f}%")
        return percent_above_ma

    def CalculateSentimentScore(self, ad_ratio, percent_above_ma, trin):
        """Combine multiple indicators into a composite sentiment score."""
        # Normalize and weight indicators (example weights: 0.4, 0.4, 0.2)
        normalized_ad_ratio = min(max(ad_ratio / 2, -1), 1)  # Normalize between -1 and 1
        normalized_percent_above_ma = (percent_above_ma - 50) / 50  # Normalize 50% as neutral
        normalized_trin = min(max(1 - trin, -1), 1)  # Invert TRIN (lower is bullish)
        
        sentiment_score = (0.4 * normalized_ad_ratio +
                           0.4 * normalized_percent_above_ma +
                           0.2 * normalized_trin)
        self.LogLimited(f"Normalized A/D Ratio: {normalized_ad_ratio:.2f}, Normalized % Above MA: {normalized_percent_above_ma:.2f}, Normalized TRIN: {normalized_trin:.2f}")
        self.LogLimited(f"Composite Sentiment Score: {sentiment_score:.2f}")
        return sentiment_score

    def LogLimited(self, message):
        current_date = self.Time.date()
        if current_date != self.last_log_date:
            self.current_log_count = 0
            self.last_log_date = current_date
        if self.current_log_count < self.max_logs_per_day:
            self.Log(message)
            self.current_log_count += 1