# Copyright 2018 Tensorforce Team. 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.
# ==============================================================================

from copy import deepcopy
from datetime import datetime
import os
import sys
import warnings

from tensorforce import TensorforceError
from tensorforce.agents import Agent
from tensorforce.core.layers import Layer
from tensorforce.environments import Environment
from tensorforce.execution import Runner
from test.unittest_environment import UnittestEnvironment


os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'


class UnittestBase(object):
    """
    Unit-test base class.
    """

    # Unittest
    num_updates = None
    num_episodes = None
    num_timesteps = None

    # Environment
    min_timesteps = 1
    states = dict(
        bool_state=dict(type='bool', shape=(1,)),
        int_state=dict(type='int', shape=(2,), num_values=4),
        float_state=dict(type='float', shape=(1, 1, 2)),
        bounded_state=dict(type='float', shape=(), min_value=-0.5, max_value=0.5)
    )
    actions = dict(
        bool_action=dict(type='bool', shape=(1,)),
        int_action=dict(type='int', shape=(2,), num_values=4),
        float_action=dict(type='float', shape=(1, 1)),
        bounded_action=dict(type='float', shape=(2,), min_value=-0.5, max_value=0.5)
    )

    # Exclude action types
    exclude_bool_action = False
    exclude_int_action = False
    exclude_float_action = False
    exclude_bounded_action = False

    # Agent
    agent = dict(
        update=4, policy=dict(network=dict(type='auto', size=8, depth=1, internal_rnn=2)),
        objective='policy_gradient', reward_estimation=dict(horizon=3)
    )

    # Tensorforce config
    require_observe = False
    require_all = False

    def setUp(self):
        warnings.filterwarnings(
            action='ignore',
            message='Converting sparse IndexedSlices to a dense Tensor of unknown shape'
        )

    def start_tests(self, name=None):
        """
        Start unit-test method.
        """
        if name is None:
            sys.stdout.write('\n{} {}: '.format(
                datetime.now().strftime('%H:%M:%S'), self.__class__.__name__[4:]
            ))
        else:
            sys.stdout.write('\n{} {} ({}): '.format(
                datetime.now().strftime('%H:%M:%S'), self.__class__.__name__[4:], name
            ))
        sys.stdout.flush()

    def finished_test(self, assertion=None):
        """
        Finished unit-test.
        """
        if assertion is None:
            assertion = True
        else:
            self.assertTrue(expr=assertion)
        if assertion:
            sys.stdout.write('.')
            sys.stdout.flush()

    def prepare(
        self, environment=None, min_timesteps=None, states=None, actions=None,
        exclude_bool_action=False, exclude_int_action=False, exclude_float_action=False,
        exclude_bounded_action=False, require_observe=False, require_all=False, **agent
    ):
        """
        Generic unit-test preparation.
        """
        Layer.layers = None

        if environment is None:
            if states is None:
                states = deepcopy(self.__class__.states)

            if actions is None:
                actions = deepcopy(self.__class__.actions)
                if exclude_bool_action or self.__class__.exclude_bool_action:
                    actions.pop('bool_action')
                if exclude_int_action or self.__class__.exclude_int_action:
                    actions.pop('int_action')
                if exclude_float_action or self.__class__.exclude_float_action:
                    actions.pop('float_action')
                if exclude_bounded_action or self.__class__.exclude_bounded_action:
                    actions.pop('bounded_action')

            if min_timesteps is None:
                min_timesteps = self.__class__.min_timesteps

            environment = UnittestEnvironment(
                states=states, actions=actions, min_timesteps=min_timesteps
            )

        elif min_timesteps is not None:
            raise TensorforceError.unexpected()

        environment = Environment.create(environment=environment, max_episode_timesteps=5)

        for key, value in self.__class__.agent.items():
            if key not in agent:
                agent[key] = value

        if self.__class__.require_all or require_all:
            config = None
        elif self.__class__.require_observe or require_observe:
            config = dict(api_functions=['reset', 'act', 'observe'])
        else:
            config = dict(api_functions=['reset', 'act'])

        agent = Agent.create(agent=agent, environment=environment, config=config)

        return agent, environment

    def unittest(
        self, num_updates=None, num_episodes=None, num_timesteps=None, environment=None,
        min_timesteps=None, states=None, actions=None, exclude_bool_action=False,
        exclude_int_action=False, exclude_float_action=False, exclude_bounded_action=False,
        require_observe=False, require_all=False, **agent
    ):
        """
        Generic unit-test.
        """
        agent, environment = self.prepare(
            environment=environment, min_timesteps=min_timesteps, states=states, actions=actions,
            exclude_bool_action=exclude_bool_action, exclude_int_action=exclude_int_action,
            exclude_float_action=exclude_float_action,
            exclude_bounded_action=exclude_bounded_action, require_observe=require_observe,
            require_all=require_all, **agent
        )

        self.runner = Runner(agent=agent, environment=environment)

        assert (num_updates is not None) + (num_episodes is not None) + \
            (num_timesteps is not None) <= 1
        if num_updates is None and num_episodes is None and num_timesteps is None:
            num_updates = self.__class__.num_updates
            num_episodes = self.__class__.num_episodes
            num_timesteps = self.__class__.num_timesteps
        if num_updates is None and num_episodes is None and num_timesteps is None:
            num_updates = 2
        assert (num_updates is not None) + (num_episodes is not None) + \
            (num_timesteps is not None) == 1

        evaluation = not any([
            require_all, require_observe, self.__class__.require_all,
            self.__class__.require_observe
        ])
        self.runner.run(
            num_episodes=num_episodes, num_timesteps=num_timesteps, num_updates=num_updates,
            use_tqdm=False, evaluation=evaluation
        )
        self.runner.close()
        agent.close()
        environment.close()

        self.finished_test()

[(75, 'warnings.filterwarnings', 'warnings.filterwarnings', False, 'import warnings\n'), (92, 'sys.stdout.flush', 'sys.stdout.flush', False, 'import sys\n'), (141, 'tensorforce.environments.Environment.create', 'Environment.create', False, 'from tensorforce.environments import Environment\n'), (154, 'tensorforce.agents.Agent.create', 'Agent.create', False, 'from tensorforce.agents import Agent\n'), (175, 'tensorforce.execution.Runner', 'Runner', False, 'from tensorforce.execution import Runner\n'), (103, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (104, 'sys.stdout.flush', 'sys.stdout.flush', False, 'import sys\n'), (134, 'test.unittest_environment.UnittestEnvironment', 'UnittestEnvironment', False, 'from test.unittest_environment import UnittestEnvironment\n'), (118, 'copy.deepcopy', 'deepcopy', False, 'from copy import deepcopy\n'), (121, 'copy.deepcopy', 'deepcopy', False, 'from copy import deepcopy\n'), (139, 'tensorforce.TensorforceError.unexpected', 'TensorforceError.unexpected', False, 'from tensorforce import TensorforceError\n'), (86, 'datetime.datetime.now', 'datetime.now', False, 'from datetime import datetime\n'), (90, 'datetime.datetime.now', 'datetime.now', False, 'from datetime import datetime\n')]