""" Matplotlib annotation helpers for technical analysis charts. This module provides utilities to draw trend lines, support/resistance levels, and other technical analysis annotations on matplotlib/mplfinance charts. """ from datetime import datetime from typing import Any, Dict, List, Optional, Tuple import matplotlib.dates as mdates import matplotlib.pyplot as plt import pandas as pd from matplotlib.axes import Axes from matplotlib.lines import Line2D from matplotlib.patches import Rectangle class ChartAnnotations: """Helper class for adding technical analysis annotations to charts.""" DEFAULT_TRENDLINE_COLOR = "blue" DEFAULT_TRENDLINE_STYLE = "--" DEFAULT_TRENDLINE_WIDTH = 1.5 DEFAULT_SUPPORT_COLOR = "green" DEFAULT_RESISTANCE_COLOR = "red" DEFAULT_LEVEL_STYLE = "-" DEFAULT_LEVEL_WIDTH = 2.0 DEFAULT_LEVEL_ALPHA = 0.7 DEFAULT_ZONE_ALPHA = 0.2 @classmethod def draw_trend_line( cls, ax: Axes, start_date: datetime, start_price: float, end_date: datetime, end_price: float, color: str = DEFAULT_TRENDLINE_COLOR, linestyle: str = DEFAULT_TRENDLINE_STYLE, linewidth: float = DEFAULT_TRENDLINE_WIDTH, label: Optional[str] = None, extend: bool = False, ) -> Line2D: """ Draw a trend line on the chart. Args: ax: Matplotlib axes to draw on start_date: Start datetime start_price: Start price end_date: End datetime end_price: End price color: Line color linestyle: Line style (-, --, :, -.) linewidth: Line width label: Optional label for legend extend: Whether to extend line to edge of chart Returns: Line2D object """ dates = [mdates.date2num(start_date), mdates.date2num(end_date)] prices = [start_price, end_price] if extend: xlim = ax.get_xlim() slope = (end_price - start_price) / (dates[1] - dates[0]) # Extend to left edge extended_start_price = start_price + slope * (xlim[0] - dates[0]) dates.insert(0, xlim[0]) prices.insert(0, extended_start_price) # Extend to right edge extended_end_price = end_price + slope * (xlim[1] - dates[1]) dates.append(xlim[1]) prices.append(extended_end_price) line = ax.plot( dates, prices, color=color, linestyle=linestyle, linewidth=linewidth, label=label, )[0] return line @classmethod def draw_support_level( cls, ax: Axes, price: float, color: str = DEFAULT_SUPPORT_COLOR, linestyle: str = DEFAULT_LEVEL_STYLE, linewidth: float = DEFAULT_LEVEL_WIDTH, alpha: float = DEFAULT_LEVEL_ALPHA, label: Optional[str] = None, ) -> Line2D: """ Draw a horizontal support level across the chart. Args: ax: Matplotlib axes price: Support price level color: Line color linestyle: Line style linewidth: Line width alpha: Transparency (0-1) label: Optional label Returns: Line2D object """ label = label or f"Support: ${price:.2f}" line = ax.axhline( y=price, color=color, linestyle=linestyle, linewidth=linewidth, alpha=alpha, label=label, ) return line @classmethod def draw_resistance_level( cls, ax: Axes, price: float, color: str = DEFAULT_RESISTANCE_COLOR, linestyle: str = DEFAULT_LEVEL_STYLE, linewidth: float = DEFAULT_LEVEL_WIDTH, alpha: float = DEFAULT_LEVEL_ALPHA, label: Optional[str] = None, ) -> Line2D: """ Draw a horizontal resistance level across the chart. Args: ax: Matplotlib axes price: Resistance price level color: Line color linestyle: Line style linewidth: Line width alpha: Transparency label: Optional label Returns: Line2D object """ label = label or f"Resistance: ${price:.2f}" line = ax.axhline( y=price, color=color, linestyle=linestyle, linewidth=linewidth, alpha=alpha, label=label, ) return line @classmethod def draw_support_resistance_zone( cls, ax: Axes, lower_price: float, upper_price: float, zone_type: str = "support", alpha: float = DEFAULT_ZONE_ALPHA, ) -> Rectangle: """ Draw a shaded zone for support or resistance. Args: ax: Matplotlib axes lower_price: Lower bound of zone upper_price: Upper bound of zone zone_type: "support" or "resistance" alpha: Transparency Returns: Rectangle patch """ xlim = ax.get_xlim() color = ( cls.DEFAULT_SUPPORT_COLOR if zone_type == "support" else cls.DEFAULT_RESISTANCE_COLOR ) rect = Rectangle( (xlim[0], lower_price), xlim[1] - xlim[0], upper_price - lower_price, facecolor=color, alpha=alpha, edgecolor=None, label=f"{zone_type.capitalize()} Zone: ${lower_price:.2f}-${upper_price:.2f}", ) ax.add_patch(rect) return rect @classmethod def draw_price_channel( cls, ax: Axes, upper_line: Tuple[datetime, float, datetime, float], lower_line: Tuple[datetime, float, datetime, float], color: str = "purple", linestyle: str = DEFAULT_TRENDLINE_STYLE, linewidth: float = DEFAULT_TRENDLINE_WIDTH, fill: bool = True, fill_alpha: float = 0.1, ) -> Tuple[Line2D, Line2D]: """ Draw a price channel with upper and lower bounds. Args: ax: Matplotlib axes upper_line: (start_date, start_price, end_date, end_price) for upper bound lower_line: (start_date, start_price, end_date, end_price) for lower bound color: Line color linestyle: Line style linewidth: Line width fill: Whether to fill area between lines fill_alpha: Fill transparency Returns: Tuple of (upper_line, lower_line) Line2D objects """ upper = cls.draw_trend_line( ax, upper_line[0], upper_line[1], upper_line[2], upper_line[3], color=color, linestyle=linestyle, linewidth=linewidth, label="Upper Channel", ) lower = cls.draw_trend_line( ax, lower_line[0], lower_line[1], lower_line[2], lower_line[3], color=color, linestyle=linestyle, linewidth=linewidth, label="Lower Channel", ) if fill: upper_dates = [ mdates.date2num(upper_line[0]), mdates.date2num(upper_line[2]), ] upper_prices = [upper_line[1], upper_line[3]] lower_dates = [ mdates.date2num(lower_line[0]), mdates.date2num(lower_line[2]), ] lower_prices = [lower_line[1], lower_line[3]] ax.fill_between( upper_dates, upper_prices, lower_prices, color=color, alpha=fill_alpha, ) return upper, lower @classmethod def annotate_signal( cls, ax: Axes, date: datetime, price: float, signal_type: str, text: Optional[str] = None, arrow_color: Optional[str] = None, ) -> None: """ Annotate a buy/sell signal on the chart. Args: ax: Matplotlib axes date: Signal datetime price: Price at signal signal_type: "buy" or "sell" text: Custom annotation text arrow_color: Custom arrow color """ is_buy = signal_type.lower() == "buy" if text is None: text = "BUY" if is_buy else "SELL" if arrow_color is None: arrow_color = "green" if is_buy else "red" # Position text above/below based on signal type xytext_offset = (0, 20) if is_buy else (0, -20) va = "bottom" if is_buy else "top" ax.annotate( text, xy=(mdates.date2num(date), price), xytext=xytext_offset, textcoords="offset points", ha="center", va=va, fontsize=10, fontweight="bold", color=arrow_color, bbox=dict( boxstyle="round,pad=0.3", facecolor="white", edgecolor=arrow_color, alpha=0.8, ), arrowprops=dict( arrowstyle="->", color=arrow_color, lw=2, ), ) @classmethod def add_legend( cls, ax: Axes, loc: str = "best", fontsize: int = 10, ) -> None: """ Add legend to chart with custom styling. Args: ax: Matplotlib axes loc: Legend location fontsize: Font size """ ax.legend( loc=loc, fontsize=fontsize, framealpha=0.9, shadow=True, ) @classmethod def find_support_resistance_levels( cls, df: pd.DataFrame, window: int = 20, num_levels: int = 3, ) -> Dict[str, List[float]]: """ Automatically identify support and resistance levels from OHLC data. Uses a simple algorithm based on local minima (support) and maxima (resistance). Args: df: OHLC DataFrame window: Rolling window for local extrema detection num_levels: Number of top levels to return for each type Returns: Dict with "support" and "resistance" lists of price levels """ supports = [] resistances = [] # Find local minima (support) for i in range(window, len(df) - window): if df["low"].iloc[i] == df["low"].iloc[i - window : i + window].min(): supports.append(df["low"].iloc[i]) # Find local maxima (resistance) for i in range(window, len(df) - window): if df["high"].iloc[i] == df["high"].iloc[i - window : i + window].max(): resistances.append(df["high"].iloc[i]) # Cluster similar levels (within 1% of each other) supports = cls._cluster_levels(supports, tolerance=0.01) resistances = cls._cluster_levels(resistances, tolerance=0.01) # Return top N most significant levels supports = sorted(supports, reverse=True)[:num_levels] resistances = sorted(resistances)[:num_levels] return { "support": supports, "resistance": resistances, } @staticmethod def _cluster_levels(levels: List[float], tolerance: float = 0.01) -> List[float]: """ Cluster price levels that are within tolerance % of each other. Args: levels: List of price levels tolerance: Percentage tolerance (0.01 = 1%) Returns: List of clustered levels (averages) """ if not levels: return [] levels = sorted(levels) clustered = [] current_cluster = [levels[0]] for level in levels[1:]: if abs(level - current_cluster[-1]) / current_cluster[-1] <= tolerance: current_cluster.append(level) else: clustered.append(sum(current_cluster) / len(current_cluster)) current_cluster = [level] # Add last cluster clustered.append(sum(current_cluster) / len(current_cluster)) return clustered