import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal

# 确保代码可以在 GPU 上跑（如果有的话）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class Critic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(Critic, self).__init__()
        
        # Q1 网络架构
        self.l1 = nn.Linear(state_dim + action_dim, 256)
        self.l2 = nn.Linear(256, 256)
        self.l3 = nn.Linear(256, 1)

        # Q2 网络架构（和 Q1 完全一样，但参数是独立初始化的）
        self.l4 = nn.Linear(state_dim + action_dim, 256)
        self.l5 = nn.Linear(256, 256)
        self.l6 = nn.Linear(256, 1)

    def forward(self, state, action):
        # 把状态和动作拼接在一起，作为 Q 网络的输入
        sa = torch.cat([state, action], 1)

        # Q1 的前向传播
        q1 = F.relu(self.l1(sa))
        q1 = F.relu(self.l2(q1))
        q1 = self.l3(q1)

        # Q2 的前向传播
        q2 = F.relu(self.l4(sa))
        q2 = F.relu(self.l5(q2))
        q2 = self.l6(q2)

        # 训练时，我们需要同时返回两个 Q 值来算误差
        return q1, q2
    

# 定义标准差的上下界，防止网络输出极端值导致计算崩溃（NaN）
LOG_SIG_MAX = 2
LOG_SIG_MIN = -20

class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, max_action):
        super(Actor, self).__init__()
        
        # 共享特征提取层
        self.l1 = nn.Linear(state_dim, 256)
        self.l2 = nn.Linear(256, 256)
        
        # 均值输出层
        self.mean_linear = nn.Linear(256, action_dim)
        # 对数标准差输出层（预测 log_std 比直接预测 std 更好优化）
        self.log_std_linear = nn.Linear(256, action_dim)
        
        # 动作的最大物理边界（比如 Pendulum 的力矩最大是 2.0）
        self.max_action = max_action

    def forward(self, state):
        x = F.relu(self.l1(state))
        x = F.relu(self.l2(x))
        
        mean = self.mean_linear(x)
        log_std = self.log_std_linear(x)
        
        # 限制 log_std 的范围，防止数值不稳定
        log_std = torch.clamp(log_std, min=LOG_SIG_MIN, max=LOG_SIG_MAX)
        return mean, log_std

    def sample(self, state):
        mean, log_std = self.forward(state)
        std = log_std.exp()
        
        # 构造一个高斯分布
        normal = Normal(mean, std)
        
        # normal.rsample() 内部执行的就是 a = mean + std * epsilon (其中 epsilon 是标准正态噪声)
        # 这就是公式 (11) 的代码实现！用 rsample 才能让梯度传导回网络。
        x_t = normal.rsample() 
        
        # 把动作压缩到 [-1, 1] 区间（这就是论文附录 C 里的 tanh 压扁函数）
        y_t = torch.tanh(x_t)
        
        # 映射到真实的物理动作区间，比如 [-2.0, 2.0]
        action = y_t * self.max_action
        
        # 计算这个动作的对数概率 log(pi(a|s))，用于后面算熵
        # 这行公式对应论文附录 C 的公式 (21)，是应用 tanh 后的概率修正
        log_prob = normal.log_prob(x_t)
        log_prob -= torch.log(self.max_action * (1 - y_t.pow(2)) + 1e-6)
        log_prob = log_prob.sum(1, keepdim=True)
        
        # mean 经过 tanh 就是测试时用的确定性动作
        mean = torch.tanh(mean) * self.max_action
        
        return action, log_prob, mean