# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../../nbs/agents/rl_agents/pre_specified_agents/11_PPO_classic.ipynb.

# %% auto 0
__all__ = ['PPOClassic']

# %% ../../../../nbs/agents/rl_agents/pre_specified_agents/11_PPO_classic.ipynb 4
# General libraries
import numpy as np

# Torch
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# Networks
from ...networks.critics import CriticNetworkStandard
from ...networks.actors import ActorNetwork

# Algorithms
from mushroom_rl.algorithms.actor_critic import PPO
from mushroom_rl.policy import GaussianTorchPolicy

# Processors
from ...processors.processors import HybridToContinuous


# %% ../../../../nbs/agents/rl_agents/pre_specified_agents/11_PPO_classic.ipynb 5
class PPOClassic():

    """
    # TODO Update docstring
    """

    def __init__(
            self,
            mdp_info,
            learning_rate_actor = 3e-4,
            learning_rate_critic = None,
            batch_size = 64,
            n_features = 128,
            std_0 = 1.,
            n_epochs_policy = 4,
            eps_ppo = 0.2,
            lam=0.95,
            ent_coeff=0.0,
            n_steps_per_fit=1000,
            optimizer = optim.Adam,
            squeeze_output = True,
            use_cuda = True,
            agent_name = None): 

        self.n_steps_per_fit=n_steps_per_fit
        
        use_cuda = use_cuda and torch.cuda.is_available()
        if use_cuda:
            self.device = "cuda"
        else:
            self.device = "cpu"

        input_shape = mdp_info.observation_space.shape
        print("input shape: ", input_shape)
        actor_output_shape = (mdp_info.action_space.shape[0],) 

        if learning_rate_critic is None:
            learning_rate_critic = learning_rate_actor

        policy_params = dict(network=ActorNetwork, # equivalent to actor params in SAC
                                input_shape=input_shape,
                                output_shape=actor_output_shape,
                                use_cuda=use_cuda,
                                st_0=std_0,
                                n_features=n_features,
                                )

        policy = GaussianTorchPolicy(**policy_params)

        actor_optimizer = {'class': optimizer,
                    'params': {'lr': learning_rate_actor}} 
    
        critic_params = dict(network=CriticNetworkStandard,
                        optimizer={'class': optim.Adam,
                                'params': {'lr': learning_rate_critic}}, 
                        loss=F.mse_loss,
                        n_features=n_features,
                        batch_size = batch_size,
                        input_shape=input_shape,
                        output_shape=(1,),
                        squeeze_output=squeeze_output,
                        use_cuda=use_cuda)
        
        self.agent = PPO(mdp_info, policy, actor_optimizer, critic_params, n_epochs_policy, batch_size, eps_ppo, lam, ent_coeff)
        
        if agent_name is None:
            self.agent.name = 'PPO_classic'
        else:
            self.agent.name = agent_name

    def __getattr__(self, attr):
        return getattr(self.agent, attr)
