import gymnasium as gym  
 from src.utils.logger import logger  
 from src.models.callback import PolicyGradientLossCallback  
 from pydantic import Field, BaseModel, ConfigDict  
 from typing import Dict, List  
 import numpy as np  
 import os  
 from pathlib import Path  
 import torch  
 import torch.nn as nn  
 import torch.optim as optim  
 import torch.nn.functional as F  
 from torch.distributions import Normal  
   
 '''演員網(wǎng)絡(luò)：估計均值和對數(shù)標準差用于熵正則化計算'''  
   
 class Actor(nn.Module):  
     def __init__(self,state_dim,action_dim):  
         super(Actor,self).__init__()  
   
         self.net = nn.Sequential(  
             nn.Linear(state_dim, 100),  
             nn.ReLU(),  
             nn.Linear(100,100),  
             nn.ReLU()  
        )  
         self.mean_linear = nn.Linear(100, action_dim)  
         self.log_std_linear = nn.Linear(100, action_dim)  
   
     def forward(self, state):  
         x = self.net(state)  
         mean = self.mean_linear(x)  
         log_std =self.log_std_linear(x)  
         log_std = torch.clamp(log_std, min=-20, max=2)  
         return mean, log_std  
       
     def sample(self, state):  
         mean, log_std = self.forward(state)  
         std = log_std.exp()  
         normal = Normal(mean, std)  
         x_t = normal.rsample() # 重參數(shù)化技巧  
         y_t = torch.tanh(x_t)  
         action = y_t  
         log_prob = normal.log_prob(x_t)  
         log_prob -= torch.log(1-y_t.pow(2)+1e-6)  
         log_prob = log_prob.sum(dim=1, keepdim =True)  
   
         return action, log_prob

'''評論家網(wǎng)絡(luò)：定義q1和q2'''  
 class Critic(nn.Module):  
     def __init__(self, state_dim, action_dim):  
         super(Critic, self).__init__()  
   
         # Q1網(wǎng)絡(luò)架構(gòu)  
         self.q1_net = nn.Sequential(  
             nn.Linear(state_dim + action_dim, 256),  
             nn.ReLU(),  
             nn.Linear(256, 256),  
             nn.ReLU(),  
             nn.Linear(256, 1),  
        )  
   
         # Q2網(wǎng)絡(luò)架構(gòu)  
         self.q2_net = nn.Sequential(  
             nn.Linear(state_dim + action_dim, 256),  
             nn.ReLU(),  
             nn.Linear(256, 256),  
             nn.ReLU(),  
             nn.Linear(256, 1),  
        )  
   
     def forward(self, state, action):  
         sa = torch.cat([state, action], dim=1)  
         q1 = self.q1_net(sa)  
         q2 = self.q2_net(sa)  
         return q1, q2

'''SAC智能體的實現(xiàn)：整合演員網(wǎng)絡(luò)和評論家網(wǎng)絡(luò)'''  
   
 class SACAgent:  
     def __init__(self, state_dim, action_dim, learning_rate, device):  
         self.device = device  
   
         self.actor = Actor(state_dim, action_dim).to(device)  
         self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=learning_rate)  
   
         self.critic = Critic(state_dim, action_dim).to(device)  
         self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=learning_rate)  
   
         # 目標網(wǎng)絡(luò)初始化  
         self.critic_target = Critic(state_dim, action_dim).to(device)  
         self.critic_target.load_state_dict(self.critic.state_dict())  
   
         # 熵溫度參數(shù)  
         self.target_entropy = -action_dim  
         self.log_alpha = torch.zeros(1, requires_grad=True, device=device)  
         self.alpha_optimizer = optim.Adam([self.log_alpha], lr=learning_rate)  
   
     def select_action(self, state, evaluate=False):  
         state = torch.FloatTensor(state).to(self.device).unsqueeze(0)  
         if evaluate:  
             with torch.no_grad():  
                 mean, _ = self.actor(state)  
                 action = torch.tanh(mean)  
                 return action.cpu().numpy().flatten()  
         else:  
             with torch.no_grad():  
                 action, _ = self.actor.sample(state)  
                 return action.cpu().numpy().flatten()  
   
     def update(self, replay_buffer, batch_size=256, gamma=0.99, tau=0.005):  
         # 從經(jīng)驗回放中采樣訓(xùn)練數(shù)據(jù)  
         batch = replay_buffer.sample_batch(batch_size)  
         state = torch.FloatTensor(batch['state']).to(self.device)  
         action = torch.FloatTensor(batch['action']).to(self.device)  
         reward = torch.FloatTensor(batch['reward']).to(self.device)  
         next_state = torch.FloatTensor(batch['next_state']).to(self.device)  
         done = torch.FloatTensor(batch['done']).to(self.device)  
   
         # 評論家網(wǎng)絡(luò)更新  
         with torch.no_grad():  
             next_action, next_log_prob = self.actor.sample(next_state)  
             q1_next, q2_next = self.critic_target(next_state, next_action)  
             q_next = torch.min(q1_next, q2_next) - torch.exp(self.log_alpha) * next_log_prob  
             target_q = reward + (1 - done) * gamma * q_next  
   
         q1_current, q2_current = self.critic(state, action)  
         critic_loss = F.mse_loss(q1_current, target_q) + F.mse_loss(q2_current, target_q)  
   
         self.critic_optimizer.zero_grad()  
         critic_loss.backward()  
         self.critic_optimizer.step()  
   
         # 演員網(wǎng)絡(luò)更新  
         action_new, log_prob = self.actor.sample(state)  
         q1_new, q2_new = self.critic(state, action_new)  
         q_new = torch.min(q1_new, q2_new)  
         actor_loss = (torch.exp(self.log_alpha) * log_prob - q_new).mean()  
   
         self.actor_optimizer.zero_grad()  
         actor_loss.backward()  
         self.actor_optimizer.step()  
   
         # 溫度參數(shù)更新  
         alpha_loss = -(self.log_alpha * (log_prob + self.target_entropy).detach()).mean()  
   
         self.alpha_optimizer.zero_grad()  
         alpha_loss.backward()  
         self.alpha_optimizer.step()  
   
         # 目標網(wǎng)絡(luò)軟更新  
         for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):  
             target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)