"""
PPO trainer for Module C: Intervention Policy.

Training stages:
  Stage 1 (Supervised warm-up): behavior cloning from a_recommend labels
  Stage 2 (PPO fine-tuning): optimize with multi-objective reward

PPO hyperparams (validated from prior work):
  clip_eps=0.2, lr=3e-4, entropy_coef=0.01
"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from typing import Dict, List, Optional
try:
    import wandb
except ImportError:
    wandb = None

from src.models.intervention_agent import InterventionAgent


class RolloutBuffer:
    """Stores PPO rollout trajectories."""

    def __init__(self, buffer_size: int, obs_dim: int, device: str = "cpu"):
        self.buffer_size = buffer_size
        self.obs_dim = obs_dim
        self.device = device
        self.reset()

    def reset(self):
        self.obs = torch.zeros(self.buffer_size, self.obs_dim)
        self.actions = torch.zeros(self.buffer_size, dtype=torch.long)
        self.log_probs = torch.zeros(self.buffer_size)
        self.rewards = torch.zeros(self.buffer_size)
        self.values = torch.zeros(self.buffer_size)
        self.dones = torch.zeros(self.buffer_size)
        self.ptr = 0
        self.full = False

    def add(
        self,
        obs: torch.Tensor,
        action: torch.Tensor,
        log_prob: torch.Tensor,
        reward: float,
        value: torch.Tensor,
        done: bool,
    ):
        self.obs[self.ptr] = obs
        self.actions[self.ptr] = action
        self.log_probs[self.ptr] = log_prob
        self.rewards[self.ptr] = float(reward)
        self.values[self.ptr] = value
        self.dones[self.ptr] = float(done)
        self.ptr = (self.ptr + 1) % self.buffer_size
        if self.ptr == 0:
            self.full = True

    def size(self) -> int:
        return self.buffer_size if self.full else self.ptr

    def compute_returns_and_advantages(
        self, gamma: float = 0.99, gae_lambda: float = 0.95
    ):
        n = self.size()
        advantages = torch.zeros(n)
        last_gae = 0.0

        for t in reversed(range(n)):
            if t + 1 < n:
                next_value = self.values[t + 1].item() * (1.0 - self.dones[t + 1].item())
            else:
                next_value = 0.0
            delta = (
                self.rewards[t].item()
                + gamma * next_value
                - self.values[t].item()
            )
            last_gae = delta + gamma * gae_lambda * (1.0 - self.dones[t].item()) * last_gae
            advantages[t] = last_gae

        returns = advantages + self.values[:n]
        return advantages.to(self.device), returns.to(self.device)

    def get(self) -> Dict[str, torch.Tensor]:
        n = self.size()
        return {
            "obs": self.obs[:n].to(self.device),
            "actions": self.actions[:n].to(self.device),
            "log_probs": self.log_probs[:n].to(self.device),
            "values": self.values[:n].to(self.device),
        }


class PPOTrainer:
    def __init__(
        self,
        agent: InterventionAgent,
        obs_dim: int,
        lr: float = 3e-4,
        clip_eps: float = 0.2,
        entropy_coef: float = 0.01,
        value_coef: float = 0.5,
        max_grad_norm: float = 0.5,
        gamma: float = 0.99,
        gae_lambda: float = 0.95,
        n_epochs: int = 4,
        batch_size: int = 64,
        buffer_size: int = 2048,
        device: str = "cpu",
        use_wandb: bool = True,
    ):
        self.agent = agent.to(device)
        self.optimizer = optim.Adam(agent.parameters(), lr=lr)
        self.device = device
        self.clip_eps = clip_eps
        self.entropy_coef = entropy_coef
        self.value_coef = value_coef
        self.max_grad_norm = max_grad_norm
        self.gamma = gamma
        self.gae_lambda = gae_lambda
        self.n_epochs = n_epochs
        self.batch_size = batch_size
        self.use_wandb = use_wandb
        self.buffer = RolloutBuffer(buffer_size, obs_dim, device)

    def behavior_cloning_warmup(
        self,
        obs_tensor: torch.Tensor,
        expert_actions: torch.Tensor,
        n_epochs: int = 5,
        lr: float = 1e-3,
    ) -> List[float]:
        """Stage 1: supervised pre-training to initialize policy."""
        optimizer = optim.Adam(self.agent.parameters(), lr=lr)
        losses = []
        dataset = torch.utils.data.TensorDataset(obs_tensor, expert_actions)
        loader = torch.utils.data.DataLoader(
            dataset, batch_size=self.batch_size, shuffle=True
        )

        for epoch in range(n_epochs):
            epoch_loss = 0.0
            self.agent.train()
            for obs_batch, act_batch in loader:
                obs_batch = obs_batch.to(self.device)
                act_batch = act_batch.to(self.device)
                loss = self.agent.behavior_clone_loss(obs_batch, act_batch)
                optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(self.agent.parameters(), self.max_grad_norm)
                optimizer.step()
                epoch_loss += loss.item()

            avg_loss = epoch_loss / max(len(loader), 1)
            losses.append(avg_loss)
            print(f"[BC] Epoch {epoch + 1}/{n_epochs}, Loss: {avg_loss:.4f}")
            if self.use_wandb:
                wandb.log({"bc/loss": avg_loss, "bc/epoch": epoch + 1})

        return losses

    def ppo_update(
        self, advantages: torch.Tensor, returns: torch.Tensor
    ) -> Dict[str, float]:
        """Single PPO update epoch across the current buffer."""
        buffer_data = self.buffer.get()
        obs = buffer_data["obs"]
        actions = buffer_data["actions"]
        old_log_probs = buffer_data["log_probs"]

        total_pg_loss = 0.0
        total_v_loss = 0.0
        total_entropy = 0.0
        n_updates = 0

        self.agent.train()
        indices = torch.randperm(len(obs), device=self.device)

        for start in range(0, len(obs), self.batch_size):
            idx = indices[start: start + self.batch_size]
            batch_obs = obs[idx]
            batch_actions = actions[idx]
            batch_old_lp = old_log_probs[idx]
            batch_adv = advantages[idx]
            batch_returns = returns[idx]

            # Normalize advantages within mini-batch
            batch_adv = (batch_adv - batch_adv.mean()) / (batch_adv.std() + 1e-8)

            log_probs, entropy, values = self.agent.evaluate_actions(
                batch_obs, batch_actions
            )

            ratio = torch.exp(log_probs - batch_old_lp)
            pg_loss1 = -batch_adv * ratio
            pg_loss2 = -batch_adv * ratio.clamp(
                1.0 - self.clip_eps, 1.0 + self.clip_eps
            )
            pg_loss = torch.max(pg_loss1, pg_loss2).mean()
            v_loss = 0.5 * (values - batch_returns).pow(2).mean()
            entropy_loss = -entropy.mean()

            loss = (
                pg_loss
                + self.value_coef * v_loss
                + self.entropy_coef * entropy_loss
            )

            self.optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(self.agent.parameters(), self.max_grad_norm)
            self.optimizer.step()

            total_pg_loss += pg_loss.item()
            total_v_loss += v_loss.item()
            total_entropy += entropy.mean().item()
            n_updates += 1

        n = max(n_updates, 1)
        return {
            "pg_loss": total_pg_loss / n,
            "v_loss": total_v_loss / n,
            "entropy": total_entropy / n,
        }

    def collect_rollout(self, env, n_steps: int = 2048):
        """
        Collect environment rollouts and fill buffer.

        Compatible with Gymnasium API:
          env.reset() → (obs, info)
          env.step(action) → (obs, reward, terminated, truncated, info)
        """
        self.buffer.reset()
        self.agent.eval()

        # Gymnasium reset returns (obs, info)
        obs_np, _ = env.reset()
        obs = torch.FloatTensor(obs_np).to(self.device)

        for _ in range(n_steps):
            with torch.no_grad():
                action, log_prob, _, value = self.agent.get_action(obs.unsqueeze(0))
                action = action.squeeze(0)
                log_prob = log_prob.squeeze(0)
                value = value.squeeze(0)

            # Gymnasium step returns 5-tuple
            next_obs_np, reward, terminated, truncated, _ = env.step(
                int(action.cpu().item())
            )
            done = terminated or truncated

            self.buffer.add(
                obs.cpu(), action.cpu(), log_prob.cpu(), reward, value.cpu(), done
            )

            if done:
                obs_np, _ = env.reset()
                obs = torch.FloatTensor(obs_np).to(self.device)
            else:
                obs = torch.FloatTensor(next_obs_np).to(self.device)

    def train(
        self,
        env,
        total_timesteps: int = 100_000,
        n_rollout_steps: int = 2048,
        checkpoint_dir: str = "checkpoints/intervention",
        save_interval: int = 10_000,
    ):
        """Full PPO training loop."""
        os.makedirs(checkpoint_dir, exist_ok=True)
        timestep = 0
        update = 0

        while timestep < total_timesteps:
            self.collect_rollout(env, n_rollout_steps)
            advantages, returns = self.buffer.compute_returns_and_advantages(
                self.gamma, self.gae_lambda
            )

            metrics = {}
            for _ in range(self.n_epochs):
                metrics = self.ppo_update(advantages, returns)

            timestep += n_rollout_steps
            update += 1

            print(
                f"[PPO] Update {update} | Steps {timestep}/{total_timesteps} | "
                f"PG: {metrics.get('pg_loss', 0):.4f}, "
                f"V: {metrics.get('v_loss', 0):.4f}, "
                f"Ent: {metrics.get('entropy', 0):.4f}"
            )

            if self.use_wandb:
                wandb.log({
                    "ppo/pg_loss": metrics.get("pg_loss", 0),
                    "ppo/v_loss": metrics.get("v_loss", 0),
                    "ppo/entropy": metrics.get("entropy", 0),
                    "ppo/timestep": timestep,
                    "ppo/update": update,
                })

            if timestep % save_interval == 0:
                ckpt_path = os.path.join(checkpoint_dir, f"agent_step{timestep}.pt")
                torch.save(self.agent.state_dict(), ckpt_path)
                print(f"Saved checkpoint: {ckpt_path}")

        final_path = os.path.join(checkpoint_dir, "final.pt")
        torch.save(self.agent.state_dict(), final_path)
        print(f"Training complete. Final model: {final_path}")
