from freqtrade.freqai.prediction_models.ReinforcementLearner import ReinforcementLearner
from freqtrade.freqai.RL.Base5ActionRLEnv import Actions, Base5ActionRLEnv, Positions
import numpy as np


class MyCoolRLModel(ReinforcementLearner):
    """
    用户创建的RL预测模型。
    
    保存此文件到 `freqtrade/user_data/freqaimodels`
    
    然后使用以下命令运行：
    freqtrade trade --freqaimodel MyCoolRLModel --config config.json --strategy FreqaiPrimer
    
    在这里，用户可以覆盖 `IFreqaiModel` 继承树中的任何函数。
    对于RL来说，最重要的是在这里覆盖 `MyRLEnv`，以定义自定义的
    `calculate_reward()` 函数，或覆盖环境的任何其他部分。
    
    此类还允许用户覆盖 IFreqaiModel 树中的任何其他部分。
    例如，用户可以覆盖 `def fit()` 或 `def train()` 或 `def predict()`
    以对这些过程进行精细控制。
    
    另一个常见的覆盖可能是 `def data_cleaning_predict()`，用户可以
    在这里对数据处理管道进行精细控制。
    """
    
    class MyRLEnv(Base5ActionRLEnv):
        """
        用户自定义环境。此类继承自 BaseEnvironment 和 gym.Env。
        用户可以覆盖这些父类中的任何函数。下面是一个
        用户自定义 `calculate_reward()` 函数的示例。
        
        警告！
        此函数是功能展示，旨在展示尽可能多的
        环境控制功能。它还设计用于在小型计算机上快速运行。
        这是一个基准，*不* 用于生产环境。
        """
        
        def calculate_reward(self, action: int) -> float:
            """
            奖励函数的示例。这是用户可能希望
            注入自己创意的一个函数。
            
            警告！
            此函数是功能展示，旨在展示尽可能多的
            环境控制功能。它还设计用于在小型计算机上快速运行。
            这是一个基准，*不* 用于生产环境。
            
            :param action: int = 智能体为当前K线做出的动作。
            :return:
            float = 给智能体当前步骤的奖励（用于优化
                神经网络中的权重）
            """
            # 首先，如果动作无效，则惩罚
            if not self._is_valid(action):
                self.tensorboard_log("invalid", category="actions")
                return -2.0

            # 获取核心状态数据
            pnl = self.get_unrealized_profit()  # 当前浮动盈亏 (百分比，如 0.01 代表 1%)
            
            # 获取持仓时间 (K线数量)
            trade_duration = self._current_tick - self._last_trade_tick
            
            # 读取配置中的最大持仓时间，默认 100 根 K 线
            max_trade_duration = self.rl_config.get("max_trade_duration_candles", 100)
            
            # 奖励累加器
            reward = 0.0
            
            # =========================================================
            # 场景 A: 决定入场 (Long Enter / Short Enter)
            # =========================================================
            if action in (Actions.Long_enter.value, Actions.Short_enter.value):
                if self._position == Positions.Neutral:
                    # 【核心修改 1】入场不给奖励，反而给微小的惩罚。
                    # 逻辑：开仓即产生手续费和滑点风险。
                    # 这迫使模型只有在预期收益远大于这个惩罚时才敢开仓。
                    return -0.05 

            # =========================================================
            # 场景 B: 观望 (Neutral)
            # =========================================================
            if action == Actions.Neutral.value and self._position == Positions.Neutral:
                # 【核心修改 2】空仓观望给 0 分，或者极小的惩罚。
                # 之前的代码给 -1 分会导致模型为了避免扣分而强行开单（Overtrading）。
                # 对于 PEPE 这种币，大部分时间是垃圾时间，观望是正确的。
                return 0.0

            # =========================================================
            # 场景 C: 持仓中 (Holding) - 每一根 K 线都会触发
            # =========================================================
            if self._position in (Positions.Short, Positions.Long):
                # 如果当前动作是继续持有 (Neutral)
                if action == Actions.Neutral.value:
                    
                    # 1. 时间衰减惩罚：拿得越久，资金占用成本越高
                    time_penalty = -0.01 * (trade_duration / max_trade_duration)
                    reward += time_penalty

                    # 2. 浮动盈亏反馈 (关键！)
                    if pnl > 0:
                        # 浮盈：给予微弱的正反馈，鼓励拿住趋势
                        # 使用 log 函数让奖励增长平缓，避免模型过于贪婪而不止盈
                        reward += np.log(1 + pnl) * 0.5
                    else:
                        # 浮亏：给予沉重的打击！
                        # 逻辑：浮亏不仅是钱的损失，更是机会成本。
                        # pnl 是负数，这里会让 reward 变得很小。
                        # 例如 pnl -0.02, 惩罚会显著大于浮盈 0.02 的奖励
                        reward -= (abs(pnl) * 2.0) 

                    # 3. 止损惩罚加速
                    # 如果浮亏超过 3%，给予额外的巨额惩罚，逼迫模型学会止损
                    if pnl < -0.03:
                        reward -= 1.0

            # =========================================================
            # 场景 D: 离场结算 (Exit)
            # =========================================================
            if (action == Actions.Long_exit.value and self._position == Positions.Long) or \
               (action == Actions.Short_exit.value and self._position == Positions.Short):
                
                # 基础因子
                factor = 10.0
                
                # 【核心修改 3】非线性结算奖励
                if pnl > 0:
                    # 盈利：奖励 = PnL * 因子
                    # 如果这笔交易是大赚 (比如 > 2%)，因子翻倍
                    if pnl > 0.02:
                        factor *= 2.0
                    
                    # 最终奖励
                    total_reward = pnl * factor
                    
                    # 额外奖励：如果这笔交易很快就赚了钱（高效率）
                    if trade_duration < (max_trade_duration * 0.2):
                         total_reward *= 1.5
                         
                    return float(total_reward)
                
                else:
                    # 亏损：惩罚 = PnL * 因子
                    # 亏损时的惩罚系数通常要比盈利系数大，模拟"损失厌恶"
                    # 这会让模型非常忌惮亏损离场
                    loss_factor = 15.0 
                    return float(pnl * loss_factor)

            return float(reward)