RationalShort

# -*- coding: utf-8 -*-
"""
RationalShort - 三态理性做空策略（V9 优化版）

适用：freqtrade 2025.10，期货 / 合约（can_short = True）

设计要点：
1. 根据 EMA200 + ADX 划分三种市场状态：
   - 1 = Bull 牛市：价格在 EMA200 上方且 ADX > 阈值
   - 2 = Side 震荡：ADX <= 阈值
   - 3 = Bear 熊市：价格在 EMA200 下方且 ADX > 阈值

2. 入场只做突破布林上轨的“首次穿越”，避免上轨附近连环抄顶。
   - 牛市：极端超买 + 短期爆拉才逆势短线做空
   - 震荡：普通超买做空
   - 熊市：小反弹就做空

3. 出场逻辑：
   - custom_exit 根据市场状态决定止盈 / 止损 / 超时强平
   - custom_stoploss 负责移动止盈，保护浮盈

4. 额外：
   - 使用 enter_tag 标记开仓时的市场状态，便于回测分析。
"""

from datetime import datetime

import numpy as np
import talib.abstract as ta
from pandas import DataFrame

from freqtrade.strategy import IStrategy, IntParameter, DecimalParameter
from freqtrade.persistence import Trade


class RationalShort(IStrategy):
    """理性做空策略 (Rational Short System) - V9 三态模型版"""

    # === 1. 策略基础配置 ===
    INTERFACE_VERSION = 3

    timeframe = "1h"
    can_short = True

    # 由 custom_stoploss / custom_exit 接管，这里留一个兜底
    stoploss = -0.10

    # ROI 完全交给 custom_exit，设置一个极大值
    minimal_roi = {"0": 100}

    # === 2. 参数区 ===
    # ADX 阈值，区分震荡 / 趋势
    adx_threshold = IntParameter(20, 40, default=25, space="buy")

    # 不同状态下的 RSI 入场门槛
    # 牛市：极端超买（80~90）
    rsi_bull = IntParameter(80, 90, default=80, space="buy")
    # 震荡：正常超买
    rsi_side = IntParameter(60, 80, default=65, space="buy")
    # 熊市：中等超买即可（反弹就空）
    rsi_bear = IntParameter(40, 65, default=55, space="buy")

    # 布林带倍数
    bb_std = DecimalParameter(1.8, 2.5, default=2.0, space="buy")

    # 超时时间（小时）
    timeout_hours = IntParameter(4, 24, default=8, space="sell")

    # === 3. 保护机制 ===
    @property
    def protections(self):
        return [
            {"method": "CooldownPeriod", "stop_duration_candles": 1},
            {
                "method": "StoplossGuard",
                "lookback_period_candles": 24,
                "trade_limit": 3,
                "stop_duration_candles": 6,
                "only_per_pair": False,
            },
        ]

    # === 4. 指标计算 ===
    def populate_indicators(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        # A. 布林带
        bollinger = ta.BBANDS(
            dataframe,
            timeperiod=20,
            nbdevup=float(self.bb_std.value),
            nbdevdn=float(self.bb_std.value),
        )
        dataframe["bb_upper"] = bollinger["upperband"]
        dataframe["bb_middle"] = bollinger["middleband"]
        dataframe["bb_lower"] = bollinger["lowerband"]

        # B. RSI & ADX
        dataframe["rsi"] = ta.RSI(dataframe, timeperiod=14)
        dataframe["adx"] = ta.ADX(dataframe)

        # C. EMA200 作为大级别趋势线
        dataframe["ema_200"] = ta.EMA(dataframe, timeperiod=200)

        # D. 短期加速指标（3 根 K 的涨幅）
        dataframe["roc_3"] = ta.ROC(dataframe, timeperiod=3)

        # E. 标记市场状态 (1=Bull, 2=Side, 3=Bear)
        thr = int(self.adx_threshold.value)
        dataframe["market_state"] = 2  # 默认震荡

        dataframe.loc[
            (dataframe["close"] > dataframe["ema_200"]) & (dataframe["adx"] > thr),
            "market_state",
        ] = 1

        dataframe.loc[
            (dataframe["close"] < dataframe["ema_200"]) & (dataframe["adx"] > thr),
            "market_state",
        ] = 3

        # F. 准备好 enter_tag / exit_short 列，方便后续逻辑 / 统计
        if "enter_tag" not in dataframe.columns:
            dataframe["enter_tag"] = ""

        if "exit_short" not in dataframe.columns:
            dataframe["exit_short"] = 0

        return dataframe

    # === 5. 开仓逻辑 (分状态路由) ===
    def populate_entry_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        # 只做“突破上轨的首次穿越”：上一根在下，这一根收盘在上
        cross_upper = (
            (dataframe["close"] > dataframe["bb_upper"]) &
            (dataframe["close"].shift(1) <= dataframe["bb_upper"].shift(1))
        )

        # 基础条件：突破上轨 + 有量
        base_cond = cross_upper & (dataframe["volume"] > 0)

        cond_bull = (
            base_cond &
            (dataframe["market_state"] == 1) &
            (dataframe["rsi"] > int(self.rsi_bull.value)) &
            (dataframe["roc_3"] > 4)
        )

        cond_side = (
            base_cond &
            (dataframe["market_state"] == 2) &
            (dataframe["rsi"] > int(self.rsi_side.value))   # 新增：只在 ema200 下方开空
        )

        cond_bear = (
            base_cond &
            (dataframe["market_state"] == 3) &
            (dataframe["rsi"] > int(self.rsi_bear.value))
        )

        dataframe.loc[cond_bull, ["enter_short", "enter_tag"]] = (1, "bull")
        dataframe.loc[cond_side, ["enter_short", "enter_tag"]] = (1, "side")
        dataframe.loc[cond_bear, ["enter_short", "enter_tag"]] = (1, "bear")

        return dataframe

    # === 6. 平仓信号（用于配合 custom_exit） ===
    def populate_exit_trend(self, dataframe: DataFrame, metadata: dict) -> DataFrame:
        """给出一个基础的 exit_short 信号：跌回布林中轨以下。"""
        if "exit_short" not in dataframe.columns:
            dataframe["exit_short"] = 0

        dataframe.loc[
            (dataframe["close"] < dataframe["bb_middle"]) & (dataframe["volume"] > 0),
            "exit_short",
        ] = 1

        return dataframe

    # === 7. 核心：三态风控引擎 ===
    def custom_exit(
        self,
        pair: str,
        trade: Trade,
        current_time: datetime,
        current_rate: float,
        current_profit: float,
        **kwargs,
    ):
        """根据市场状态决定不同的止盈 / 止损 / 超时逻辑。"""

        # 取当前时间之前最近一根 K 线，防止使用未来数据
        dataframe, _ = self.dp.get_analyzed_dataframe(pair, self.timeframe)
        try:
            df = dataframe.loc[dataframe["date"] <= current_time]
            if df.empty:
                last_candle = dataframe.iloc[-1]
            else:
                last_candle = df.iloc[-1]
        except Exception:
            # 兜底：直接用最后一根
            last_candle = dataframe.iloc[-1]

        state = int(last_candle.get("market_state", 2))  # 默认震荡
        exit_signal = int(last_candle.get("exit_short", 0))

        trade_duration = (current_time - trade.open_date_utc).total_seconds() / 3600
        base_timeout = int(self.timeout_hours.value)

        # 不同状态下最大持仓时间
        if state == 1:      # 牛市：短线
            max_hours = max(3, int(base_timeout * 0.6))
        elif state == 3:    # 熊市：趋势单可多拿一会儿
            max_hours = int(base_timeout * 1.5)
        else:               # 震荡：使用基础超时时间
            max_hours = base_timeout

        # --------------------------------------
        # 1. 牛市：短线 scalping
        # --------------------------------------
        if state == 1:
            tag = getattr(trade, "enter_tag", "")

            # 这单不是在 Bull 状态下开的（Side / Bear 单被拉上来）
            if tag != "bull":
                # 状态反转保护：小亏小赚就走，避免拖成大亏
                if current_profit > -0.003:   # 亏不到 0.3% 就直接平
                    return "regime_flip"
                # 亏得更深的情况交给全局止损 / timeout，不再走 bull_stop_loss 逻辑
            else:
                # 只有真正的 bull 单才执行原来的 scalping + bull_stop_loss 逻辑
                if current_profit <= -0.01:
                    return "bull_stop_loss"

                if current_profit >= 0.012:
                    return "bull_scalp_profit"

                if exit_signal == 1 and current_profit > 0:
                    return "exit_signal"


        # --------------------------------------
        # 2. 震荡：布林带回归
        # --------------------------------------
        elif state == 2:
            # 止损：约 -3%
            if current_profit <= -0.03:
                return "stop_loss"

            # exit_signal 触发且 >0 就离场
            if exit_signal == 1 and current_profit > 0:
                return "exit_signal"

            # 兜底止盈：> 4% 直接平仓
            if current_profit >= 0.04:
                return "side_take_profit"

        # --------------------------------------
        # 3. 熊市：趋势行情
        # --------------------------------------
        elif state == 3:
            # 止损：约 -4%
            if current_profit <= -0.04:
                return "bear_stop_loss"

            # 如果已经吃到比较大的趋势 (>3%)，且出现 exit_signal，则落袋为安
            if exit_signal == 1 and current_profit > 0.03:
                return "exit_signal"

        # --------------------------------------
        # 4. 超时强制平仓
        # --------------------------------------
        if trade_duration > max_hours:
            # 有利润就直接平仓
            if current_profit > 0:
                return f"timeout_exit_({int(max_hours)}h)"
            # 轻微亏损也不要再拖
            if current_profit > -0.01:
                return f"timeout_exit_({int(max_hours)}h)"

        # 不触发任何硬退出
        return None

    # === 8. 动态移动止盈 ===
    def custom_stoploss(
        self,
        pair: str,
        trade: Trade,
        current_time: datetime,
        current_rate: float,
        current_profit: float,
        **kwargs,
    ) -> float:
        """根据浮盈动态上调止损价位，保护利润。

        返回值是相对于开仓价的最大允许亏损（负数）。
        这里采用分段式移动止盈：
        - 盈利 > 1.5%  -> 止损抬到 +0.2%（接近保本）
        - 盈利 > 3%    -> 锁定 2%
        - 盈利 > 6%    -> 锁定 3%
        - 盈利 > 10%   -> 锁定 7%
        其它情况返回 1，让基础 stoploss / custom_exit 去处理。
        """

        # 阶段 1：赚 1.5% -> 基本保本
        if current_profit > 0.015:
            return 0.002

        # 阶段 2：赚 3% -> 锁定 2%
        if current_profit > 0.03:
            return 0.02

        # 阶段 3：赚 6% -> 锁定 3%
        if current_profit > 0.06:
            return 0.03

        # 阶段 4：赚 10% -> 锁定 7%
        if current_profit > 0.10:
            return 0.07

        # 其它情况交给 custom_exit / 基础 stoploss
        return 1