"""TD3 based agent"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/30_agents/51_RL_agents/10_TD3_agents.ipynb.

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

# %% ../../../nbs/30_agents/51_RL_agents/10_TD3_agents.ipynb 3
import logging

# set logging level to INFO
logging.basicConfig(level=logging.INFO)

from abc import ABC, abstractmethod
from typing import Union, Optional, List, Tuple
import numpy as np
import os

from ...envs.base import BaseEnvironment
from .mushroom_rl import MushroomBaseAgent
from ...utils import MDPInfo, Parameter
from ..obsprocessors import FlattenTimeDimNumpy
from ...RL_approximators import MLPStateAction, MLPActor
from ...envs.actionprocessors import ClipAction

from ...dataloaders.base import BaseDataLoader

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

import torch
import torch.optim as optim
import torch.nn.functional as F
from torchinfo import summary

import time

# %% ../../../nbs/30_agents/51_RL_agents/10_TD3_agents.ipynb 4
class TD3Agent(MushroomBaseAgent):

    """
    XXX
    """

    # TODO: Make structure same to SAC with TD3 base class

    def __init__(self, 
                environment_info: MDPInfo,

                learning_rate_actor: float = 3e-4,
                learning_rate_critic: float | None = None, # If none, then it is set to learning_rate_actor
                initial_replay_size: int = 1024,
                max_replay_size: int = 50000,
                batch_size: int = 64,
                hidden_layers: List = None, # if None, then default is [64, 64]
                activation: str = "relu", # "relu", "sigmoid", "tanh", "leakyrelu", "elu"
                tau: float = 0.005,
                policy_delay: int = 2,
                noise_std: float = 0.2,
                sigma_scale: float = 0.5,
                theta: float = 0.15,
                dt=0.02,

                drop_prob: float = 0.0,
                batch_norm: bool = False,
                init_method: str = "xavier_uniform", # "xavier_uniform", "xavier_normal", "he_normal", "he_uniform", "normal", "uniform"

                optimizer: str = "Adam", # "Adam" or "SGD" or "RMSprop"  
                loss: str = "MSE", # currently only MSE is supported     
                obsprocessors: list | None = None,      # default: []
                device: str = "cpu", # "cuda" or "cpu"
                agent_name: str | None = "SAC",
                ):

        # The standard TD3 agent needs a 2D input, so we need to flatten the time dimension
        flatten_time_dim_processor = FlattenTimeDimNumpy(allow_2d=True, batch_dim_included=False)
        obsprocessors = (obsprocessors or []) + [flatten_time_dim_processor]

        # TODO: Add option to also use Gaussian policy
        self.policy_class = OrnsteinUhlenbeckPolicy
        self.policy_params = dict(sigma=np.ones(1) * sigma_scale, theta=theta, dt=dt)

        use_cuda = self.set_device(device)

        hidden_layers = hidden_layers or [64, 64]
        self.warmup_training_steps = initial_replay_size


        OptimizerClass=self.get_optimizer_class(optimizer)
        learning_rate_critic = learning_rate_critic or learning_rate_actor
        lossfunction = self.get_loss_function(loss)

        input_shape = self.get_input_shape(environment_info.observation_space)
        output_shape = environment_info.action_space.shape

        input_shape = self.convert_recursively_to_int(input_shape)
        output_shape = self.convert_recursively_to_int(output_shape)

        actor_input_shape = input_shape
        actor_output_shape = output_shape
        critic_input_shape = (actor_input_shape[0] + actor_output_shape[0],) # check how this works when RNN and mixed agents are used

        actor_params = dict(network=MLPActor,

                                input_shape=actor_input_shape,
                                output_shape=actor_output_shape,

                                hidden_layers=hidden_layers,
                                activation=activation,
                                drop_prob=drop_prob,
                                batch_norm=batch_norm,
                                init_method=init_method,

                                use_cuda=use_cuda,
                                dropout=self.dropout,

                                )

        actor_optimizer = {'class': OptimizerClass,
            'params': {'lr': learning_rate_actor}} 

        critic_params = dict(network=MLPStateAction,
                optimizer={'class': OptimizerClass,
                        'params': {'lr': learning_rate_critic}}, 
                loss=lossfunction,
                input_shape=critic_input_shape,
                output_shape=(1,),

                hidden_layers=hidden_layers,
                activation=activation,
                drop_prob=drop_prob,
                batch_norm=batch_norm,
                init_method=init_method,

                use_cuda=use_cuda,
                dropout=self.dropout,)

        self.agent = TD3(
            mdp_info=environment_info,
            policy_class=self.policy_class,
            policy_params=self.policy_params,
            actor_params=actor_params,
            actor_optimizer=actor_optimizer,
            critic_params=critic_params,
            batch_size=batch_size,
            initial_replay_size=initial_replay_size,
            max_replay_size=max_replay_size,
            tau=tau,
            policy_delay=policy_delay,
            noise_std=noise_std,
            critic_fit_params=None
        )

        super().__init__(
            environment_info=environment_info,
            obsprocessors=obsprocessors,
            device=device,
            agent_name=agent_name
        )

        logging.info("Actor network:")
        if logging.getLogger().isEnabledFor(logging.INFO):
            input_size = self.add_batch_dimension_for_shape(actor_input_shape)
            print(summary(self.actor, input_size=input_size))
            time.sleep(.2)

        logging.info("Critic network:")
        if logging.getLogger().isEnabledFor(logging.INFO):
            input_size = self.add_batch_dimension_for_shape([actor_input_shape, actor_output_shape])
            print(summary(self.critic, input_size=input_size))

    def get_network_list(self, set_actor_critic_attributes: bool = True):
        """ Get the list of networks in the agent for the save and load functions
        Get the actor for the predict function in eval mode """

        networks = []
        ensemble_critic = self.agent._critic_approximator._impl.model
        for i, model in enumerate(ensemble_critic):
            networks.append(model.network)
        networks.append(self.agent.policy._approximator._impl.model.network)

        actor = self.agent.policy._approximator._impl.model.network
        critic = ensemble_critic[0].network

        if set_actor_critic_attributes:
            return networks, actor, critic
        else:
            return networks
