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

class ValueNetwork(nn.Module):
    """
    状态值函数网络 (Critic)，用于估计状态的内在价值 V(s)
    增加了隐藏层容量，以确保在多轮迭代中能够更准确地拟合优势函数
    """
    def __init__(self, state_dim, hidden_dim=128):
        super(ValueNetwork, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, 1)
        )

    def forward(self, state):
        # 返回对当前状态的价值评估
        return self.net(state)


class PolicyNetwork(nn.Module):
    """
    策略网络 (Actor)，输出高斯分布的均值和标准差，适用于连续动作空间
    """
    def __init__(self, state_dim, action_dim, action_bound, hidden_dim=128):
        super(PolicyNetwork, self).__init__()
        self.action_bound = action_bound # 环境允许的最大物理动作幅度
        
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, action_dim),
            nn.Tanh() # 关键修复：强制均值输出在 [-1, 1] 之间，防止动作空间爆炸
        )
        # 将初始对数标准差设为 -0.5 (对应的标准差约为 0.6)
        # 较小的初始方差有助于防止初期探索步子迈得太大导致系统崩溃
        self.action_log_std = nn.Parameter(torch.full((1, action_dim), -0.5))

    def forward(self, state):
        # 计算动作均值并将其缩放到实际的物理边界内
        action_mean = self.net(state) * self.action_bound
        # 限制标准差的范围以提高数值稳定性
        action_std = torch.exp(self.action_log_std.expand_as(action_mean))
        return action_mean, action_std

    def evaluate(self, state):
        # 评估当前状态，返回构建好的高斯动作分布
        mean, std = self.forward(state)
        dist = Normal(mean, std)
        return dist