from enum import Enum
from typing import Any, Dict, Type, Union

from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.logger import HParam

from freqtrade.freqai.RL.BaseEnvironment import BaseActions, BaseEnvironment


class TensorboardCallback(BaseCallback):
    """
    Custom callback for plotting additional values in tensorboard and
    episodic summary reports.
    """
    def __init__(self, verbose=1, actions: Type[Enum] = BaseActions):
        super(TensorboardCallback, self).__init__(verbose)
        self.model: Any = None
        self.logger = None  # type: Any
        self.training_env: BaseEnvironment = None  # type: ignore
        self.actions: Type[Enum] = actions

    def _on_training_start(self) -> None:
        hparam_dict = {
            "algorithm": self.model.__class__.__name__,
            "learning_rate": self.model.learning_rate,
            # "gamma": self.model.gamma,
            # "gae_lambda": self.model.gae_lambda,
            # "batch_size": self.model.batch_size,
            # "n_steps": self.model.n_steps,
        }
        metric_dict: Dict[str, Union[float, int]] = {
            "eval/mean_reward": 0,
            "rollout/ep_rew_mean": 0,
            "rollout/ep_len_mean": 0,
            "train/value_loss": 0,
            "train/explained_variance": 0,
        }
        self.logger.record(
            "hparams",
            HParam(hparam_dict, metric_dict),
            exclude=("stdout", "log", "json", "csv"),
        )

    def _on_step(self) -> bool:

        local_info = self.locals["infos"][0]
        custom_info = self.training_env.get_attr("custom_info")[0]

        for info in local_info:
            if info not in ["episode", "terminal_observation"]:
                self.logger.record(f"_info/{info}", local_info[info])

        for info in custom_info:
            if info in [action.name for action in self.actions]:
                self.logger.record(f"_actions/{info}", custom_info[info])
            else:
                self.logger.record(f"_custom/{info}", custom_info[info])

        return True