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

class OffPACAgent:
    def __init__(self, state_dim, action_dim, device, actor_lr=0.001, critic_lr=0.002, gamma=0.99, epsilon=0.1):
        self.device = device
        self.gamma = gamma
        self.action_dim = action_dim
        # epsilon 用于构建行为策略 beta 的探索率
        self.epsilon = epsilon 
        
        # 沿用 A2C 的网络结构
        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)
            # 获取目标策略 pi 的动作概率分布
            pi_probs = self.actor(state_tensor).squeeze(0)
            
            # 构建行为策略 beta：在 pi 的基础上加入 epsilon 的均匀随机噪声
            beta_probs = (1 - self.epsilon) * pi_probs + self.epsilon / self.action_dim
            
            # 根据行为策略 beta 采样动作，用于与环境交互
            action = distributions.Categorical(beta_probs).sample().item()
        return action

    def update(self, state, action, reward, next_state, next_action, done):
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
        next_state_tensor = torch.FloatTensor(next_state).unsqueeze(0).to(self.device)
        reward_tensor = torch.FloatTensor([reward]).unsqueeze(0).to(self.device)
        action_tensor = torch.tensor([action]).to(self.device)
        
        # --- 计算重要性采样权重 (Importance Weight) ---
        # 重新获取当前最新目标策略 pi 下的动作概率
        pi_probs = self.actor(state_tensor).squeeze(0)
        pi_prob_a = pi_probs[action] 
        
        # 重建行为策略 beta 选出该动作的概率 (作为分母)
        beta_prob_a = (1 - self.epsilon) * pi_prob_a.detach() + self.epsilon / self.action_dim
        
        # 计算重要性权重 rho = pi(a|s) / beta(a|s)
        rho = (pi_prob_a / beta_prob_a).detach()
        
        # 【工程防崩技巧】：如果两个策略偏差过大，rho 会极大导致梯度爆炸。
        # 工业界通用的做法是对 rho 进行截断 (Clip)，这也是 PPO 算法的前身思想。
        rho = torch.clamp(rho, 0.1, 10.0) 

        # --- Critic 更新 (Off-Policy) ---
        v_value = self.critic(state_tensor)
        next_v_value = self.critic(next_state_tensor).detach()
        
        # 计算 TD 目标和优势函数
        td_target = reward_tensor + self.gamma * next_v_value * (1 - int(done))
        advantage = td_target - v_value
        
        # Critic 损失加入了重要性权重 rho 
        critic_loss = (rho * F.mse_loss(v_value, td_target, reduction='none')).mean()
        
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()
        
        # --- Actor 更新 (Off-Policy) ---
        action_probs = self.actor(state_tensor)
        dist = distributions.Categorical(action_probs)
        log_prob = dist.log_prob(action_tensor)
        
        # Actor 梯度上升目标：rho * ln(pi) * Advantage 
        actor_loss = -(rho * log_prob * advantage.detach()).mean()
        
        self.actor_optimizer.zero_grad()
        self.actor_optimizer.step()