StatArbZScoreStrategy

# pragma pylint: disable=missing-docstring, invalid-name, pointless-string-statement
# flake8: noqa: F401
# isort: skip_file
# --- Do not remove these imports ---
import numpy as np
import pandas as pd
from datetime import datetime, timedelta, timezone
from pandas import DataFrame
from typing import Dict, Optional, Union, Tuple

from freqtrade.strategy import (
    IStrategy,
    Trade,
    Order,
    PairLocks,
    informative,  # @informative decorator
    # Hyperopt Parameters
    BooleanParameter,
    CategoricalParameter,
    DecimalParameter,
    IntParameter,
    RealParameter,
    # timeframe helpers
    timeframe_to_minutes,
    timeframe_to_next_date,
    timeframe_to_prev_date,
    # Strategy helper functions
    merge_informative_pair,
    stoploss_from_absolute,
    stoploss_from_open,
    AnnotationType,
)

# --------------------------------
# Add your lib to import here
import talib.abstract as ta
from technical import qtpylib


class StatArbZScoreStrategy(IStrategy):
    """
    基于你 plot.py 中 StatisticalArbitrageStrategy 的 Z-score 统计套利策略，
    迁移到 freqtrade 策略格式。

    核心逻辑：
    - 对收盘价做滚动均值 / 标准差，计算 z-score
    - z-score 下穿 -threshold 开多（价格相对均值偏低）
    - z-score 上穿 +threshold 开空（需要在期货模式并开启 can_short）
    - 当 |z-score| < exit_threshold 时，平掉对应方向的仓位
    """

    INTERFACE_VERSION = 3

    # 和你原脚本一致，用 15m
    timeframe = "1m"

    # 允许做空（如果你只做多，可以改成 False，并注释掉 enter_short / exit_short）
    can_short: bool = True

    # ROI 和止损可以先用比较宽松的参数，后面你再根据回测调整
    minimal_roi = {
        "0": 0.20
    }

    stoploss = -0.10

    trailing_stop = False

    process_only_new_candles = True

    use_exit_signal = True
    exit_profit_only = False
    ignore_roi_if_entry_signal = False

    # 至少需要这么多 K 线才有稳定的 z-score
    startup_candle_count: int = 100

    max_open_trades=10
    # -----------------------------
    # 超参（可以用于 hyperopt）
    # -----------------------------

    # 回看周期 = 你原脚本的 lookback_period
    lookback = IntParameter(20, 288, default=96, space="buy")

    # 进场 z-score 阈值 = 你原脚本的 zscore_threshold
    entry_zscore = RealParameter(1.0, 4.0, default=2.0, space="buy")

    # 退出 z-score 阈值 = 你原脚本的 exit_threshold
    exit_zscore = RealParameter(0.1, 1.0, default=0.5, space="sell")

    # 是否使用 EMA 代替 SMA
    use_ema = CategoricalParameter([False, True], default=False, space="buy")

    # 绘图配置（freqtrade plot 命令用）
    @property
    def plot_config(self):
        return {
            "main_plot": {
                "close": {},
                "mean_price": {"color": "orange"},
                "upper_band": {"color": "grey"},
                "lower_band": {"color": "grey"},
            },
            "subplots": {
                "Z-Score": {
                    "zscore": {"color": "blue"},
                    "entry_pos": {"color": "green"},
                    "entry_neg": {"color": "red"},
                }
            }
        }

    # -----------------------------
    # 指标计算：把 plot.py 的 calculate_zscore 搬进来
    # -----------------------------
    def _calculate_zscore(
        self, series: pd.Series, lookback: int, use_ema: bool
    ) -> Tuple[pd.Series, pd.Series, pd.Series]:
        """
        计算 Z-score、均值、标准差
        对应你 plot.py 里的 calculate_zscore。
        """
        if use_ema:
            # 参考你原来的 half-life 算法
            half_life = lookback / 2
            alpha = 1 - np.exp(np.log(0.5) / half_life)
            mean = series.ewm(alpha=alpha, adjust=False).mean()
            std = series.ewm(alpha=alpha, adjust=False).std()
        else:
            mean = series.rolling(window=lookback, min_periods=1).mean()
            std = series.rolling(window=lookback, min_periods=1).std()

        std = std.replace(0, np.nan).ffill().fillna(0.0001)
        zscore = (series - mean) / std

        return zscore, mean, std

    # -----------------------------
    # Freqtrade 必须实现的三个函数
    # -----------------------------
    def populate_indicators(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        """
        在这里计算所有需要的指标（zscore / 均值 / 上下轨等）
        """
        if dataframe.empty:
            return dataframe

        lb = int(self.lookback.value)
        uz = float(self.entry_zscore.value)
        use_ema = bool(self.use_ema.value)

        # 计算 Z-score
        z, mean, std = self._calculate_zscore(
            dataframe["close"], lookback=lb, use_ema=use_ema
        )

        dataframe["zscore"] = z
        dataframe["mean_price"] = mean
        dataframe["std_price"] = std

        # 上下轨（和你原策略一致：mean ± zscore_threshold * std）
        dataframe["upper_band"] = dataframe["mean_price"] + uz * dataframe["std_price"]
        dataframe["lower_band"] = dataframe["mean_price"] - uz * dataframe["std_price"]

        # 方便画图看看正负区域
        dataframe["entry_pos"] = uz
        dataframe["entry_neg"] = -uz

        # 保证 volume 存在（有些数据源 volume 可能为 NaN）
        if "volume" not in dataframe.columns:
            dataframe["volume"] = 1.0
        else:
            dataframe["volume"] = dataframe["volume"].fillna(0)

        return dataframe

    def populate_entry_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        """
        对应你 plot.py 中 generate_signals 里的开仓逻辑：
        - zscore 从上方穿越 -threshold：开多
        - zscore 从下方穿越 +threshold：开空
        """
        uz = float(self.entry_zscore.value)

        # Long：zscore 从上方 >= -threshold 到 当前 < -threshold
        long_cond = (
            # (dataframe["zscore"].shift(1) >= -uz) &
            (dataframe["zscore"] < -uz) &
            (dataframe["volume"] > 0)
        )

        dataframe.loc[long_cond, "enter_long"] = 1

        if self.can_short:
            # Short：zscore 从下方 <= threshold 到 当前 > threshold
            short_cond = (
                # (dataframe["zscore"].shift(1) <= uz) &
                (dataframe["zscore"] > uz) &
                (dataframe["volume"] > 0)
            )
            dataframe.loc[short_cond, "enter_short"] = 1

        return dataframe

    def populate_exit_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        """
        对应你 plot.py 中“Z-score 回归到 exit_threshold 内时平仓”的逻辑：
        - 当 |zscore| < exit_threshold 时，平多 / 平空。
        在 freqtrade 里，我们用同一条件触发 exit_long / exit_short。
        """
        ez = float(self.exit_zscore.value)

        exit_cond = (
            (dataframe["zscore"].abs() < ez) &
            (dataframe["volume"] > 0)
        )

        dataframe.loc[exit_cond, "exit_long"] = 1

        if self.can_short:
            dataframe.loc[exit_cond, "exit_short"] = 1

        return dataframe