etf/visualization/charts/indicators.py

"""
技术指标绘制组件
"""

import pandas as pd
import numpy as np
import talib as ta
import random
from lightweight_charts import Chart


def get_fixed_color(num: int) -> str:
    """根据数字生成固定颜色"""
    random.seed(num)
    r = random.randint(0, 255)
    g = random.randint(0, 255)
    b = random.randint(0, 255)
    color = "#{:02x}{:02x}{:02x}".format(r, g, b)
    random.seed(None)
    return color


def add_ema(
    chart: Chart,
    df: pd.DataFrame,
    period: int = 20,
    color: str = None,
    price_label: bool = False,
):
    """添加EMA指标线"""
    name = f"EMA_{period}"
    df[name] = ta.EMA(df["close"], timeperiod=period)

    line_color = color or get_fixed_color(period)
    line = chart.create_line(
        name, color=line_color, width=2,
        price_label=price_label, price_line=False
    )
    line.set(df[["time", name]])
    return line


def add_cci(
    chart: Chart,
    df: pd.DataFrame,
    period: int = 14,
    height: float = 0.15,
    position: str = "bottom",
):
    """添加CCI副图"""
    cci = ta.CCI(df["high"], df["low"], df["close"], timeperiod=period)
    df[f"CCI_{period}"] = cci

    # 创建副图
    cci_chart = chart.create_subchart(
        position=position, width=1, height=height, sync=True
    )
    cci_chart.layout(font_family="Times New Roman")
    cci_chart.legend(visible=True, font_size=14, color="#FFFFFF")
    cci_chart.time_scale(visible=False)

    # CCI线
    cci_line = cci_chart.create_line(
        name=f"CCI_{period}", color="#FF0000", width=2
    )
    cci_line.set(df[["time", f"CCI_{period}"]])

    # 水平参考线
    for level, label in [(100, "+100"), (-100, "-100")]:
        df[f"cci_{label}"] = level
        ref_line = cci_chart.create_line(
            name=label, color="#D4C21C", width=1,
            style="dashed", price_label=False, price_line=False
        )
        ref_line.set(df[["time", f"cci_{label}"]])

    return cci_chart


def add_macd(
    chart: Chart,
    df: pd.DataFrame,
    fastperiod: int = 12,
    slowperiod: int = 26,
    signalperiod: int = 9,
    height: float = 0.15,
    position: str = "bottom",
):
    """添加MACD副图"""
    macd, signal, hist = ta.MACD(
        df["close"],
        fastperiod=fastperiod,
        slowperiod=slowperiod,
        signalperiod=signalperiod,
    )

    df["DIF"] = macd
    df["DEA"] = signal
    macd_name = f"MACD_{fastperiod}_{slowperiod}_{signalperiod}"
    df[macd_name] = hist * 2

    # 创建副图
    macd_chart = chart.create_subchart(
        position=position, width=1, height=height, sync=True
    )
    macd_chart.layout(font_family="Times New Roman")
    macd_chart.legend(visible=True, font_size=14, color="#FFFFFF")
    macd_chart.time_scale(visible=False)

    # 柱状图
    histogram = macd_chart.create_histogram(name=macd_name)
    hist_data = df[["time", macd_name]].copy()
    hist_data["prev_value"] = hist_data[macd_name].shift(1)

    def get_color(row):
        current, prev = row[macd_name], row["prev_value"]
        is_hollow = (current >= 0 and current < prev) or (current < 0 and current > prev)
        if current >= 0:
            return "rgba(255, 0, 0, 0.5)" if is_hollow else "#ff0000"
        else:
            return "rgba(0, 255, 0, 0.5)" if is_hollow else "#00FF00"

    hist_data["color"] = hist_data.apply(get_color, axis=1)
    hist_data = hist_data.drop("prev_value", axis=1)
    histogram.set(hist_data)

    # DIF线
    dif_line = macd_chart.create_line(
        name="DIF", color="#2962FF", width=2, price_label=False, price_line=False
    )
    dif_line.set(df[["time", "DIF"]])

    # DEA线
    dea_line = macd_chart.create_line(
        name="DEA", color="#FF0000", width=2, price_label=False, price_line=False
    )
    dea_line.set(df[["time", "DEA"]])

    return macd_chart


def add_td_sequence(chart: Chart, df: pd.DataFrame):
    """添加TD序列标记"""
    close = df["close"].to_list()
    td = [0, 0, 0, 0]
    up = 0
    down = 0

    for i in range(4, len(close)):
        if close[i] > close[i - 4]:
            up += 1
            down = 0
            td.append(up)
        else:
            down -= 1
            up = 0
            td.append(down)

    df["TD"] = td

    # 添加标记
    markers = []
    for _, row in df.iterrows():
        td_val = row["TD"]
        if td_val in [9, 13]:
            markers.append({
                "time": row["time"].strftime("%Y-%m-%d %H:%M:%S"),
                "position": "above",
                "shape": "arrow_down",
                "color": "#00FF00",
                "text": str(td_val),
            })
        elif td_val in [-9, -13]:
            markers.append({
                "time": row["time"].strftime("%Y-%m-%d %H:%M:%S"),
                "position": "below",
                "shape": "arrow_up",
                "color": "#FF0000",
                "text": str(abs(td_val)),
            })

    chart.marker_list(markers)


def add_buy_sell_signals(chart: Chart, df: pd.DataFrame):
    """添加买卖信号标记"""
    if "buy" not in df.columns and "sell" not in df.columns:
        return

    markers = []
    for _, row in df.iterrows():
        if row.get("buy") == 1:
            markers.append({
                "time": row["time"].strftime("%Y-%m-%d"),
                "position": "below",
                "shape": "arrow_up",
                "color": "#00FF00",
                "text": "B",
            })
        elif row.get("sell") == 1:
            markers.append({
                "time": row["time"].strftime("%Y-%m-%d"),
                "position": "above",
                "shape": "arrow_down",
                "color": "#FF0000",
                "text": "S",
            })

    chart.marker_list(markers)


class IndicatorOverlay:
    """指标叠加器"""

    def __init__(self, chart: Chart):
        self.chart = chart

    def add_default_indicators(self, df: pd.DataFrame):
        """添加默认指标组合"""
        # 短期EMA
        for period in [3, 5, 8, 10, 12, 15]:
            add_ema(self.chart, df, period=period, color=5)

        # 长期EMA
        for period in [30, 35, 40, 45, 50, 60]:
            add_ema(self.chart, df, period=period, color=10)

        # 年线
        add_ema(self.chart, df, period=260)

        # MACD
        add_macd(self.chart, df, fastperiod=30, slowperiod=90)

        # CCI
        add_cci(self.chart, df, period=120)

        # TD序列
        add_td_sequence(self.chart, df)