#
# Copyright (c) 2017 Intel 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 typing import List, Tuple

from rl_coach.base_parameters import Frameworks, AgentParameters
from rl_coach.logger import failed_imports
from rl_coach.saver import SaverCollection
from rl_coach.spaces import SpacesDefinition
from rl_coach.utils import force_list


class NetworkWrapper(object):
    """
    The network wrapper contains multiple copies of the same network, each one with a different set of weights which is
    updating in a different time scale. The network wrapper will always contain an online network.
    It will contain an additional slow updating target network if it was requested by the user,
    and it will contain a global network shared between different workers, if Coach is run in a single-node
    multi-process distributed mode. The network wrapper contains functionality for managing these networks and syncing
    between them.
    """
    def __init__(self, agent_parameters: AgentParameters, has_target: bool, has_global: bool, name: str,
                 spaces: SpacesDefinition, replicated_device=None, worker_device=None):
        self.ap = agent_parameters
        self.network_parameters = self.ap.network_wrappers[name]
        self.has_target = has_target
        self.has_global = has_global
        self.name = name
        self.sess = None

        if self.network_parameters.framework == Frameworks.tensorflow:
            try:
                import tensorflow as tf
            except ImportError:
                raise Exception('Install tensorflow before using it as framework')
            from rl_coach.architectures.tensorflow_components.general_network import GeneralTensorFlowNetwork
            general_network = GeneralTensorFlowNetwork.construct
        elif self.network_parameters.framework == Frameworks.mxnet:
            try:
                import mxnet as mx
            except ImportError:
                raise Exception('Install mxnet before using it as framework')
            from rl_coach.architectures.mxnet_components.general_network import GeneralMxnetNetwork
            general_network = GeneralMxnetNetwork.construct
        else:
            raise Exception("{} Framework is not supported"
                            .format(Frameworks().to_string(self.network_parameters.framework)))

        variable_scope = "{}/{}".format(self.ap.full_name_id, name)

        # Global network - the main network shared between threads
        self.global_network = None
        if self.has_global:
            # we assign the parameters of this network on the parameters server
            self.global_network = general_network(variable_scope=variable_scope,
                                                  devices=force_list(replicated_device),
                                                  agent_parameters=agent_parameters,
                                                  name='{}/global'.format(name),
                                                  global_network=None,
                                                  network_is_local=False,
                                                  spaces=spaces,
                                                  network_is_trainable=True)

        # Online network - local copy of the main network used for playing
        self.online_network = None
        self.online_network = general_network(variable_scope=variable_scope,
                                              devices=force_list(worker_device),