import torch
import torch.optim as optim
import torch.nn.functional as F
import torch.distributions as distributions
from networks import Actor, VCritic

class A2CAgent:
    def __init__(self, state_dim, action_dim, device, actor_lr=0.001, critic_lr=0.002, gamma=0.99):
        self.device = device
        self.gamma = gamma
        
        # A2C 使用 VCritic
        self.actor = Actor(state_dim, action_dim).to(self.device)
        self.critic = VCritic(state_dim).to(self.device)
        
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=actor_lr)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=critic_lr)

    def select_action(self, state):
        with torch.no_grad():
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
            action_probs = self.actor(state_tensor)
            action = distributions.Categorical(action_probs).sample().item()
        return action

    # 注意：A2C 的更新不需要 next_action，但为了与 QAC 的接口统一，此处用 *args 吸收多余参数
    def update(self, state, action, reward, next_state, next_action, done):
        state = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        next_state = torch.FloatTensor(next_state).unsqueeze(0).to(self.device)
        reward = torch.FloatTensor([reward]).unsqueeze(0).to(self.device)
        
        # --- Critic 更新 ---
        v_value = self.critic(state)
        next_v_value = self.critic(next_state).detach()
        
        td_target = reward + self.gamma * next_v_value * (1 - int(done))
        # 计算优势函数 (Advantage)
        advantage = td_target - v_value
        
        critic_loss = F.mse_loss(v_value, td_target)
        
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()
        
        # --- Actor 更新 ---
        action_probs = self.actor(state)
        dist = distributions.Categorical(action_probs)
        log_prob = dist.log_prob(torch.tensor([action]).to(self.device))

        # 计算策略的熵，鼓励探索
        entropy = dist.entropy()
        
        # Actor 梯度上升目标：ln(pi) * Advantage，方差更小
        actor_loss = -(log_prob * advantage.detach()).mean()- 0.01 * entropy.mean()
        
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()