#!/usr/bin/env python
#################################################################################
# Copyright 2018 ROBOTIS CO., LTD.
#
# 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.
#################################################################################

# Authors: Gilbert, Trevor Sherrard #

import rospy
import numpy as np
import math
from math import pi
from geometry_msgs.msg import Twist, Point, Pose
from sensor_msgs.msg import LaserScan
from nav_msgs.msg import Odometry
from std_srvs.srv import Empty
from tf.transformations import euler_from_quaternion, quaternion_from_euler
from respawnGoal import Respawn
global numTurns
numTurns = 0

class Env():
    def __init__(self, action_size):
        self.goal_x = 0
        self.goal_y = 0
        self.heading = 0
        self.action_size = action_size
        self.initGoal = True
        self.get_goalbox = False
        self.position = Pose()
        self.pub_cmd_vel = rospy.Publisher('cmd_vel', Twist, queue_size=5)
        self.sub_odom = rospy.Subscriber('odom', Odometry, self.getOdometry)
        self.reset_proxy = rospy.ServiceProxy('gazebo/reset_simulation', Empty)
        self.unpause_proxy = rospy.ServiceProxy('gazebo/unpause_physics', Empty)
        self.pause_proxy = rospy.ServiceProxy('gazebo/pause_physics', Empty)
        self.respawn_goal = Respawn()

    def getGoalDistace(self):
        """
        This function computes the euclidean distance from the agent's
        current position to the goal state.

        params:
            self

        returns:
            distance (float): the distance to the goal state from the current state
        """
        goal_distance = round(math.hypot(self.goal_x - self.position.x, self.goal_y - self.position.y), 2)

        return goal_distance

    def getOdometry(self, odom):
        """
        This function computes the overall heading of the agent from recieved odom
        messages, and updates the env object heading field with that value.

        params:
            self
            odom (Odometry ROS message): the recieved ROS odometry message

        returns:
            None
        """
        self.position = odom.pose.pose.position
        orientation = odom.pose.pose.orientation
        orientation_list = [orientation.x, orientation.y, orientation.z, orientation.w]
        _, _, yaw = euler_from_quaternion(orientation_list)

        goal_angle = math.atan2(self.goal_y - self.position.y, self.goal_x - self.position.x)

        heading = goal_angle - yaw
        if heading > pi:
            heading -= 2 * pi

        elif heading < -pi:
            heading += 2 * pi

        self.heading = round(heading, 2)

    def getState(self, scan, cmdVel):
        """
        This function updates the env object with current information about the state
        of the agent within the enviornment. This function also detects if a collision has occured
        or if the agent has reached the goal state.

        params:
            self
            scan (ROS LaserScan Message): The recieved ROS laser scan message from the TurtleBot3 Lidar
            cmdVel (ROS Twist Message): The recieved Twist message containing target velocities for
            the TurtleBot3 agent.

        returns:
            A list of various state variables
        """
        scan_range = []
        heading = self.heading
        min_range = 0.13
        done = False

        for i in range(len(scan.ranges)):
            if scan.ranges[i] == float('Inf'):
                scan_range.append(3.5)
            elif np.isnan(scan.ranges[i]):
                scan_range.append(0)
            else:
                scan_range.append(scan.ranges[i])

        if min_range > min(scan_range) > 0:
            done = True

        current_distance = round(math.hypot(self.goal_x - self.position.x, self.goal_y - self.position.y),2)
        if current_distance < 0.2:
            self.get_goalbox = True

        return scan_range + [cmdVel.linear.x, cmdVel.angular.z] + [heading, current_distance], done

    def setReward(self, state, done, action):
        """
        This function computes the overal reward incurred by the agent in a given time step.

        params:
            self
            state (list of various types): The list of state variables generated in the above function
            done (Boolean): A boolean indicating a collision
            action (int): An integer corresponding to the action taken in the given time step

        returns:
            reward (float): The computed reward value
        """
        global numTurns
        runningTotal = 0
        punishmentStep = 0.05
        speedPunish = 0
        spinPunish = 0
        spinPunishStep = 5
        distancePunishFactor = 4

        current_distance = state[-1]
        heading = state[-2]
        angSpeed = state[-3]
        linSpeed = state[-4]
        laserScanPoints = state[0:360]

        for elem in laserScanPoints:
            if(elem < 0.3):
                runningTotal += punishmentStep

        if(abs(angSpeed) > 1.3):
            numTurns += 1

        if(numTurns > 30):
            spinPunish = spinPunishStep
            # reset numTurns
            numTurns = 0

        reward = -runningTotal - (current_distance * distancePunishFactor) - speedPunish - spinPunish

        if done:
            rospy.loginfo("Collision!!")
            reward = -200
            self.pub_cmd_vel.publish(Twist())
            numTurns = 0

        if self.get_goalbox:
            rospy.loginfo("Goal!!")
            reward = 200
            self.pub_cmd_vel.publish(Twist())
            self.goal_x, self.goal_y = self.respawn_goal.getPosition(True, delete=True)
            self.goal_distance = self.getGoalDistace()
            self.get_goalbox = False
            numTurns = 0

        return reward

    def step(self, action):
        """
        This function computes and publisher the overall Twist message values based off of the
        action recieved from the model. This function then uses the above functions to compute
        the current state and reward after the given action is completed.

        params:
            self
            action (int): An integer indicating what action the agent should take in the
                          given time step.

        returns:
            A list of various types containing compute state variables, the reward for
            a given action, and the status of a collision.
        """
        max_angular_vel = 1.5
        ang_vel = ((self.action_size - 1)/2 - action) * max_angular_vel * 0.5

        vel_cmd = Twist()
        vel_cmd.linear.x = 0.20
        vel_cmd.angular.z = ang_vel
        self.pub_cmd_vel.publish(vel_cmd)

        data = None
        while data is None:
            try:
                data = rospy.wait_for_message('scan', LaserScan, timeout=5)
            except:
                pass

        state, done = self.getState(data, vel_cmd)
        reward = self.setReward(state, done, action)

        return np.asarray(state), reward, done

    def reset(self):
        """
        This function resets the current simulation enviornment and all associate enviornment
        state variables.

        params:
            self

        returns:
            A list of various types containing various state variables
        """
        rospy.wait_for_service('gazebo/reset_simulation')
        try:
            self.reset_proxy()
        except (rospy.ServiceException) as e:
            print("gazebo/reset_simulation service call failed")

        data = None
        velData = None
        while data is None:
            try:
                data = rospy.wait_for_message('scan', LaserScan, timeout=5)
                velData = rospy.wait_for_message('cmd_vel', Twist, timeout=5)
            except:
                velData = Twist()

        if self.initGoal:
            self.goal_x, self.goal_y = self.respawn_goal.getPosition()
            self.initGoal = False

        self.goal_distance = self.getGoalDistace()
        state, done = self.getState(data, velData)

        return np.asarray(state)
