#################################################################################
#   Copyright Amazon.com, Inc. or its affiliates. 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 typing import Any, Callable, Optional, Iterable
from unittest import TestCase
from unittest.mock import patch, MagicMock, call
import inspect

from deepsim_envs.agents.abs_agent import AbstractAgent

from deepsim import (
    DeepSimCallbackError,
    Tag,
    AbstractModelSpawner,
    Pose
)

from rosgraph_msgs.msg import Clock

myself: Callable[[], Any] = lambda: inspect.stack()[1][3]


class DummyAgent(AbstractAgent):
    def __init__(self, name: str,
                 spawner: Optional[AbstractModelSpawner] = None, *,
                 initial_pose: Optional[Pose] = None,
                 max_retry_attempts: int = 10,
                 backoff_time_sec: float = 1.0,
                 spawn_on_init: bool = True,
                 enable_on_init: bool = True,
                 tags: Optional[Iterable[str]] = None,
                 **kwargs) -> None:
        self.mock = MagicMock()
        super(DummyAgent, self).__init__(name=name,
                                         spawner=spawner,
                                         initial_pose=initial_pose,
                                         max_retry_attempts=max_retry_attempts,
                                         backoff_time_sec=backoff_time_sec,
                                         spawn_on_init=spawn_on_init,
                                         enable_on_init=enable_on_init,
                                         tags=tags,
                                         **kwargs)

    def on_enable(self) -> None:
        self.mock.on_enable()

    def on_disable(self) -> None:
        self.mock.on_disable()

    def on_delete(self) -> None:
        self.mock.on_delete()

    def on_update(self) -> None:
        self.mock.on_update()

    def on_fixed_update(self, delta_time: float, sim_time: Clock) -> None:
        self.mock.on_fixed_update(delta_time, sim_time)

    def on_reset_agent(self) -> None:
        self.mock.on_reset_agent()

    def on_action_received(self, action: Any) -> None:
        self.mock.on_action_received(action)

    def get_next_state(self) -> Any:
        return self.mock.get_next_state()

    def on_episode_begin(self) -> None:
        self.mock.on_episode_begin()

    def on_episode_end(self) -> None:
        self.mock.on_episode_end()


@patch("deepsim.behaviours.behaviour_manager.BehaviourManager")
@patch("deepsim.behaviours.deepsim_behaviour.Transform")
class AbstractAgentTest(TestCase):
    def setUp(self) -> None:
        pass

    def test_initialize(self, transform_mock, behaviour_manager_mock):
        name = myself()

        agent = DummyAgent(name)
        assert Tag.AGENT in agent.tags
        assert not agent.done
        assert agent.last_step_reward == 0.0
        assert agent.episode_reward == 0.0
        assert agent.last_action is None

    def test_reward_setter(self, transform_mock, behaviour_manager_mock):
        name = myself()

        agent = DummyAgent(name)
        assert agent.last_step_reward == 0.0
        assert agent.episode_reward == 0.0

        agent.add_reward(5.0)
        assert agent.last_step_reward == 5.0
        assert agent.episode_reward == 5.0

        # On next step
        agent.add_reward(3.0)
        assert agent.last_step_reward == 3.0
        assert agent.episode_reward == 8.0

    def test_act(self, transform_mock, behaviour_manager_mock):
        name = myself()

        agent = DummyAgent(name)

        action = MagicMock()

        agent.act(action)

        assert agent.last_action == action
        agent.mock.on_action_received.assert_called_once_with(action)

    def test_act_callback_error(self, transform_mock, behaviour_manager_mock):
        name = myself()

        agent = DummyAgent(name)
        agent.mock.on_action_received.side_effect = Exception()

        action = MagicMock()
        with self.assertRaises(DeepSimCallbackError):
            agent.act(action)

        assert agent.last_action == action

    def test_end_episode(self, transform_mock, behaviour_manager_mock):
        name = myself()

        agent = DummyAgent(name)

        assert not agent.done
        agent.end_episode()
        assert agent.done

        agent.mock.on_episode_end.assert_called_once()

    def test_end_episode_callback_error(self, transform_mock, behaviour_manager_mock):
        name = myself()

        agent = DummyAgent(name)
        agent.mock.on_episode_end.side_effect = Exception()

        with self.assertRaises(DeepSimCallbackError):
            assert not agent.done
            agent.end_episode()

        assert agent.done

    def test_reset(self, transform_mock, behaviour_manager_mock):
        name = myself()

        agent = DummyAgent(name)

        agent.end_episode()
        agent.add_reward(3.0)
        agent.add_reward(5.0)
        assert agent.done
        assert agent.last_step_reward == 5.0
        assert agent.episode_reward == 8.0
        agent.reset()
        assert agent.last_step_reward == 0.0
        assert agent.episode_reward == 0.0
        assert not agent.done

        agent.mock.on_reset_agent.assert_called_once()

    def test_reset_callback_error(self, transform_mock, behaviour_manager_mock):
        name = myself()

        agent = DummyAgent(name)
        agent.mock.on_reset_agent.side_effect = Exception()

        with self.assertRaises(DeepSimCallbackError):
            agent.end_episode()
            agent.add_reward(3.0)
            agent.add_reward(5.0)
            assert agent.done
            assert agent.last_step_reward == 5.0
            assert agent.episode_reward == 8.0
            agent.reset()
        assert not agent.done
        assert agent.last_step_reward == 0.0
        assert agent.episode_reward == 0.0