utils/charts/annotations.py

"""
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