# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass
from typing import Union

import gym

import nnabla_rl.model_trainers as MT
from nnabla_rl.algorithms.categorical_dqn import (CategoricalDQN, CategoricalDQNConfig, DefaultReplayBufferBuilder,
                                                  DefaultSolverBuilder, DefaultValueDistFunctionBuilder)
from nnabla_rl.builders import ModelBuilder, ReplayBufferBuilder, SolverBuilder
from nnabla_rl.environments.environment_info import EnvironmentInfo
from nnabla_rl.models import ValueDistributionFunction
from nnabla_rl.utils.misc import sync_model


@dataclass
class CategoricalDDQNConfig(CategoricalDQNConfig):
    pass


class CategoricalDDQN(CategoricalDQN):
    '''Categorical Double DQN algorithm.

    This class implements the Categorical Double DQN algorithm
    introduced by M. Bellemare, et al. in the paper: "A Distributional Perspective on Reinfocement Learning"
    For details see: https://arxiv.org/abs/1707.06887

    Args:
        env_or_env_info \
        (gym.Env or :py:class:`EnvironmentInfo <nnabla_rl.environments.environment_info.EnvironmentInfo>`):
            the environment to train or environment info
        config (:py:class:`CategoricalDDQNConfig <nnabla_rl.algorithms.categorical_ddqn.CategoricalDDQNConfig>`):
            configuration of the CategoricalDDQN algorithm
        value_distribution_builder (:py:class:`ModelBuilder[ValueDistributionFunctionFunction] \
            <nnabla_rl.builders.ModelBuilder>`): builder of value distribution function models
        value_distribution_solver_builder (:py:class:`SolverBuilder <nnabla_rl.builders.SolverBuilder>`):
            builder of value distribution function solvers
        replay_buffer_builder (:py:class:`ReplayBufferBuilder <nnabla_rl.builders.ReplayBufferBuilder>`):
            builder of replay_buffer
    '''

    def __init__(self, env_or_env_info: Union[gym.Env, EnvironmentInfo],
                 config: CategoricalDQNConfig = CategoricalDDQNConfig(),
                 value_distribution_builder: ModelBuilder[ValueDistributionFunction]
                 = DefaultValueDistFunctionBuilder(),
                 value_distribution_solver_builder: SolverBuilder = DefaultSolverBuilder(),
                 replay_buffer_builder: ReplayBufferBuilder = DefaultReplayBufferBuilder()):
        super(CategoricalDDQN, self).__init__(env_or_env_info,
                                              config=config,
                                              value_distribution_builder=value_distribution_builder,
                                              value_distribution_solver_builder=value_distribution_solver_builder,
                                              replay_buffer_builder=replay_buffer_builder)

    def _setup_value_distribution_function_training(self, env_or_buffer):
        trainer_config = MT.q_value_trainers.CategoricalDDQNQTrainerConfig(
            v_min=self._config.v_min,
            v_max=self._config.v_max,
            num_atoms=self._config.num_atoms)

        model_trainer = MT.q_value_trainers.CategoricalDDQNQTrainer(
            train_function=self._atom_p,
            solvers={self._atom_p.scope_name: self._atom_p_solver},
            target_function=self._target_atom_p,
            env_info=self._env_info,
            config=trainer_config)

        # NOTE: Copy initial parameters after setting up the training
        # Because the parameter is created after training graph construction
        sync_model(self._atom_p, self._target_atom_p)
        return model_trainer
