# -*- coding: utf-8 -*-
"""
RationalShort - 三态理性做空策略（MathModel V2）

改动目标：
1. 用“期望 E = p * W - (1-p) * L”约束策略结构：
   - 胜率 p 可以在 40%~60% 波动；
   - 但要求平均盈利 W 至少 ~2%，平均亏损 L 控制在 ~2% 左右。
2. 尽量放大利润单、缩小亏损单：
   - 止损改为 ATR + 上限控制，避免 -4%~-5% 的大亏损；
   - exit_signal / 趋势行情允许利润拉大，再用动态移动止盈出场；
   - RegimeFlip 保护：Side/Bear 单被拉成 Bull 时尽快小亏/打平离场。
3. 保留“三态模型（Bull / Side / Bear）”的大框架结构，方便后续用 hyperopt 做参数搜索。

注意：
- 该版本适用于 freqtrade 2025.10，期货 / 合约环境（can_short = True）。
- 若与旧版本共存，请将文件名改为 `RationalShortV2.py`，类名改为 `RationalShortV2`，
  并在命令行用 `--strategy RationalShortV2` 运行。
"""

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 RationalShortv2(IStrategy):
    """理性做空策略 (Rational Short System) - MathModel V2"""

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

    timeframe = "1h"
    can_short = True

    # 兜底硬止损（真正逻辑由 custom_stoploss / custom_exit 接管）
    stoploss = -0.15

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

    # === 2. 参数区（可后续用 hyperopt 调） ===
    # ADX 阈值，区分震荡 / 趋势
    adx_threshold = IntParameter(20, 40, default=25, space="buy")

    # 不同状态下的 RSI 入场门槛
    rsi_bull = IntParameter(80, 90, default=84, space="buy")   # 牛市：极端超买
    rsi_side = IntParameter(60, 80, default=68, space="buy")   # 震荡：正常超买
    rsi_bear = IntParameter(45, 70, default=58, space="buy")   # 熊市：反弹就空

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

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

    # ATR 止损倍数（配合上限使用）
    atr_stop_mult = DecimalParameter(2.0, 4.0, default=3.0, space="sell")

    # === 3. 保护机制 ===
    @property
    def protections(self):
        return [
            {"method": "CooldownPeriod", "stop_duration_candles": 1},
            {
                "method": "StoplossGuard",
                "lookback_period_candles": 24,
                "trade_limit": 5,
                "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. ATR：用于动态止损
        dataframe["atr"] = ta.ATR(dataframe, timeperiod=14)

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

        # F. 标记市场状态 (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

        # G. 预备列
        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)

        # === 牛市：极端超买 + 短期爆拉，逆势短线做空（scalping） ===
        cond_bull = (
            base_cond &
            (dataframe["market_state"] == 1) &
            (dataframe["rsi"] > int(self.rsi_bull.value)) &
            (dataframe["roc_3"] > 4)
        )

        # === 震荡：正常超买，但要求价格不在 EMA200 上方太远 ===
        cond_side = (
            base_cond &
            (dataframe["market_state"] == 2) &
            (dataframe["rsi"] > int(self.rsi_side.value)) &
            (dataframe["close"] <= dataframe["ema_200"] * 1.01)
        )

        # === 熊市：略微反弹就做空 ===
        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

    # === 工具函数：根据 ATR 计算当前状态下的最大允许亏损 ===
    def _calc_max_loss_pct(self, state: int, last_candle: DataFrame) -> float:
        atr = float(last_candle.get("atr", 0))
        close = float(last_candle.get("close", 0)) or 1.0
        atr_pct = atr / close if close > 0 else 0.01

        mult = float(self.atr_stop_mult.value)
        raw = atr_pct * mult

        # 不同状态给不同上限（单位：小数，0.02=2%）
        if state == 1:      # Bull：只允许小亏
            max_cap = 0.012
        elif state == 2:    # Side：中等亏损
            max_cap = 0.02
        else:               # Bear：略宽一点
            max_cap = 0.025

        # 同时保证不低于 0.006（避免极小 ATR 导致随便被打掉）
        return max(0.006, min(raw, max_cap))

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

        dataframe, _ = self.dp.get_analyzed_dataframe(pair, self.timeframe)
        try:
            df = dataframe.loc[dataframe["date"] <= current_time]
            last_candle = df.iloc[-1] if not df.empty else dataframe.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

        # 动态最大允许亏损（基于 ATR + 上限）
        max_loss_pct = self._calc_max_loss_pct(state, last_candle)

        enter_tag = getattr(trade, "enter_tag", "") or ""

        # --------------------------------------
        # 0. Regime Flip 保护：Side/Bear 单被拉成 Bull
        # --------------------------------------
        if state == 1 and enter_tag not in ("bull", ""):
            # 状态切换到 Bull，且这单不是在 Bull 下开的
            # 小亏 / 打平就走，避免被拖成大亏
            if current_profit > -0.003:  # -0.3% 以内
                return "regime_flip"
            # 更深的亏损交给后面的止损逻辑处理

        # --------------------------------------
        # 1. 牛市：短线 scalping
        # --------------------------------------
        if state == 1:
            # 小止损：基于 ATR 的动态值
            if current_profit <= -max_loss_pct:
                return "bull_stop_loss"

            # 小止盈：目标 1.5% 左右
            if current_profit >= 0.015:
                return "bull_scalp_profit"

            # 如果已经有利润，并且触发布林中轨 exit 信号，则落袋
            if exit_signal == 1 and current_profit > 0.005:
                return "exit_signal"

        # --------------------------------------
        # 2. 震荡：布林带回归
        # --------------------------------------
        elif state == 2:
            # 止损：动态 + 上限
            if current_profit <= -max_loss_pct:
                return "stop_loss"

            # exit_signal 触发且 >0.5% 就离场，保证盈利单不是太小
            if exit_signal == 1 and current_profit > 0.005:
                return "exit_signal"

            # 兜底止盈：> 3.5% 直接平仓
            if current_profit >= 0.035:
                return "side_trend_profit"

        # --------------------------------------
        # 3. 熊市：趋势行情
        # --------------------------------------
        elif state == 3:
            # 止损：动态 + 上限
            if current_profit <= -max_loss_pct:
                return "bear_stop_loss"

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

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

        # 不触发任何硬退出
        return None

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

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

        # 阶段 1：赚 2% -> 锁定 0.5%
        if current_profit > 0.02:
            return 0.005

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

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

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

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