import torch
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
from networks_cont import ContActor, ContQCritic

class DPACAgent:
    def __init__(self, state_dim, action_dim, max_action, device, actor_lr=0.001, critic_lr=0.002, gamma=0.99):
        self.device = device
        self.gamma = gamma
        self.max_action = max_action
        
        self.actor = ContActor(state_dim, action_dim, max_action).to(self.device)
        self.critic = ContQCritic(state_dim, action_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 = self.actor(state_tensor).cpu().data.numpy().flatten()
            
            # 探索噪声 (行为策略 beta 与目标策略 mu 的区别就在这里)
            noise = np.random.normal(0, 0.1 * self.max_action, size=action.shape)
            action = np.clip(action + noise, -self.max_action, self.max_action)
        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.FloatTensor(action).unsqueeze(0).to(self.device) 
        
        # --- Critic 更新 (Algorithm 10.4 核心逻辑) ---
        # 1. 目标策略 mu 在下一个状态的理想输出
        next_mu_action = self.actor(next_state_tensor).detach()
        # 2. 评估这个理想动作的 Q 值
        next_q_value = self.critic(next_state_tensor, next_mu_action).detach()
        # 3. 计算 TD 目标
        td_target = reward_tensor + self.gamma * next_q_value * (1 - int(done))
        
        # 4. 当前实际采取动作的 Q 值
        current_q_value = self.critic(state_tensor, action_tensor)
        
        critic_loss = F.mse_loss(current_q_value, td_target)
        
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()
        
        # --- Actor 更新 (链式法则) ---
        # 1. 目标策略当前状态的输出
        mu_action = self.actor(state_tensor)
        
        # 2. 拿到这个动作去问 Critic："给我打分"。
        # 为了让 Q 值最大化，我们加负号转化为梯度下降
        actor_loss = -self.critic(state_tensor, mu_action).mean()
        
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

    def save(self, path):
        torch.save({
            'actor': self.actor.state_dict(),
            'critic': self.critic.state_dict(),
        }, path)

    def load(self, path):
        checkpoint = torch.load(path, map_location=self.device)
        self.actor.load_state_dict(checkpoint['actor'])
        self.critic.load_state_dict(checkpoint['critic'])