"""PPO Trainer with GAE advantage estimation."""
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import numpy as np


class PPOTrainer:
    """PPO trainer handling the training loop."""

    def __init__(
        self,
        actor,
        critic,
        rollout_buffer,
        device,
        clip_eps=0.2,
        gamma=0.99,
        gae_lambda=0.95,
        lr=3e-4,
        ent_coef=0.01,
        vf_coef=0.5,
        max_grad_norm=0.5,
        ppo_epochs=4,
        mini_batch_size=64,
    ):
        self.actor = actor
        self.critic = critic
        self.buffer = rollout_buffer
        self.device = device
        self.clip_eps = clip_eps
        self.gamma = gamma
        self.gae_lambda = gae_lambda
        self.ent_coef = ent_coef
        self.vf_coef = vf_coef
        self.max_grad_norm = max_grad_norm
        self.ppo_epochs = ppo_epochs
        self.mini_batch_size = mini_batch_size

        # Separate optimizers
        self.actor_optim = optim.Adam(actor.parameters(), lr=lr)
        self.critic_optim = optim.Adam(critic.parameters(), lr=lr)

        self.loss_history = {'actor': [], 'critic': [], 'entropy': [], 'total': []}

    def update(self, last_value):
        """Perform one PPO update."""
        states, actions, rewards, dones, values, log_probs_old = self.buffer.get()

        # Compute returns and advantages
        returns, advantages = self.buffer.compute_returns(
            last_value, self.gamma, self.gae_lambda
        )

        # Normalize advantages
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        # Convert to tensors
        states_t = torch.from_numpy(states).float().to(self.device)
        actions_t = torch.from_numpy(actions).float().to(self.device)
        log_probs_old_t = torch.from_numpy(log_probs_old).float().to(self.device)
        returns_t = torch.from_numpy(returns).float().to(self.device)
        advantages_t = torch.from_numpy(advantages).float().to(self.device)

        dataset = TensorDataset(states_t, actions_t, log_probs_old_t, returns_t, advantages_t)
        loader = DataLoader(dataset, batch_size=self.mini_batch_size, shuffle=True)

        total_actor_loss = 0
        total_critic_loss = 0
        total_entropy = 0
        count = 0

        for _ in range(self.ppo_epochs):
            for batch in loader:
                s, a, log_pi_old, ret, adv = batch

                # Get current policy distribution
                mu, std = self.actor(s)
                dist = torch.distributions.Normal(mu, std)
                log_pi = dist.log_prob(a).sum(dim=-1, keepdim=True)
                entropy = dist.entropy().sum(dim=-1, keepdim=True)

                # Probability ratio
                ratio = torch.exp(log_pi - log_pi_old)

                # Clipped surrogate objective
                surr1 = ratio * adv
                surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * adv
                actor_loss = -torch.min(surr1, surr2).mean()

                # Value loss
                value = self.critic(s)
                critic_loss = nn.MSELoss()(value.squeeze(), ret)

                # Total loss
                loss = actor_loss + self.vf_coef * critic_loss - self.ent_coef * entropy.mean()

                # Update
                self.actor_optim.zero_grad()
                self.critic_optim.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(self.actor.parameters(), self.max_grad_norm)
                nn.utils.clip_grad_norm_(self.critic.parameters(), self.max_grad_norm)
                self.actor_optim.step()
                self.critic_optim.step()

                total_actor_loss += actor_loss.item()
                total_critic_loss += critic_loss.item()
                total_entropy += entropy.mean().item()
                count += 1

        avg_actor = total_actor_loss / count
        avg_critic = total_critic_loss / count
        avg_entropy = total_entropy / count

        self.loss_history['actor'].append(avg_actor)
        self.loss_history['critic'].append(avg_critic)
        self.loss_history['entropy'].append(avg_entropy)
        self.loss_history['total'].append(avg_actor + avg_critic)

        self.buffer.reset()
        return avg_actor, avg_critic, avg_entropy