import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from typing import Tuple, List

class ActorCritic(nn.Module):
    """Actor-Critic网络"""
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super().__init__()
        self.actor = 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()
        )
        self.critic = 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):
        action = self.actor(state)
        value = self.critic(state)
        return action, value

class PPO:
    """近端策略优化算法"""
    def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99, eps_clip=0.2, epochs=10):
        self.gamma = gamma
        self.eps_clip = eps_clip
        self.epochs = epochs
        
        self.policy = ActorCritic(state_dim, action_dim)
        self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)
        self.policy_old = ActorCritic(state_dim, action_dim)
        self.policy_old.load_state_dict(self.policy.state_dict())
        
        self.memory = []
    
    def select_action(self, state):
        state = torch.FloatTensor(state).unsqueeze(0)
        action, _ = self.policy_old(state)
        return action.detach().numpy()[0]
    
    def store_transition(self, transition):
        self.memory.append(transition)
    
    def update(self):
        if len(self.memory) == 0:
            return
        
        # 准备数据
        states = torch.FloatTensor([t[0] for t in self.memory])
        actions = torch.FloatTensor([t[1] for t in self.memory])
        rewards = [t[2] for t in self.memory]
        dones = [t[3] for t in self.memory]
        next_states = torch.FloatTensor([t[4] for t in self.memory])
        
        # 计算回报和优势
        returns = []
        discounted_reward = 0
        for reward, done in zip(reversed(rewards), reversed(dones)):
            if done:
                discounted_reward = 0
            discounted_reward = reward + self.gamma * discounted_reward
            returns.insert(0, discounted_reward)
        returns = torch.FloatTensor(returns)
        
        # 旧策略的动作概率
        with torch.no_grad():
            old_actions, old_values = self.policy_old(states)
        
        # 多轮优化
        for _ in range(self.epochs):
            actions_pred, values = self.policy(states)
            
            # 计算比率
            ratio = torch.exp(actions_pred - old_actions.detach())
            
            # 计算优势
            advantages = returns - values.squeeze()
            
            # PPO裁剪目标
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1 - self.eps_clip, 1 + self.eps_clip) * advantages
            actor_loss = -torch.min(surr1, surr2).mean()
            
            # 评论家损失
            critic_loss = advantages.pow(2).mean()
            
            # 总损失
            loss = actor_loss + 0.5 * critic_loss
            
            # 反向传播
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
        
        # 更新旧策略
        self.policy_old.load_state_dict(self.policy.state_dict())
        self.memory = []