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

# %% auto 0
__all__ = ['SACRNN', 'SACRNN_hybrid']

# %% ../../../../nbs/agents/rl_agents/pre_specified_agents/12_SAC_RNN.ipynb 5
# 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 CriticNetworkStateActionRNN
from ...networks.actors import ActorNetworkRNN_MLP_hybrid, ActorNetworkRNN 
from ...processors.processors import GetTimeSeriesAndStaticFeatures

# Algorithms
from mushroom_rl.algorithms.actor_critic.deep_actor_critic import SAC

# Processors
from ...processors.processors import HybridToContinuous


# %% ../../../../nbs/agents/rl_agents/pre_specified_agents/12_SAC_RNN.ipynb 6
class SACRNN():

    """
    # TODO Update docstring
    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.
        warmup_transitions (int): Number of transitions to replay before starting to update 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,
            num_time_series_features,
            num_non_time_series_features,
            window_size,
            learning_rate_actor = 3e-4,
            learning_rate_critic = None,
            initial_replay_size = 64,
            max_replay_size = 50000,
            batch_size = 64,
            n_features = 64,
            warmup_transitions = 100,
            lr_alpha = 3e-4,
            tau = 0.005,
            log_std_min = -20,
            log_std_max = 2,
            target_entropy = None,
            RNN_unit = "GRU", # TODO: Make RNN and LSTM possible
            optimizer = optim.Adam,
            squeeze_output = True,
            use_cuda = True,
            agent_name = None):
        
        use_cuda = use_cuda and torch.cuda.is_available()

        input_shape_data = mdp_info.observation_space.shape
        input_shape_RNN = (num_time_series_features, window_size)
        input_shape_MLP = (num_non_time_series_features,)

        actor_output_shape = (mdp_info.action_space.shape[0],)

        if learning_rate_critic is None:
            learning_rate_critic = learning_rate_actor


        actor_mu_params = dict(network=ActorNetworkRNN,
                                n_features=n_features,
                                input_shape=input_shape_data,
                                input_shape_RNN=input_shape_RNN,
                                input_shape_MLP=input_shape_MLP,
                                output_shape=actor_output_shape,
                                input_processor=GetTimeSeriesAndStaticFeatures(num_time_series_features,window_size),
                                use_cuda=use_cuda)

        actor_sigma_params = dict(network=ActorNetworkRNN,
                                    n_features=n_features,
                                    input_shape= input_shape_data,
                                    input_shape_RNN=input_shape_RNN,
                                    input_shape_MLP=input_shape_MLP,
                                    output_shape=actor_output_shape,
                                    input_processor=GetTimeSeriesAndStaticFeatures(num_time_series_features,window_size),
                                    use_cuda=use_cuda)
        
        actor_optimizer = {'class': optimizer,
                    'params': {'lr': learning_rate_actor}} 
        
        critic_input_shape = (input_shape_data[0] + actor_output_shape[0],)
        input_shape_RNN = (num_time_series_features, window_size)
        input_shape_MLP = (num_non_time_series_features + actor_output_shape[0],)

        critic_params = dict(network=CriticNetworkStateActionRNN,
                        optimizer={'class': optim.Adam,
                                'params': {'lr': learning_rate_critic}}, 
                        loss=F.mse_loss,
                        n_features=n_features,
                        input_shape=critic_input_shape,
                        input_shape_RNN = input_shape_RNN,
                        input_shape_MLP = input_shape_MLP,
                        output_shape=(1,),
                        input_processor=GetTimeSeriesAndStaticFeatures(num_time_series_features,window_size),
                        squeeze_output=squeeze_output,
                        use_cuda=use_cuda)
        
        self.agent = SAC(mdp_info, actor_mu_params, actor_sigma_params,
                    actor_optimizer, critic_params, batch_size, initial_replay_size,
                    max_replay_size, warmup_transitions, tau, lr_alpha, log_std_min, log_std_max, target_entropy,
                    critic_fit_params=None)
        
        if agent_name is None:
            self.agent.name = 'SAC_RNN'
        else:
            self.agent.name = agent_name

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

# %% ../../../../nbs/agents/rl_agents/pre_specified_agents/12_SAC_RNN.ipynb 7
class SACRNN_hybrid():

    """
    # TODO Update docstring
    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.
        warmup_transitions (int): Number of transitions to replay before starting to update 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,
            num_time_series_features,
            num_non_time_series_features,
            window_size,
            learning_rate_actor = 3e-4,
            learning_rate_critic = None,
            initial_replay_size = 64,
            max_replay_size = 50000,
            batch_size = 64,
            n_features = 64,
            warmup_transitions = 100,
            lr_alpha = 3e-4,
            tau = 0.005,
            log_std_min = -20,
            log_std_max = 2,
            target_entropy = None,
            RNN_unit = "GRU", # TODO: Make RNN and LSTM possible
            optimizer = optim.Adam,
            squeeze_output = True,
            use_cuda = True,
            agent_name = None):
        
        use_cuda = use_cuda and torch.cuda.is_available()

        input_shape_data = mdp_info.observation_space.shape
        input_shape_RNN = (num_time_series_features, window_size)
        input_shape_MLP = (num_non_time_series_features,)

        actor_output_shape = (mdp_info.action_space.shape[0],)

        if learning_rate_critic is None:
            learning_rate_critic = learning_rate_actor


        actor_mu_params = dict(network=ActorNetworkRNN_MLP_hybrid,
                                n_features=n_features,
                                input_shape=input_shape_data,
                                input_shape_RNN=input_shape_RNN,
                                input_shape_MLP=input_shape_MLP,
                                output_shape=actor_output_shape,
                                input_processor=GetTimeSeriesAndStaticFeatures(num_time_series_features,window_size),
                                use_cuda=use_cuda)

        actor_sigma_params = dict(network=ActorNetworkRNN_MLP_hybrid,
                                    n_features=n_features,
                                    input_shape= input_shape_data,
                                    input_shape_RNN=input_shape_RNN,
                                    input_shape_MLP=input_shape_MLP,
                                    output_shape=actor_output_shape,
                                    input_processor=GetTimeSeriesAndStaticFeatures(num_time_series_features,window_size),
                                    use_cuda=use_cuda)
        
        actor_optimizer = {'class': optimizer,
                    'params': {'lr': learning_rate_actor}} 
        
        critic_input_shape = (input_shape_data[0] + actor_output_shape[0],)
        input_shape_RNN = (num_time_series_features, window_size)
        print(num_non_time_series_features)
        print(actor_output_shape[0])
        print("XXXXX")
        input_shape_MLP = (num_non_time_series_features + actor_output_shape[0],)

        critic_params = dict(network=CriticNetworkStateActionRNN,
                        optimizer={'class': optim.Adam,
                                'params': {'lr': learning_rate_critic}}, 
                        loss=F.mse_loss,
                        n_features=n_features,
                        input_shape=critic_input_shape,
                        input_shape_RNN = input_shape_RNN,
                        input_shape_MLP = input_shape_MLP,
                        output_shape=(1,),
                        input_processor=GetTimeSeriesAndStaticFeatures(num_time_series_features,window_size),
                        squeeze_output=squeeze_output,
                        use_cuda=use_cuda)
        
        self.agent = SAC(mdp_info, actor_mu_params, actor_sigma_params,
                    actor_optimizer, critic_params, batch_size, initial_replay_size,
                    max_replay_size, warmup_transitions, tau, lr_alpha, log_std_min, log_std_max, target_entropy,
                    critic_fit_params=None)
        
        if agent_name is None:
            self.agent.name = 'SAC_RNN_hybrid'
        else:
            self.agent.name = agent_name

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