import torch
import torch.nn as nn
import torch.optim as optim

# Neural Network for Actor-Critic
import torch.nn as nn
import torch.nn.functional as F
from nanoppo.policy.actor_critic import ActorCritic
import wandb


# PPO Agent
class PPOAgent:
    def __init__(
        self,
        state_dim,
        action_dim,
        n_latent_var,
        policy_lr,
        value_lr,
        betas,
        gamma,
        K_epochs,
        eps_clip,
        state_normalizer,
        action_low,
        action_high,
        vl_coef=0.5,
        el_coef=0.001,
        device="cpu",
        wandb_log=False,
    ):
        self.gamma = gamma
        self.eps_clip = eps_clip
        self.K_epochs = K_epochs
        self.state_normalizer = state_normalizer
        self.action_low = action_low
        self.action_high = action_high
        self.vl_coef = vl_coef
        self.el_coef = el_coef
        self.device = device
        action_low_tensor = torch.tensor(action_low, dtype=torch.float32).to(device)
        action_high_tensor = torch.tensor(action_high, dtype=torch.float32).to(device)
        self.policy = ActorCritic(
            state_dim, action_dim, n_latent_var, action_low_tensor, action_high_tensor
        ).float().to(device)
        self.policy_old = ActorCritic(
            state_dim, action_dim, n_latent_var, action_low_tensor, action_high_tensor
        ).float().to(device)
        self.policy_old.load_state_dict(self.policy.state_dict())

        # Separate the parameters of the actor and critic networks
        actor_params = list(self.policy.action_mu.parameters()) + list(
            self.policy.action_log_std.parameters()
        )
        critic_params = list(self.policy.value_layer.parameters())

        # Create an optimizer with different learning rates for the actor and critic
        self.optimizer = torch.optim.Adam(
            [
                {"params": actor_params, "lr": policy_lr},
                {"params": critic_params, "lr": value_lr},
            ],
            betas=betas,
        )

        self.mse_loss = nn.MSELoss()
        self.wandb_log = wandb_log

    def update(self, states, actions, returns, next_states, dones):
        for _ in range(self.K_epochs):
            # Getting predicted values and log probs for given states and actions
            logprobs, state_values = self.policy.evaluate(states, actions)
            # Calculate the advantages
            advantages = returns - state_values.detach()
            # Normalize the advantages (optional, but can help in training stability)
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-5)

            # Compute ratio for PPO
            old_logprobs, _ = self.policy_old.evaluate(states, actions)
            ratio = torch.exp(logprobs - old_logprobs.detach())

            # Surrogate loss
            surr1 = ratio * advantages
            surr2 = (
                torch.clamp(ratio, 1 - self.eps_clip, 1 + self.eps_clip) * advantages
            )
            policy_loss = -torch.min(surr1, surr2).mean()

            # Value loss
            value_loss = self.vl_coef * self.mse_loss(state_values, returns)

            # Entropy (for exploration)
            entropy_loss = -self.el_coef * logprobs.mean()

            # Total loss
            loss = policy_loss + value_loss + entropy_loss

            # Optimize policy network
            self.optimizer.zero_grad()
            loss.backward()

            # Clip gradients to prevent exploding gradients
            # Large gradients can cause the weights to update too aggressively. 
            # Actions can have extream values e.g.(100, -0.10, ...)
            torch.nn.utils.clip_grad_norm_(self.policy.parameters(), max_norm=1.0)
        
            self.optimizer.step()
            if self.wandb_log:
                wandb.log(
                    {
                        "policy_loss": policy_loss.item(),
                        "value_loss": value_loss.item(),
                        "entropy_loss": entropy_loss.item(),
                        "total_loss": loss.item(),
                    }
                )

        # Copy new weights into old policy
        self.policy_old.load_state_dict(self.policy.state_dict())

    def save(self, path):
        # Saving
        torch.save(
            {
                "model_state_dict": self.policy.state_dict(),
                "optimizer_state_dict": self.optimizer.state_dict(),
                "state_normalizer_state": self.state_normalizer.get_state(),
            },
            path,
        )

    def load(self, path):
        checkpoint = torch.load(path, map_location=self.device)
        self.policy.load_state_dict(checkpoint["model_state_dict"])
        self.policy_old.load_state_dict(self.policy.state_dict())
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        self.state_normalizer.set_state(checkpoint["state_normalizer_state"])
