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

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

# %% ../../../../nbs/agents/rl_agents/pre_specified_agents/11_TD3_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 CriticNetworkStateAction
from ...networks.actors import ActorNetwork

# Algorithms
from mushroom_rl.algorithms.actor_critic.deep_actor_critic import TD3
from mushroom_rl.policy import OrnsteinUhlenbeckPolicy

# Processors
from ...processors.processors import HybridToContinuous


# %% ../../../../nbs/agents/rl_agents/pre_specified_agents/11_TD3_classic.ipynb 5
class TD3Classic():

    """
    Soft Actor Critic (SAC) agent with hybrid action, both based on Gaussian. The binary action is 
    0 if the output of the network is less or equal than 0, and 1 otherwise.

    Args:
        mdp_info (MDPInfo): Contains relevant information about the environment.
        learning_rate_actor (float): Learning rate for the actor.
        learning_rate_critic (float): Learning rate for the critic.
        learning_rate_alpha (float): Learning rate for the temperature parameter.
        initial_replay_size (int): Number of transitions to save in the replay buffer during startup.
        max_replay_size (int): Maximum number of transitions to save in the replay buffer.
        batch_size (int): Number of transitions to sample each time experience is replayed.
        n_features (int): Number of features for the hidden layers of the networks.
        lr_alpha (float): Learning rate for the temperature parameter.
        tau (float): Parameter for the soft update of the target networks.
        optimizer (torch.optim): Optimizer to use for the networks.
        squeeze_output (bool): Whether to squeeze the output of the actor network or not.
        use_cuda (bool): Whether to use CUDA or not. If True and not available, it will use CPU.
        agent_name (str): Name of the agent. If set to None will use some default name.

    """

    def __init__(
            self,
            mdp_info,
            learning_rate_actor = 3e-4,
            learning_rate_critic = None,
            initial_replay_size = 1024,
            max_replay_size = 50000,
            batch_size = 64,
            n_features = 64,
            tau = 0.005,
            policy_delay = 2,
            noise_std = 0.2,
            optimizer = optim.Adam,
            sigma_scale = 0.5,

            theta=0.15,
            dt=0.02,
            squeeze_output = True,
            use_cuda = False,
            agent_name = None): 
        
        # print("in init fubction")
        
        self.policy_class = OrnsteinUhlenbeckPolicy
        self.policy_params = dict(sigma=np.ones(1) * sigma_scale, theta=theta, dt=dt)
        
        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
        actor_output_shape = (mdp_info.action_space.shape[0],) 

        if learning_rate_critic is None:
            learning_rate_critic = learning_rate_actor

        actor_params = dict(network=ActorNetwork,
                                n_features=n_features,
                                input_shape=input_shape,
                                output_shape=actor_output_shape,
                                use_cuda=use_cuda)
        
        # print("setting optimizer class")
        actor_optimizer = {'class': optimizer,
                    'params': {'lr': learning_rate_actor}} 
        
        critic_input_shape = (input_shape[0] + actor_output_shape[0],)
        critic_params = dict(network=CriticNetworkStateAction,
                        optimizer={'class': optimizer,
                                'params': {'lr': learning_rate_critic}}, 
                        loss=F.mse_loss,
                        n_features=n_features,
                        input_shape=critic_input_shape,
                        output_shape=(1,),
                        squeeze_output=squeeze_output,
                        use_cuda=use_cuda)
        
        # print("creating agent from mushroom")
        
        self.agent = TD3(mdp_info, self.policy_class, self.policy_params,
                    actor_params, actor_optimizer, critic_params, batch_size,
                    initial_replay_size, max_replay_size, tau, policy_delay, noise_std)
    
        # print("created agent from mushroom")

        if agent_name is None:
            self.agent.name = 'TD3_classic'
        else:
            self.agent.name = agent_name

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