# python3
# Copyright 2021 InstaDeep Ltd. All rights reserved.
#
# 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.

"""MAPPO system implementation."""

import functools
from typing import Any, Callable, Dict, Optional, Type, Union

import acme
import dm_env
import launchpad as lp
import reverb
import sonnet as snt
from acme import specs as acme_specs
from acme.utils import counting, loggers

import mava
from mava import core
from mava import specs as mava_specs
from mava.components.tf.architectures import DecentralisedValueActorCritic
from mava.environment_loop import ParallelEnvironmentLoop
from mava.systems.tf import savers as tf2_savers
from mava.systems.tf.mappo import builder, execution, training
from mava.utils import lp_utils
from mava.utils.loggers import MavaLogger, logger_utils
from mava.wrappers import DetailedPerAgentStatistics


class MAPPO:
    """MAPPO system."""

    def __init__(
        self,
        environment_factory: Callable[[bool], dm_env.Environment],
        network_factory: Callable[[acme_specs.BoundedArray], Dict[str, snt.Module]],
        logger_factory: Callable[[str], MavaLogger] = None,
        architecture: Type[
            DecentralisedValueActorCritic
        ] = DecentralisedValueActorCritic,
        trainer_fn: Type[training.MAPPOTrainer] = training.MAPPOTrainer,
        executor_fn: Type[core.Executor] = execution.MAPPOFeedForwardExecutor,
        num_executors: int = 1,
        num_caches: int = 0,
        environment_spec: mava_specs.MAEnvironmentSpec = None,
        shared_weights: bool = True,
        agent_net_keys: Dict[str, str] = {},
        executor_variable_update_period: int = 100,
        optimizer: Union[snt.Optimizer, Dict[str, snt.Optimizer]] = snt.optimizers.Adam(
            learning_rate=5e-4
        ),
        discount: float = 0.99,
        lambda_gae: float = 0.95,
        clipping_epsilon: float = 0.2,
        entropy_cost: float = 0.01,
        baseline_cost: float = 1.0,
        max_gradient_norm: Optional[float] = 0.5,
        max_queue_size: Optional[int] = None,
        batch_size: int = 512,
        minibatch_size: int = None,
        num_epochs: int = 10,
        sequence_length: int = 10,
        sequence_period: Optional[int] = None,
        max_executor_steps: int = None,
        checkpoint: bool = True,
        checkpoint_subpath: str = "~/mava/",
        checkpoint_minute_interval: int = 5,
        logger_config: Dict = {},
        train_loop_fn: Callable = ParallelEnvironmentLoop,
        eval_loop_fn: Callable = ParallelEnvironmentLoop,
        train_loop_fn_kwargs: Dict = {},
        eval_loop_fn_kwargs: Dict = {},
        evaluator_interval: Optional[dict] = None,
        learning_rate_scheduler_fn: Optional[Callable[[int], None]] = None,
        normalize_advantage: bool = False,
    ):
        """Initialise the system

        Args:
            environment_factory: function to
                instantiate an environment.
            network_factory : function to instantiate system networks.
            logger_factory : function to
                instantiate a system logger. Defaults to None.
            architecture : system
                architecture, e.g. decentralised or centralised. Defaults to
                DecentralisedValueActorCritic.
            trainer_fn : training type
                associated with executor and architecture, e.g. centralised training.
                Defaults to training.MAPPOTrainer.
            executor_fn : executor type, e.g. feedforward
                or recurrent. Defaults to execution.MAPPOFeedForwardExecutor.
            num_executors : number of executor processes to run in
                parallel. Defaults to 1.
            num_caches : number of trainer node caches. Defaults to 0.
            environment_spec : description of
                the action, observation spaces etc. for each agent in the system.
                Defaults to None.
            shared_weights : whether agents should share weights or not.
                When agent_net_keys are provided the value of shared_weights is ignored.
                Defaults to True.
            agent_net_keys: : specifies what network each agent uses.
                Defaults to {}.
            executor_variable_update_period : number of steps before
                updating executor variables from the variable source. Defaults to 100.
            optimizer : optimizer(s) for updating networks.
                Defaults to snt.optimizers.Adam(learning_rate=5e-4).
            discount : discount factor to use for TD updates. Defaults
                to 0.99.
            lambda_gae : scalar determining the mix of bootstrapping
                vs further accumulation of multi-step returns at each timestep.
                Defaults to 0.99.
            clipping_epsilon : Hyper-parameter for clipping in the
                policy objective. Defaults to 0.2.
            entropy_cost : contribution of entropy regularization to
                the total loss. Defaults to 0.01.
            baseline_cost : contribution of the value loss to the
                total loss. Defaults to 0.5.
            max_gradient_norm: value to specify the maximum clipping value for the
            gradient norm during optimization.
            max_queue_size : maximum number of items in the queue.
                Should be larger than batch size.
            batch_size: sample batch size for updates.
                Defaults to 512. Minibatches are sampled from this data.
            minibatch_size: size of minibatch that is sampled
                from the training batch. Minibatches are used for each gradient step.
            num_epochs: number of epochs every training step.
                Recommendation as per https://arxiv.org/pdf/2103.01955.pdf, "15
                epochs for easy tasks,and 10 or 5 epochs for difficult tasks."
            sequence_length: sequence rollout length. Defaults
                to 10.
            sequence_period: consecutive starting points for
                overlapping rollouts across a sequence. Defaults to sequence length -1.
            max_executor_steps: maximum number of steps and executor
                can in an episode. Defaults to None.
            checkpoint : whether to checkpoint models. Defaults to
                False.
            checkpoint_subpath: subdirectory specifying where to store
                checkpoints. Defaults to "~/mava/".
            checkpoint_minute_interval: The number of minutes to wait between
                checkpoints.
            logger_config: additional configuration settings for the
                logger factory. Defaults to {}.
            train_loop_fn: function to instantiate a train loop.
                Defaults to ParallelEnvironmentLoop.
            eval_loop_fn: function to instantiate an evaluation
                loop. Defaults to ParallelEnvironmentLoop.
            train_loop_fn_kwargs: possible keyword arguments to send
                to the training loop. Defaults to {}.
            eval_loop_fn_kwargs: possible keyword arguments to send to
                the evaluation loop. Defaults to {}.
            learning_rate_scheduler_fn: an optional learning rate scheduler for
                the optimiser.
            evaluator_interval: An optional condition that is used to
                evaluate/test system performance after [evaluator_interval]
                condition has been met. If None, evaluation will
                happen at every timestep.
                E.g. to evaluate a system after every 100 executor episodes,
                evaluator_interval = {"executor_episodes": 100}.
            normalize_advantage: whether to normalize the advantage estimate. This can
                hurt peformance when shared weights are used.
        """
        # minibatch size defaults to train batch size
        if minibatch_size:
            self._minibatch_size = minibatch_size
        else:
            self._minibatch_size = batch_size

        assert batch_size % self._minibatch_size == 0, (
            "batch_size must be divisible by minibatch_size."
            + f"Got batch_size={batch_size},"
            + f"minibatch_size={self._minibatch_size}"
        )

        if max_queue_size:
            self._max_queue_size = max_queue_size
        else:
            # NOTE: Based on https://github.com/deepmind/acme/blob/6bf350df1d9dd16cd85217908ec9f47553278976/acme/agents/jax/ppo/builder.py#L75 # noqa: E501
            extra_capacity_to_avoid_single_machine_deadlocks = 1000
            self._max_queue_size = (
                batch_size + extra_capacity_to_avoid_single_machine_deadlocks
            )

        assert self._max_queue_size > batch_size, (
            "Max queue size should be larger than batch size."
            + f"Got max_queue_size={self._max_queue_size},"
            + f"batch_size={batch_size}"
        )

        # An extra step is used for bootstrapping when computing advantages.
        self._sequence_length = sequence_length + 1

        if sequence_period:
            self._sequence_period = sequence_period
        else:
            self._sequence_period = self._sequence_length - 1

        if not environment_spec:
            environment_spec = mava_specs.MAEnvironmentSpec(
                environment_factory(evaluation=False)  # type: ignore
            )

        # set default logger if no logger provided
        if not logger_factory:
            logger_factory = functools.partial(
                logger_utils.make_logger,
                directory="~/mava",
                to_terminal=True,
                time_delta=10,
            )

        self._architecture = architecture
        self._environment_factory = environment_factory
        self._network_factory = network_factory
        self._logger_factory = logger_factory
        self._environment_spec = environment_spec
        # Setup agent networks
        self._agent_net_keys = agent_net_keys
        if not agent_net_keys:
            agents = environment_spec.get_agent_ids()
            self._agent_net_keys = {
                agent: agent.split("_")[0] if shared_weights else agent
                for agent in agents
            }
        self._num_exectors = num_executors
        self._num_caches = num_caches
        self._max_executor_steps = max_executor_steps
        self._checkpoint_subpath = checkpoint_subpath
        self._checkpoint = checkpoint
        self._logger_config = logger_config
        self._train_loop_fn = train_loop_fn
        self._train_loop_fn_kwargs = train_loop_fn_kwargs
        self._eval_loop_fn = eval_loop_fn
        self._eval_loop_fn_kwargs = eval_loop_fn_kwargs
        self._checkpoint_minute_interval = checkpoint_minute_interval
        self._evaluator_interval = evaluator_interval

        self._builder = builder.MAPPOBuilder(
            config=builder.MAPPOConfig(
                environment_spec=environment_spec,
                agent_net_keys=self._agent_net_keys,
                executor_variable_update_period=executor_variable_update_period,
                discount=discount,
                lambda_gae=lambda_gae,
                clipping_epsilon=clipping_epsilon,
                entropy_cost=entropy_cost,
                baseline_cost=baseline_cost,
                max_gradient_norm=max_gradient_norm,
                max_queue_size=self._max_queue_size,
                batch_size=batch_size,
                minibatch_size=self._minibatch_size,
                num_epochs=num_epochs,
                sequence_length=self._sequence_length,
                sequence_period=self._sequence_period,
                checkpoint=checkpoint,
                optimizer=optimizer,
                checkpoint_subpath=checkpoint_subpath,
                checkpoint_minute_interval=checkpoint_minute_interval,
                evaluator_interval=evaluator_interval,
                learning_rate_scheduler_fn=learning_rate_scheduler_fn,
                normalize_advantage=normalize_advantage,
            ),
            trainer_fn=trainer_fn,
            executor_fn=executor_fn,
        )

    def replay(self) -> Any:
        """Replay data storage.

        Returns:
            Any: replay data table built according the environment specification.
        """

        return self._builder.make_replay_tables(self._environment_spec)

    def counter(self, checkpoint: bool) -> Any:
        """Step counter

        Args:
            checkpoint : whether to checkpoint the counter.

        Returns:
            Any: step counter object.
        """

        if checkpoint:
            return tf2_savers.CheckpointingRunner(
                counting.Counter(),
                time_delta_minutes=self._checkpoint_minute_interval,
                directory=self._checkpoint_subpath,
                subdirectory="counter",
            )
        else:
            return counting.Counter()

    def coordinator(self, counter: counting.Counter) -> Any:
        """Coordination helper for a distributed program

        Args:
            counter: step counter object.

        Returns:
            Any: step limiter object.
        """

        return lp_utils.StepsLimiter(counter, self._max_executor_steps)  # type: ignore

    def trainer(
        self,
        replay: reverb.Client,
        counter: counting.Counter,
    ) -> mava.core.Trainer:
        """System trainer

        Args:
            replay: replay data table to pull data from.
            counter: step counter object.

        Returns:
            mava.core.Trainer: system trainer.
        """

        # Create the networks to optimize (online)
        networks = self._network_factory(  # type: ignore
            environment_spec=self._environment_spec,
            agent_net_keys=self._agent_net_keys,
        )

        # Create system architecture with target networks.
        system_networks = self._architecture(
            environment_spec=self._environment_spec,
            observation_networks=networks["observations"],
            policy_networks=networks["policies"],
            critic_networks=networks["critics"],
            agent_net_keys=self._agent_net_keys,
        ).create_system()

        # create logger
        trainer_logger_config = {}
        if self._logger_config:
            if "trainer" in self._logger_config:
                trainer_logger_config = self._logger_config["trainer"]
        trainer_logger = self._logger_factory(  # type: ignore
            "trainer", **trainer_logger_config
        )

        dataset = self._builder.make_dataset_iterator(replay)
        counter = counting.Counter(counter, "trainer")

        return self._builder.make_trainer(
            networks=system_networks,
            dataset=dataset,
            counter=counter,
            logger=trainer_logger,
        )

    def executor(
        self,
        executor_id: str,
        replay: reverb.Client,
        variable_source: acme.VariableSource,
        counter: counting.Counter,
    ) -> mava.ParallelEnvironmentLoop:
        """System executor

        Args:
            executor_id: id to identify the executor process for logging purposes.
            replay: replay data table to push data to.
            variable_source (acme.VariableSource): variable server for updating
                network variables.
            counter: step counter object.

        Returns:
            mava.ParallelEnvironmentLoop: environment-executor loop instance.
        """

        # Create the behavior policy.
        networks = self._network_factory(  # type: ignore
            environment_spec=self._environment_spec,
            agent_net_keys=self._agent_net_keys,
        )

        # Create system architecture with target networks.
        system = self._architecture(
            environment_spec=self._environment_spec,
            observation_networks=networks["observations"],
            policy_networks=networks["policies"],
            critic_networks=networks["critics"],
            agent_net_keys=self._agent_net_keys,
        )

        # create variables
        _ = system.create_system()

        # behaviour policy networks (obs net + policy head)
        behaviour_policy_networks = system.create_behaviour_policy()

        # Create the executor.
        executor = self._builder.make_executor(
            policy_networks=behaviour_policy_networks,
            adder=self._builder.make_adder(replay),
            variable_source=variable_source,
            evaluator=False,
        )

        # TODO (Arnu): figure out why factory function are giving type errors
        # Create the environment.
        environment = self._environment_factory(evaluation=False)  # type: ignore

        # Create logger and counter; actors will not spam bigtable.
        counter = counting.Counter(counter, "executor")

        # Create executor logger
        executor_logger_config = {}
        if self._logger_config:
            if "executor" in self._logger_config:
                executor_logger_config = self._logger_config["executor"]
        exec_logger = self._logger_factory(  # type: ignore
            f"executor_{executor_id}", **executor_logger_config
        )

        # Create the loop to connect environment and executor.
        train_loop = self._train_loop_fn(
            environment,
            executor,
            counter=counter,
            logger=exec_logger,
            **self._train_loop_fn_kwargs,
        )

        train_loop = DetailedPerAgentStatistics(train_loop)

        return train_loop

    def evaluator(
        self,
        variable_source: acme.VariableSource,
        counter: counting.Counter,
        logger: loggers.Logger = None,
    ) -> Any:
        """System evaluator - an executor process not connected to a dataset.

        Args:
            variable_source : variable server for updating network variables.
            counter : step counter object.
            logger : logger object. Defaults to None.

        Returns:
            environment-executor evaluation loop instance for evaluating the
                performance of a system.
        """

        # Create the behavior policy.
        networks = self._network_factory(  # type: ignore
            environment_spec=self._environment_spec,
            agent_net_keys=self._agent_net_keys,
        )

        # Create system architecture with target networks.
        system = self._architecture(
            environment_spec=self._environment_spec,
            observation_networks=networks["observations"],
            policy_networks=networks["policies"],
            critic_networks=networks["critics"],
            agent_net_keys=self._agent_net_keys,
        )

        # create variables
        _ = system.create_system()

        # behaviour policy networks (obs net + policy head)
        behaviour_policy_networks = system.create_behaviour_policy()

        # Create the agent.
        executor = self._builder.make_executor(
            policy_networks=behaviour_policy_networks,
            variable_source=variable_source,
            evaluator=True,
        )

        # Make the environment.
        environment = self._environment_factory(evaluation=True)  # type: ignore

        # Create logger and counter.
        counter = counting.Counter(counter, "evaluator")
        evaluator_logger_config = {}
        if self._logger_config:
            if "evaluator" in self._logger_config:
                evaluator_logger_config = self._logger_config["evaluator"]
        eval_logger = self._logger_factory(  # type: ignore
            "evaluator", **evaluator_logger_config
        )

        # Create the run loop and return it.
        # Create the loop to connect environment and executor.
        eval_loop = self._eval_loop_fn(
            environment,
            executor,
            counter=counter,
            logger=eval_logger,
            **self._eval_loop_fn_kwargs,
        )

        eval_loop = DetailedPerAgentStatistics(eval_loop)
        return eval_loop

    def build(self, name: str = "mappo") -> Any:
        """Build the distributed system as a graph program.

        Args:
            name : system name. Defaults to "mappo".

        Returns:
            Any: graph program for distributed system training.
        """
        program = lp.Program(name=name)

        with program.group("replay"):
            replay = program.add_node(lp.ReverbNode(self.replay))

        with program.group("counter"):
            counter = program.add_node(lp.CourierNode(self.counter, self._checkpoint))

        if self._max_executor_steps:
            with program.group("coordinator"):
                _ = program.add_node(lp.CourierNode(self.coordinator, counter))

        with program.group("trainer"):
            trainer = program.add_node(lp.CourierNode(self.trainer, replay, counter))

        with program.group("evaluator"):
            program.add_node(lp.CourierNode(self.evaluator, trainer, counter))

        if not self._num_caches:
            # Use the trainer as a single variable source.
            sources = [trainer]
        else:
            with program.group("cacher"):
                # Create a set of trainer caches.
                sources = []
                for _ in range(self._num_caches):
                    cacher = program.add_node(
                        lp.CacherNode(
                            trainer, refresh_interval_ms=2000, stale_after_ms=4000
                        )
                    )
                    sources.append(cacher)

        with program.group("executor"):
            # Add executors which pull round-robin from our variable sources.
            for executor_id in range(self._num_exectors):
                source = sources[executor_id % len(sources)]
                program.add_node(
                    lp.CourierNode(self.executor, executor_id, replay, source, counter)
                )

        return program
