StatisticalArbitrageStrategy

# 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,
    BooleanParameter,
    CategoricalParameter,
    DecimalParameter,
    IntParameter,
    RealParameter,
    timeframe_to_minutes,
    timeframe_to_next_date,
    timeframe_to_prev_date,
    merge_informative_pair,
    stoploss_from_absolute,
    stoploss_from_open,
    AnnotationType,
)



class StatisticalArbitrageStrategy(IStrategy):
    """
    ETH/BTC统计套利策略
    基于Z-score进行均值回归交易
    """
    
    # 策略接口版本
    INTERFACE_VERSION = 3

    # 最优时间框架 - 根据原策略设置为15分钟
    timeframe = "15m"

    # 允许做空
    can_short: bool = True

    # 最小ROI - 可以根据需要进行调整
    minimal_roi = {
        "0": 0.10,  # 如果达到10%的利润就立即退出
        "1440": 0.05,  # 24小时后减少到5%
        "2880": 0.02,  # 48小时后减少到2%
        "4320": 0.01,  # 72小时后减少到1%
    }
    # minimal_roi = {
    #     "0": 0.10,  # 如果达到10%的利润就立即退出
    #     "24": 0.05,  # 24小时后减少到5%
    #     "48": 0.02,  # 48小时后减少到2%
    #     "72": 0.10,  # 72小时后减少到1%
    # }
    # 止损
    stoploss = -0.10

    leverage_rate = 0.0

    # if leverage_rate:
    #     minimal_roi = {
    #         k: v * 10.0 for k, v in minimal_roi.items()
    #     }
   
    #     stoploss *= leverage_rate

    # 追踪止损
    trailing_stop = False
    trailing_only_offset_is_reached = False
    trailing_stop_positive = 0.01
    trailing_stop_positive_offset = 0.02

    # 只处理新蜡烛
    process_only_new_candles = True

    # 使用退出信号
    use_exit_signal = True
    exit_profit_only = False
    ignore_roi_if_entry_signal = False

    # 策略需要的启动蜡烛数（基于回看周期）
    startup_candle_count: int = 200

    # 策略参数（可以在超参数优化中调整）
    lookback_period = IntParameter(50, 200, default=96, space="buy")
    entry_threshold = DecimalParameter(1.5, 3.0, default=2.0, decimals=1, space="buy")
    exit_threshold = DecimalParameter(0.2, 1.0, default=0.5, decimals=1, space="sell")
    use_ema = BooleanParameter(default=False, space="buy")
    
    # 风险控制参数
    position_adjustment_enable = True
    max_entry_position_adjustment = 3

    # 订单类型
    order_types = {
        "entry": "limit",
        "exit": "limit",
        "stoploss": "market",
        "stoploss_on_exchange": False
    }

    # 订单有效时间
    order_time_in_force = {
        "entry": "GTC",
        "exit": "GTC"
    }

    @property
    def plot_config(self):
        """
        配置策略可视化
        """
        return {
            "main_plot": {
                "close": {"color": "blue", "fill_to": "lower_band"},
                "mean": {"color": "orange", "type": "line"},
                "upper_band": {"color": "red", "type": "line"},
                "lower_band": {"color": "green", "type": "line"},
            },
            "subplots": {
                "Z-Score": {
                    "zscore": {"color": "purple"},
                    "entry_threshold": {"color": "red", "type": "line"},
                    "exit_threshold": {"color": "green", "type": "line"},
                    "neg_entry_threshold": {"color": "red", "type": "line"},
                    "neg_exit_threshold": {"color": "green", "type": "line"},
                },
                "Position": {
                    "position": {"color": "blue", "type": "bar"}
                }
            }
        }

    def informative_pairs(self):
        """
        如果需要额外的信息性交易对，可以在这里定义
        """
        return []
    def leverage(self, pair: str, current_time: datetime, current_rate: float,
                 proposed_leverage: float, max_leverage: float, entry_tag: str | None, side: str,
                 **kwargs) -> float:
        """
        Customize leverage for each new trade. This method is only called in futures mode.

        :param pair: Pair that's currently analyzed
        :param current_time: datetime object, containing the current datetime
        :param current_rate: Rate, calculated based on pricing settings in exit_pricing.
        :param proposed_leverage: A leverage proposed by the bot.
        :param max_leverage: Max leverage allowed on this pair
        :param entry_tag: Optional entry_tag (buy_tag) if provided with the buy signal.
        :param side: "long" or "short" - indicating the direction of the proposed trade
        :return: A leverage amount, which is between 1.0 and max_leverage.
        """
        return self.leverage_rate
    def calculate_zscore(self, dataframe, price_col='close'):
        """
        计算Z-score（价格标准化）
        """
        lookback = self.lookback_period.value
        use_ema = self.use_ema.value
        
        if use_ema:
            # 使用EMA（对近期数据给予更高权重）
            half_life = lookback / 2
            alpha = 1 - np.exp(np.log(0.5) / half_life)
            mean = dataframe[price_col].ewm(alpha=alpha).mean()
            std = dataframe[price_col].ewm(alpha=alpha).std()
        else:
            # 使用简单移动平均
            mean = dataframe[price_col].rolling(window=lookback, min_periods=1).mean()
            std = dataframe[price_col].rolling(window=lookback, min_periods=1).std()
        
        # 避免除以零
        std = std.replace(0, np.nan).fillna(method='ffill').fillna(0.0001)
        
        zscore = (dataframe[price_col] - mean) / std
        
        return zscore, mean, std

    def populate_indicators(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        """
        添加指标到DataFrame
        """
        # 计算Z-score和相关指标
        zscore, mean, std = self.calculate_zscore(dataframe)
        
        # 将计算结果添加到dataframe
        dataframe['zscore'] = zscore
        dataframe['mean'] = mean
        dataframe['std'] = std
        
        # 计算布林带（基于Z-score阈值）
        entry_threshold = self.entry_threshold.value
        exit_threshold = self.exit_threshold.value
        
        dataframe['upper_band'] = dataframe['mean'] + entry_threshold * dataframe['std']
        dataframe['lower_band'] = dataframe['mean'] - entry_threshold * dataframe['std']
        dataframe['upper_exit'] = dataframe['mean'] + exit_threshold * dataframe['std']
        dataframe['lower_exit'] = dataframe['mean'] - exit_threshold * dataframe['std']
        
        # 添加信号线
        dataframe['zscore_signal'] = 0
        dataframe['position'] = 0
        
        # 计算信号（用于可视化，实际交易信号在populate_entry/exit_trend中计算）
        for i in range(1, len(dataframe)):
            current_zscore = dataframe['zscore'].iloc[i]
            prev_zscore = dataframe['zscore'].iloc[i-1]
            
            # 简单信号计算（用于可视化）
            if prev_zscore >= -entry_threshold and current_zscore < -entry_threshold:
                dataframe.loc[dataframe.index[i], 'zscore_signal'] = 1
            elif prev_zscore <= entry_threshold and current_zscore > entry_threshold:
                dataframe.loc[dataframe.index[i], 'zscore_signal'] = -1
        
        return dataframe

    def populate_entry_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        """
        生成入场信号
        """
        entry_threshold = self.entry_threshold.value
        
        # 初始化信号列
        dataframe.loc[:, 'enter_long'] = 0
        dataframe.loc[:, 'enter_short'] = 0
        
        # 计算Z-score的差值用于判断交叉
        dataframe['zscore_prev'] = dataframe['zscore'].shift(1)
        
        # 多头入场条件：Z-score从上方穿过负阈值
        long_condition = (
            (dataframe['zscore_prev'] >= -entry_threshold) &
            (dataframe['zscore'] < -entry_threshold) &
            (dataframe['volume'] > 0)
        )
        
        # 空头入场条件：Z-score从下方穿过正阈值
        short_condition = (
            (dataframe['zscore_prev'] <= entry_threshold) &
            (dataframe['zscore'] > entry_threshold) &
            (dataframe['volume'] > 0)
        )
        
        # 应用入场信号
        dataframe.loc[long_condition, 'enter_long'] = 1
        dataframe.loc[short_condition, 'enter_short'] = 1
        
        # 添加入场原因
        dataframe.loc[long_condition, 'enter_tag'] = 'zscore_long'
        dataframe.loc[short_condition, 'enter_tag'] = 'zscore_short'
        
        return dataframe

    def populate_exit_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        """
        生成退出信号
        """
        exit_threshold = self.exit_threshold.value
        
        # 初始化退出列
        dataframe.loc[:, 'exit_long'] = 0
        dataframe.loc[:, 'exit_short'] = 0
        
        # 多头退出条件：Z-score回归到退出阈值内（从负值变为大于-exit_threshold）
        long_exit_condition = (
            (dataframe['zscore'] > -exit_threshold) &
            (dataframe['zscore'].shift(1) <= -exit_threshold) &
            (dataframe['volume'] > 0)
        )
        
        # 空头退出条件：Z-score回归到退出阈值内（从正值变为小于exit_threshold）
        short_exit_condition = (
            (dataframe['zscore'] < exit_threshold) &
            (dataframe['zscore'].shift(1) >= exit_threshold) &
            (dataframe['volume'] > 0)
        )
        
        # 应用退出信号
        dataframe.loc[long_exit_condition, 'exit_long'] = 1
        dataframe.loc[short_exit_condition, 'exit_short'] = 1
        
        return dataframe

    def custom_stake_amount(self, pair: str, current_time: datetime, current_rate: float,
                           proposed_stake: float, min_stake: Optional[float], max_stake: float,
                           leverage: float, entry_tag: Optional[str], side: str,
                           **kwargs) -> float:
        """
        自定义仓位大小
        可以根据Z-score的偏离程度调整仓位大小
        """
        # 获取最新的Z-score
        dataframe, _ = self.dp.get_analyzed_dataframe(pair, self.timeframe)
        if len(dataframe) > 0:
            latest_zscore = dataframe['zscore'].iloc[-1]
            
            # Z-score偏离越大，仓位可以相对更大（但不要超过最大仓位限制）
            zscore_abs = abs(latest_zscore)
            
            if zscore_abs > 2.5:
                # Z-score偏离很大，可以使用较大仓位
                stake_multiplier = min(1.2, 1 + (zscore_abs - 2.5) * 0.1)
            elif zscore_abs > 2.0:
                # Z-score偏离适中
                stake_multiplier = 1.0
            else:
                # Z-score偏离较小，减少仓位
                stake_multiplier = 0.8
            
            adjusted_stake = proposed_stake * stake_multiplier
            
            # 确保在最小和最大仓位限制内
            if min_stake is not None:
                adjusted_stake = max(adjusted_stake, min_stake)
            
            adjusted_stake = min(adjusted_stake, max_stake)
            
            return adjusted_stake
        
        return proposed_stake

    def confirm_trade_entry(self, pair: str, order_type: str, amount: float, rate: float,
                           time_in_force: str, current_time: datetime, entry_tag: Optional[str],
                           side: str, **kwargs) -> bool:
        """
        在交易入场前进行确认
        """
        # 获取最新的Z-score
        dataframe, _ = self.dp.get_analyzed_dataframe(pair, self.timeframe)
        if len(dataframe) > 0:
            latest_zscore = dataframe['zscore'].iloc[-1]
            entry_threshold = self.entry_threshold.value
            
            # 检查当前Z-score是否仍然符合入场条件
            if side == 'long' and latest_zscore >= -entry_threshold:
                # Z-score已经回归，取消买入
                return False
            elif side == 'short' and latest_zscore <= entry_threshold:
                # Z-score已经回归，取消卖出
                return False
        
        return True

    def custom_exit(self, pair: str, trade: Trade, current_time: datetime, current_rate: float,
                   current_profit: float, **kwargs) -> Optional[Union[str, bool]]:
        """
        自定义退出逻辑
        可以基于Z-score或利润目标提前退出
        """
        # 获取最新的Z-score
        dataframe, _ = self.dp.get_analyzed_dataframe(pair, self.timeframe)
        if len(dataframe) > 0:
            latest_zscore = dataframe['zscore'].iloc[-1]
            
            # 如果Z-score已经回归到0附近，可以考虑提前退出
            if abs(latest_zscore) < 0.2 and current_profit > 0:
                return "zscore_converged"
            
            # 如果利润达到2%且Z-score开始反向运动，考虑退出
            if current_profit > 0.02:
                if len(dataframe) > 1:
                    prev_zscore = dataframe['zscore'].iloc[-2]
                    if trade.is_short and latest_zscore > prev_zscore:
                        return "profit_and_zscore_reversal_short"
                    elif not trade.is_short and latest_zscore < prev_zscore:
                        return "profit_and_zscore_reversal_long"
        
        return None

# 可选：如果需要额外的风险管理，可以添加以下方法
def adjust_trade_position(self, trade: Trade, current_time: datetime,
                         current_rate: float, current_profit: float,
                         min_stake: float, max_stake: float,
                         **kwargs) -> Optional[float]:
    """
    调整交易仓位（加仓/减仓）
    """
    # 如果Z-score偏离继续扩大，可以考虑加仓
    dataframe, _ = self.dp.get_analyzed_dataframe(trade.pair, self.timeframe)
    if len(dataframe) > 0:
        latest_zscore = dataframe['zscore'].iloc[-1]
        
        # 确保当前仓位不超过最大调整次数
        if trade.nr_of_successful_entries < self.max_entry_position_adjustment:
            # 如果是多头且Z-score继续向下偏离
            if trade.is_long and latest_zscore < -self.entry_threshold.value * 1.2:
                # 加仓
                return current_rate * trade.amount * 0.5  # 加仓50%
            
            # 如果是空头且Z-score继续向上偏离
            elif trade.is_short and latest_zscore > self.entry_threshold.value * 1.2:
                # 加仓
                return current_rate * trade.amount * 0.5  # 加仓50%
    
    return None