#   +-----------------------------------------------+
#   | RL-ROBOT. Reinforcement Learning for Robotics |
#   | Angel Martinez-Tenor                          |
#   | MAPIR. University of Malaga. 2016             |
#   +-----------------------------------------------+
""" Generic Robot module """

import math
import sys

import numpy as np

import exp

link_dict = {"ROS": "rl_ros", "VREP": "rl_vrep", "MODEL": "rl_vrep"}
try:
    link_module = link_dict[exp.ENVIRONMENT_TYPE]
except KeyError:
    sys.exit("ENVIRONMENT_TYPE " + exp.ENVIRONMENT_TYPE + " undefined \n")
link = __import__(link_module)

# 2D navigation

mobilebase_pose2d = np.full(3, -1, dtype=np.float64)  # x(m),y(m),theta(rad)
last_mobilebase_pose2d = np.full(3, -1, dtype=np.float64)
mobilebase_displacement2d = 0  # between two consecutive steps (m)

dist_obstacle = np.full(link.N_LASERS, -1, dtype=np.float64)

last_distance_mobilebase_goal = -1
distance_mobilebase_goal = -1
mobilebase_goal_displacement2d = -1

# 3D motion
gripper_pose3d = np.full(3, -1, dtype=np.float64)
last_gripper_pose3d = np.full(3, -1, dtype=np.float64)
gripper_displacement3d = -1

goal_pose3d = np.full(3, -1, dtype=np.float64)
last_goal_pose3d = np.full(3, -1, dtype=np.float64)
goal_displacement3d = -1  # Just In case we want to move an object

last_distance_gripper_goal = -1
distance_gripper_goal = -1
gripper_goal_displacement3d = -1

AGENT_ELEMENTS = []
ENV_ELEMENTS = []

initiated = False

sensor = {}


def setup(agent_elem, environment_elem):
    """initialize robot"""
    global AGENT_ELEMENTS, ENV_ELEMENTS
    global mobilebase_pose2d, last_mobilebase_pose2d, mobilebase_displacement2d
    global dist_obstacle
    global gripper_pose3d, last_gripper_pose3d, gripper_displacement3d
    global goal_pose3d, last_goal_pose3d, goal_displacement3d
    global last_distance_mobilebase_goal, distance_mobilebase_goal
    global mobilebase_goal_displacement2d
    global last_distance_gripper_goal, distance_gripper_goal
    global gripper_goal_displacement3d
    global initiated

    global sensor

    AGENT_ELEMENTS = agent_elem
    ENV_ELEMENTS = environment_elem

    mobilebase_pose2d = np.full(3, -1, dtype=np.float64)
    last_mobilebase_pose2d = np.full(3, -1, dtype=np.float64)

    mobilebase_displacement2d = 0

    dist_obstacle = np.full(link.N_LASERS, -1, dtype=np.float64)

    gripper_pose3d = np.full(3, -1, dtype=np.float64)
    last_gripper_pose3d = np.full(3, -1, dtype=np.float64)
    gripper_displacement3d = -1

    goal_pose3d = np.full(3, -1, dtype=np.float64)
    last_goal_pose3d = np.full(3, -1, dtype=np.float64)
    goal_displacement3d = -1

    last_distance_mobilebase_goal = -1
    distance_mobilebase_goal = -1
    mobilebase_goal_displacement2d = -1

    last_distance_gripper_goal = -1
    distance_gripper_goal = -1
    gripper_goal_displacement3d = -1

    update()
    update()  # To obtain the first relative displacement

    return


def update():
    """update robot & environment state (sensors, locations...)"""
    global mobilebase_pose2d, last_mobilebase_pose2d, mobilebase_displacement2d
    global dist_obstacle
    global gripper_pose3d, last_gripper_pose3d, gripper_displacement3d
    global goal_pose3d, last_goal_pose3d, goal_displacement3d
    global last_distance_mobilebase_goal, distance_mobilebase_goal
    global mobilebase_goal_displacement2d
    global last_distance_gripper_goal, distance_gripper_goal
    global gripper_goal_displacement3d

    global sensor

    if "DISTANCE_SENSOR" in AGENT_ELEMENTS:
        dist_obstacle = get_distance_obstacle()

    if "MOBILE_BASE" in AGENT_ELEMENTS:
        last_mobilebase_pose2d = mobilebase_pose2d
        mobilebase_pose2d = get_mobilebase_pose2d()
        mobilebase_displacement2d = distance2d(
            mobilebase_pose2d, last_mobilebase_pose2d
        )

    if "ARM" in AGENT_ELEMENTS:
        last_gripper_pose3d = gripper_pose3d
        gripper_pose3d = link.get_gripper_pose3d()
        gripper_displacement3d = distance3d(gripper_pose3d, last_gripper_pose3d)

    if "GOAL_OBJECT" in ENV_ELEMENTS:
        last_goal_pose3d = link.get_goal_pose_3d()
        goal_pose3d = link.get_goal_pose_3d()
        goal_displacement3d = distance3d(goal_pose3d, last_goal_pose3d)

        if "MOBILE_BASE" in AGENT_ELEMENTS:
            last_distance_mobilebase_goal = distance_mobilebase_goal
            distance_mobilebase_goal = distance2d(goal_pose3d, mobilebase_pose2d)
            mobilebase_goal_displacement2d = (
                distance_mobilebase_goal - last_distance_mobilebase_goal
            )
            # Negative: the mobilebase is getting closer

        if "ARM" in AGENT_ELEMENTS:
            last_distance_gripper_goal = distance_gripper_goal
            distance_gripper_goal = distance3d(goal_pose3d, gripper_pose3d)
            gripper_goal_displacement3d = (
                distance_gripper_goal - last_distance_gripper_goal
            )
            # Negative: the arm is getting closer

    sensor["mobile_x"] = mobilebase_pose2d[0]
    sensor["mobile_y"] = mobilebase_pose2d[1]
    sensor["mobile_theta"] = mobilebase_pose2d[2]

    sensor["gripper_x"] = gripper_pose3d[0]
    sensor["gripper_y"] = gripper_pose3d[1]
    sensor["gripper_z"] = gripper_pose3d[2]

    sensor["goal_x"] = goal_pose3d[0]
    sensor["goal_y"] = goal_pose3d[0]

    sensor["laser_front"] = dist_obstacle[0]
    sensor["laser_front_left"] = dist_obstacle[1]
    sensor["laser_left"] = dist_obstacle[2]
    sensor["laser_rear_left"] = dist_obstacle[3]
    sensor["laser_rear"] = dist_obstacle[4]
    sensor["laser_rear_right"] = dist_obstacle[5]
    sensor["laser_right"] = dist_obstacle[6]
    sensor["laser_front_right"] = dist_obstacle[7]


def move_wheels(left_wheel, right_wheel):
    """move base wheels (inputs: rad/s)"""
    link.move_wheels(left_wheel, right_wheel)
    return


def stop_motion():
    """stop the mobilebase"""
    link.stop_motion()
    return


def move_full_arm(arm, biceps, forearm):
    """move robotic arm"""
    link.move_arm(arm)
    link.move_biceps(biceps)
    link.move_forearm(forearm)
    return


def get_distance_obstacle():
    """return distances to objects measured by laser"""
    di = link.get_distance_obstacle()
    return di


def get_mobilebase_pose2d():
    """return 2d pose of the mobilebase (x,y,theta)"""
    po = link.get_mobilebase_pose2d()
    return po


def get_gripper_pose3d():
    """return the position of the gripper:  [ x(m), y(m), z(m) ]"""
    po = link.get_gripper_pose3d()
    return po


def get_goal_pose3d():
    """return the position of the goal object:  [ x(m), y(m), z(m) ]"""
    po = link.get_goal_pose_3d()
    return po


def distance2d(pose_a, pose_b):
    """get distance 2d (x,y axis) between 2 poses"""
    delta_pose2d = abs(pose_a - pose_b)
    displacement2d = math.sqrt(delta_pose2d[0] ** 2 + delta_pose2d[1] ** 2)
    return displacement2d


def distance3d(pose_a, pose_b):
    """get distance 3d between 2 poses"""

    delta_pose3d = abs(pose_a - pose_b)
    displacement3d = math.sqrt(
        delta_pose3d[0] ** 2 + delta_pose3d[1] ** 2 + delta_pose3d[2] ** 2
    )
    return displacement3d


def start():
    """Start robot"""
    link.start()
    return


def stop():
    """stop robot"""
    link.stop()
    return


def setup_devices():
    """setup robot's devices"""
    link.setup_devices()
    return


def connect():
    """connect to the robot"""
    link.connect()
    return


def disconnect():
    """disconnect from the robot"""
    link.disconnect()
    return