import torch
import torch.nn.functional as F
import torch.optim as optim
from models.networks import Actor, Critic
import copy

class SAC(object):
    def __init__(self, state_dim, action_dim, max_action, config):
        """
        初始化 Soft Actor-Critic 算法
        """
        # 设备配置 (CPU 或 GPU)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # 从 config 字典中加载超参数
        self.gamma = config.get('gamma', 0.99)        # 折扣因子
        self.tau = config.get('tau', 0.005)           # 目标网络软更新系数 (论文公式 9 下方)
        self.alpha = config.get('alpha', 0.2)         # 熵的温度参数 (控制探索的随机性)
        
        # 1. 实例化策略网络 (Actor)
        self.actor = Actor(state_dim, action_dim, max_action).to(self.device)
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=config.get('lr', 3e-4))
        
        # 2. 实例化价值网络 (Critic) - 内部已经包含了 Q1 和 Q2
        self.critic = Critic(state_dim, action_dim).to(self.device)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=config.get('lr', 3e-4))
        
        # 3. 实例化目标价值网络 (Target Critic)
        # 用 copy.deepcopy 完美复制一份初始参数，并冻结其梯度计算
        self.critic_target = copy.deepcopy(self.critic)
        for param in self.critic_target.parameters():
            param.requires_grad = False

    def select_action(self, state, evaluate=False):
        """
        与环境交互时使用的动作选择函数
        """
        # 将输入的状态转换为 PyTorch Tensor
        state = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        
        # 论文技巧：评估（测试）时使用均值动作，训练时使用采样动作
        with torch.no_grad():
            if evaluate:
                _, _, action = self.actor.sample(state) # 第三个返回值是均值
            else:
                action, _, _ = self.actor.sample(state) # 第一个返回值是加了噪声的采样值
                
        # 转换回 numpy 数组，送给 Gym 环境执行
        return action.cpu().data.numpy().flatten()

    def update(self, replay_buffer, batch_size):
        """
        算法的核心心跳：从经验池采样并更新神经网络参数
        """
        # 从 Replay Buffer 中随机抽取一个 Batch 的数据
        state, action, reward, next_state, not_done = replay_buffer.sample(batch_size)

        # ================================================================= #
        #                           1. 更新 Critic                          #
        # ================================================================= #
        with torch.no_grad():
            # 拿到下一个状态的动作和其对应的对数概率 (用于计算熵)
            next_action, next_log_prob, _ = self.actor.sample(next_state)
            
            # 使用目标网络计算下一个状态的 Q 值 (Q1 和 Q2)
            target_Q1, target_Q2 = self.critic_target(next_state, next_action)
            
            # 【核心对抗高估】：取两个 Q 值的最小值
            target_Q = torch.min(target_Q1, target_Q2)
            
            # 【软贝尔曼方程】：目标 Q 值 = 奖励 + gamma * (目标 Q - alpha * 熵)
            # 注意这里加上了 -self.alpha * next_log_prob，这就是论文中“把熵当做奖励”的体现
            target_Q = reward + not_done * self.gamma * (target_Q - self.alpha * next_log_prob)

        # 获取当前状态和动作对应的 Q 值预测
        current_Q1, current_Q2 = self.critic(state, action)

        # 计算 Critic 的损失 (均方误差 MSE)
        critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q)

        # 优化 Critic 网络
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        # ================================================================= #
        #                           2. 更新 Actor                           #
        # ================================================================= #
        # 让当前 Actor 对**当前状态**重新采样一个动作 (注意：不能用 Buffer 里的旧动作)
        pi_action, log_prob, _ = self.actor.sample(state)
        
        # 拿到更新后的 Critic 对这个新动作的打分
        q1_pi, q2_pi = self.critic(state, pi_action)
        min_q_pi = torch.min(q1_pi, q2_pi)

        # 计算 Actor 的损失：最小化 (alpha * log_prob - min_Q)
        # 等价于最大化 (min_Q - alpha * log_prob) -> 既要 Q 值大，又要熵大（分布广）
        actor_loss = (self.alpha * log_prob - min_q_pi).mean()

        # 优化 Actor 网络
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # ================================================================= #
        #                    3. 软更新 Target Critic                        #
        # ================================================================= #
        # 使用 EMA (指数移动平均) 缓慢将前线 Critic 的参数移交给 Target Critic
        for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)