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

class ActorNet(nn.Module):
    """
    策略网络 (Actor)：输入环境状态，输出连续动作的高斯分布参数 (均值和标准差)
    """
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super(ActorNet, self).__init__()
        # 定义两层隐藏层，提取状态特征
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        
        # 输出层：计算动作的均值 (Mean)
        self.mean_layer = nn.Linear(hidden_dim, action_dim)
        
        # 定义对数标准差 (Log Standard Deviation) 为可训练的独立参数
        # 注意：在 TRPO/PPO 中，通常让标准差独立于状态，这能让训练更稳定
        # 初始值设为 0，意味着初始标准差为 exp(0) = 1.0
        self.log_std = nn.Parameter(torch.zeros(1, action_dim))

    def forward(self, state):
        # 前向传播提取特征
        x = F.tanh(self.fc1(state))
        x = F.tanh(self.fc2(x))
        
        # 计算均值
        mean = self.mean_layer(x)
        
        # 将对数标准差扩展到与 batch size 相同的维度
        log_std = self.log_std.expand_as(mean)
        # 转化为标准差
        std = torch.exp(log_std)
        
        return mean, std

    def get_action(self, state):
        """
        根据当前状态采样动作，并返回对应的对数概率 (log probability)
        """
        # 获取当前状态的均值和标准差
        mean, std = self.forward(state)
        
        # 构建正态分布 (高斯分布)
        dist = Normal(mean, std)
        
        # 从分布中采样一个动作
        action = dist.sample()
        
        # 返回采样动作及其对数概率 (后续计算重要性采样权重时必须用到)
        return action, dist.log_prob(action)


class CriticNet(nn.Module):
    """
    价值网络 (Critic)：输入环境状态，评估该状态的预期收益 (标量 V 值)
    """
    def __init__(self, state_dim, hidden_dim=64):
        super(CriticNet, self).__init__()
        # 定义隐藏层
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        
        # 输出层：计算状态价值 V(s)，输出维度为 1
        self.value_layer = nn.Linear(hidden_dim, 1)

    def forward(self, state):
        # 前向传播计算价值
        x = F.tanh(self.fc1(state))
        x = F.tanh(self.fc2(x))
        value = self.value_layer(x)
        return value