import rospy
import gym
from gym.utils import seeding
from .controllers_connection import ControllersConnection
from geometry_msgs.msg import Pose
from tf2_msgs.msg import TFMessage
from openai_ros.msg import RLExperimentInfo

class RobotRealEnv(gym.Env):

    def __init__(self, robot_name_space, controllers_list, reset_controls) -> None:
        # To reset Simulations
        rospy.logdebug("START init RobotRealEnv")
        self.controllers_object = ControllersConnection(namespace=robot_name_space, controllers_list=controllers_list)
        self.reset_controls = reset_controls
        self.seed()

        self.episode_num = 0
        self.cumulated_episode_reward = 0
        rospy.Subscriber("/tf", TFMessage, self._tf_real_sub)
        self._reward_pub = rospy.Publisher('/openai/reward', RLExperimentInfo, queue_size=1)
        self._tcp_real_pose_pub = rospy.Publisher('/tcp_pose_real', Pose, queue_size=10)
        self._tcp_real_msg = Pose()

        if self.reset_controls:
           self.controllers_object.reset_controllers()
        rospy.logdebug("END init RobotRealEnv")

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def step(self, action):
        """
        Function executed each time step.
        Here we get the action execute it in a time step and retrieve the
        observations generated by that action.
        :param action:
        :return: obs, reward, done, info
        """

        """
        Here we should convert the action num to movement action, execute the action in the
        simulation and get the observations result of performing that action.
        """
        rospy.logdebug("START STEP OpenAIROS")

        self._set_action(action)

        obs = self._get_obs()
        done = self._is_done(obs)
        info = {}
        reward = self._compute_reward(obs, done)
        self.cumulated_episode_reward += reward

        rospy.logdebug("END STEP OpenAIROS")

        return obs, reward, done, info

    def reset(self):
        rospy.logdebug("Reseting RobotRealEnvironment")
        self._reset_env()
        self._init_env_variables()
        self._update_episode()
        obs = self._get_obs()
        rospy.logdebug("END Reseting RobotRealEnvironment")
        return obs

    def close(self):
        """
        Function executed when closing the environment.
        Use it for closing GUIS and other systems that need closing.
        :return:
        """
        rospy.logdebug('Closing RobotRealEnvironment')
        rospy.signal_shutdown('Closing RobotRealEnvironment')

    def _update_episode(self):
        """
        Publishes the cumulated reward of the episode and
        increases the episode number by one.
        :return:
        """
        rospy.logwarn("PUBLISHING REWARD...")
        self._publish_reward_topic(
                                    self.cumulated_episode_reward,
                                    self.episode_num
                                    )
        rospy.logwarn("PUBLISHING REWARD...DONE="+str(self.cumulated_episode_reward)+",EP="+str(self.episode_num))

        self.episode_num += 1
        self.cumulated_episode_reward = 0

    def _tf_real_sub(self, data):
        if data.transforms[0].child_frame_id == "tool0_controller":
            self._tcp_real_msg.position.x = data.transforms[0].transform.translation.x
            self._tcp_real_msg.position.y = data.transforms[0].transform.translation.y
            self._tcp_real_msg.position.z = data.transforms[0].transform.translation.z
            self._tcp_real_msg.orientation = data.transforms[0].transform.rotation

    def _publish_reward_topic(self, reward, episode_number=1):
        """
        This function publishes the given reward in the reward topic for
        easy access from ROS infrastructure.
        :param reward:
        :param episode_number:
        :return:
        """
        reward_msg = RLExperimentInfo()
        reward_msg.episode_number = episode_number
        reward_msg.episode_reward = reward
        self._reward_pub.publish(reward_msg)

    def _reset_env(self):
        """Resets a simulation
        """
        rospy.logdebug("RESET ENV START")
        rospy.logwarn("DONT RESET CONTROLLERS")
        self._check_all_systems_ready()
        self._set_init_pose()
        self._check_all_systems_ready()

        rospy.logdebug("RESET ENV END")
        return True

    def _set_init_pose(self):
        """Sets the Robot in its init pose
        """
        raise NotImplementedError()

    def _check_all_systems_ready(self):
        """
        Checks that all the sensors, publishers and other simulation systems are
        operational.
        """
        raise NotImplementedError()

    def _get_obs(self):
        """Returns the observation.
        """
        raise NotImplementedError()

    def _init_env_variables(self):
        """Inits variables needed to be initialised each time we reset at the start
        of an episode.
        """
        raise NotImplementedError()

    def _set_action(self, action):
        """Applies the given action to the simulation.
        """
        raise NotImplementedError()

    def _is_done(self, observations):
        """Indicates whether or not the episode is done ( the robot has fallen for example).
        """
        raise NotImplementedError()

    def _compute_reward(self, observations, done):
        """Calculates the reward to give based on the observations given.
        """
        raise NotImplementedError()

    def _env_setup(self, initial_qpos):
        """Initial configuration of the environment. Can be used to configure initial state
        and extract information from the simulation.
        """
        raise NotImplementedError()