#! /usr/bin/env python

"""Implements the tabletop strategy."""

from math import pi, sin, cos
import rospy
import tf
import tf2_ros
from std_msgs.msg import Bool
from geometry_msgs.msg import PoseStamped, Twist
from sensor_msgs.msg import LaserScan
from tf.transformations import quaternion_from_euler, euler_from_quaternion
from ros_arduino_msgs.msg import Digital
from diff_drive.pose import Pose
import actionlib
import smach
import smach_ros
from smach import StateMachine
from smach_ros import SimpleActionState
import diff_drive
from diff_drive.msg import GoToPoseAction, GoToPoseGoal


class Names:
    DIFF_DRIVE_SERVER = 'diff_drive_go_to_goal'


class Outcomes:
    SUCCESS = 'success'
    FAILURE = 'failure'
    OPTION1 = 'option1'
    OPTION2 = 'option2'
    REPEAT = 'repeat'


class WaitForButton(smach.State):

    def __init__(self):
        smach.State.__init__(
            self, outcomes=[Outcomes.OPTION1, Outcomes.OPTION2],
            input_keys=['node'],
            output_keys=['node'])

    def execute(self, userdata):
        rate = rospy.Rate(10.0)
        while not rospy.is_shutdown() and userdata.node.button_a==0 \
           and userdata.node.button_b==0:
            rate.sleep()

        # Wait a bit for user to release button.
        a_value = userdata.node.button_a
        rospy.sleep(0.5)
        if a_value != 0:
            return Outcomes.OPTION1
        else:
            return Outcomes.OPTION2


class FindObject(smach.State):

    def __init__(self):
        smach.State.__init__(
            self, outcomes=[Outcomes.SUCCESS, Outcomes.FAILURE, Outcomes.OPTION1, Outcomes.OPTION2, Outcomes.REPEAT],
            input_keys=['node', 'frame', 'delay'],
            output_keys=['node', 'frame', 'pose', 'delay'])

    def execute(self, userdata):
        rospy.sleep(userdata.delay)
        pose = userdata.node.get_frame_pose(userdata.frame)
        rospy.loginfo('%s pose: %s', userdata.frame, pose)
        if pose is None:
            return Outcomes.FAILURE
        else:
            userdata.pose = pose
            return Outcomes.SUCCESS


class MoveToPose(smach.State):

    def __init__(self):
        smach.State.__init__(
            self, outcomes=[Outcomes.SUCCESS, Outcomes.FAILURE],
            input_keys=['object_pose', 'node', 'offset', 'angle'],
            output_keys=['object_pose', 'node', 'offset', 'angle'])

    def execute(self, userdata):
        object_pose = userdata.object_pose
        desired_pose = Pose()
        desired_pose.x = object_pose.x + userdata.offset
        desired_pose.y = object_pose.y
        desired_pose.theta = userdata.node.normalize_angle(userdata.angle)

        q_pose = userdata.node.get_quaternion_pose(desired_pose)

        action_goal = GoToPoseGoal()
        action_goal.pose.pose.position.x = q_pose.pose.position.x
        action_goal.pose.pose.position.y = q_pose.pose.position.y
        action_goal.pose.pose.position.z = q_pose.pose.position.z
        action_goal.pose.pose.orientation.x = q_pose.pose.orientation.x
        action_goal.pose.pose.orientation.y = q_pose.pose.orientation.y
        action_goal.pose.pose.orientation.z = q_pose.pose.orientation.z
        action_goal.pose.pose.orientation.w = q_pose.pose.orientation.w
        success = userdata.node.send_goal(action_goal)
        if success:
            return Outcomes.SUCCESS
        else:
            return Outcomes.FAILURE


class MoveToRelativePose(smach.State):

    def __init__(self):
        smach.State.__init__(
            self, outcomes=[Outcomes.SUCCESS, Outcomes.FAILURE],
            input_keys=['object_pose', 'node', 'offset', 'offset_angle'],
            output_keys=['object_pose', 'node', 'offset', 'offset_angle'])

    def execute(self, userdata):
        object_pose = userdata.object_pose
        desired_pose = Pose()
        desired_pose.x = object_pose.x + userdata.offset*cos(object_pose.theta)
        desired_pose.y = object_pose.y + userdata.offset*sin(object_pose.theta)
        desired_pose.theta = userdata.node.normalize_angle(
            object_pose.theta + userdata.offset_angle)

        q_pose = userdata.node.get_quaternion_pose(desired_pose)

        action_goal = GoToPoseGoal()
        action_goal.pose.pose.position.x = q_pose.pose.position.x
        action_goal.pose.pose.position.y = q_pose.pose.position.y
        action_goal.pose.pose.position.z = q_pose.pose.position.z
        action_goal.pose.pose.orientation.x = q_pose.pose.orientation.x
        action_goal.pose.pose.orientation.y = q_pose.pose.orientation.y
        action_goal.pose.pose.orientation.z = q_pose.pose.orientation.z
        action_goal.pose.pose.orientation.w = q_pose.pose.orientation.w
        success = userdata.node.send_goal(action_goal)
        if success:
            return Outcomes.SUCCESS
        else:
            return Outcomes.FAILURE


class DriveToEdge(smach.State):

    def __init__(self):
        smach.State.__init__(
            self, outcomes=[Outcomes.SUCCESS, Outcomes.FAILURE],
            input_keys=['node'],
            output_keys=['node'])

    def execute(self, userdata):
        rate = rospy.Rate(10.0)
        while not rospy.is_shutdown():
            if userdata.node.ranges is not None:
                mid = (len(userdata.node.ranges) + 1) / 2
                dist = userdata.node.ranges[mid]
                if dist < 0.2:
                    userdata.node.set_speed(0, 0)
                    rospy.sleep(1.0)
                    dist = userdata.node.ranges[mid]
                    if dist < 0.2:
                        break;

                if dist > 0.5:
                    diff = 0
                    theta_vel = 0
                else:
                    right = self.median(userdata.node.ranges, mid-2)
                    left = self.median(userdata.node.ranges, mid+2)
                    diff = right - left
                    theta_vel = max(min(diff * userdata.node.k_angle,
                                        userdata.node.max_theta_vel),
                                    -userdata.node.max_theta_vel)

                rospy.loginfo('dist=%f diff=%f theta_vel=%f',
                              dist, diff, theta_vel)
                userdata.node.set_speed(0.07, theta_vel)

            rate.sleep()

        return Outcomes.SUCCESS

    def median(self, vector, i):
        left, mid, right = (vector[i-1], vector[i], vector[i+1])
        if left<=mid and mid<=right:
            return mid
        elif right<=mid and mid<=right:
            return mid
        elif mid<=left and left<=right:
            return left
        elif right<=left and left<=mid:
            return left
        else:
            return right


class Strategy:

    def run(self):
        rospy.init_node('tabletop_strategy')

        self.tf_buffer = tf2_ros.Buffer()
        self.listener = tf2_ros.TransformListener(self.tf_buffer)

        self.action_client = actionlib.SimpleActionClient(
            Names.DIFF_DRIVE_SERVER, GoToPoseAction)

        rate = rospy.Rate(rospy.get_param('~rate', 10.0))

        self.button_a = 0
        self.button_b = 0
        self.block_frame = rospy.get_param('~block_frame', 'block')
        self.goal_frame = rospy.get_param('~block_frame', 'goal')
        self.reference_frame = rospy.get_param('~reference_frame', 'odom')
        self.k_speed = rospy.get_param('~k_speed', 2.0)
        self.k_angle = rospy.get_param('~k_angle', 10.0)
        self.max_x_vel = rospy.get_param('~max_linear_vel', 0.1)
        self.max_theta_vel = rospy.get_param('~max_theta_vel', 0.3)
        self.ranges = None

        self.goal_pub = rospy.Publisher('move_base_simple/goal', PoseStamped,
                                        queue_size=1)

        self.vel_pub = rospy.Publisher('robot/cmd_vel', Twist, queue_size=1)

        rospy.Subscriber('goal_achieved', Bool, self.on_goal_achieved)
        rospy.Subscriber('arduino/sensor/button_a', Digital,
                         self.on_button_a)
        rospy.Subscriber('arduino/sensor/button_b', Digital,
                         self.on_button_b)
        rospy.Subscriber('scan', LaserScan, self.on_laser_scan)

        sm = StateMachine([Outcomes.SUCCESS, Outcomes.FAILURE,
                           Outcomes.OPTION1, Outcomes.OPTION2,
                           Outcomes.REPEAT])
        sm.userdata.node = self
        sm.userdata.home_pose = PoseStamped()
        sm.userdata.robot_pose = PoseStamped()
        sm.userdata.block_pose = PoseStamped()
        sm.userdata.goal_pose = PoseStamped()
        sm.userdata.robot_frame = 'base_link'
        sm.userdata.block_frame = 'block'
        sm.userdata.goal_frame = 'goal'

        sm.userdata.object_find_delay = 2.0
        sm.userdata.no_delay = 0.0
        sm.userdata.half_turn = pi
        sm.userdata.zero_offset = 0
        sm.userdata.zero_angle = 0
        sm.userdata.qtr_turn = pi/2
        sm.userdata.neg_qtr_turn = -pi/2

        sm.userdata.face_block_offset = -0.33
        sm.userdata.approach_block_offset = -0.18
        sm.userdata.engage_block_offset = -0.1

        sm.userdata.face_goal_offset = -0.55
        sm.userdata.approach_goal_offset = -0.40
        sm.userdata.leave_block_offset = -0.20
        sm.userdata.retreat_offset = -0.50
        sm.userdata.home_offset = -1.6
        sm.userdata.face_end_offset = 0.18

        with sm:
            smach.StateMachine.add('WAIT_FOR_BUTTON', WaitForButton(),
                                   transitions={
                                       Outcomes.OPTION1: 'RG_START',
                                       Outcomes.OPTION2: 'NEW_START'
                                   })

            smach.StateMachine.add('RG_START', WaitForButton(),
                                   transitions={
                                       Outcomes.OPTION1: 'FIND_HOME',
                                       Outcomes.OPTION2: 'FIND_RETURN_POS'
                                   })

            smach.StateMachine.add('FIND_HOME', FindObject(),
                                   remapping={
                                       'frame': 'robot_frame',
                                       'pose': 'home_pose',
                                       'delay': 'object_find_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'FIND_BLOCK',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('FIND_BLOCK', FindObject(),
                                   remapping={
                                       'frame': 'block_frame',
                                       'pose': 'block_pose',
                                       'delay': 'no_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'FACE_BLOCK',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('FACE_BLOCK', MoveToPose(),
                                   remapping={
                                       'object_pose': 'block_pose',
                                       'offset': 'face_block_offset',
                                       'angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'FIND_GOAL'
                                   })

            smach.StateMachine.add('FIND_GOAL', FindObject(),
                                   remapping={
                                       'frame': 'goal_frame',
                                       'pose': 'goal_pose',
                                       'delay': 'object_find_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'REFIND_BLOCK',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('REFIND_BLOCK', FindObject(),
                                   remapping={
                                       'frame': 'block_frame',
                                       'pose': 'block_pose',
                                       'delay': 'no_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'APPROACH_BLOCK',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('APPROACH_BLOCK', MoveToPose(),
                                   remapping={
                                       'object_pose': 'block_pose',
                                       'offset': 'approach_block_offset',
                                       'angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'PICK_UP_BLOCK'
                                   })


            smach.StateMachine.add('PICK_UP_BLOCK', MoveToPose(),
                                   remapping={
                                       'object_pose': 'block_pose',
                                       'offset': 'engage_block_offset',
                                       'angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'FACE_GOAL'
                                   })

            smach.StateMachine.add('FACE_GOAL', MoveToPose(),
                                   remapping={
                                       'object_pose': 'goal_pose',
                                       'offset': 'face_goal_offset',
                                       'angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'REFIND_GOAL',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('REFIND_GOAL', FindObject(),
                                   remapping={
                                       'frame': 'goal_frame',
                                       'pose': 'goal_pose',
                                       'delay': 'object_find_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'LEAVE_BLOCK',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('LEAVE_BLOCK', MoveToPose(),
                                   remapping={
                                       'object_pose': 'goal_pose',
                                       'offset': 'leave_block_offset',
                                       'angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'RETREAT_FROM_GOAL'
                                   })

            smach.StateMachine.add('RETREAT_FROM_GOAL', MoveToPose(),
                                   remapping={
                                       'object_pose': 'goal_pose',
                                       'offset': 'retreat_offset',
                                       'angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('FIND_RETURN_POS', FindObject(),
                                   remapping={
                                       'frame': 'robot_frame',
                                       'pose': 'home_pose',
                                       'delay': 'object_find_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'FIND_TABLE_END',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('FIND_TABLE_END', FindObject(),
                                   remapping={
                                       'frame': 'goal_frame',
                                       'pose': 'goal_pose',
                                       'delay': 'no_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'FACE_TABLE_END',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('FACE_TABLE_END', MoveToPose(),
                                   remapping={
                                       'object_pose': 'goal_pose',
                                       'offset': 'face_goal_offset',
                                       'angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'REFIND_TABLE_END'
                                   })

            smach.StateMachine.add('REFIND_TABLE_END', FindObject(),
                                   remapping={
                                       'frame': 'goal_frame',
                                       'pose': 'goal_pose',
                                       'delay': 'object_find_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'DRIVE_TO_END',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('DRIVE_TO_END', MoveToPose(),
                                   remapping={
                                       'object_pose': 'goal_pose',
                                       'offset': 'leave_block_offset',
                                       'angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'DO_RETURN'
                                   })

            smach.StateMachine.add('DO_RETURN', MoveToPose(),
                                   remapping={
                                       'object_pose': 'home_pose',
                                       'offset': 'zero_offset',
                                       'angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'WAIT_FOR_BUTTON'
                                   })

            #
            # New states for HBRC demo.
            #
            
            smach.StateMachine.add('NEW_START', WaitForButton(),
                                   transitions={
                                       Outcomes.OPTION1: 'DRIVE_TO_SIDE_EDGE',
                                       Outcomes.OPTION2: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('DRIVE_TO_SIDE_EDGE', DriveToEdge(),
                                   transitions={
                                       Outcomes.SUCCESS: 'FIND_SIDE_EDGE_POSE'
                                   })

            smach.StateMachine.add('FIND_SIDE_EDGE_POSE', FindObject(),
                                   remapping={
                                       'frame': 'robot_frame',
                                       'pose': 'robot_pose',
                                       'delay': 'no_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'TURN_TO_RIGHT',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('TURN_TO_RIGHT', MoveToRelativePose(),
                                   remapping={
                                       'object_pose': 'robot_pose',
                                       'offset': 'zero_offset',
                                       'offset_angle': 'neg_qtr_turn'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'DRIVE_TO_NEAR_END'
                                   })

            smach.StateMachine.add('DRIVE_TO_NEAR_END', DriveToEdge(),
                                   transitions={
                                       Outcomes.SUCCESS: 'FIND_NEAR_END_POSE'
                                   })

            smach.StateMachine.add('FIND_NEAR_END_POSE', FindObject(),
                                   remapping={
                                       'frame': 'robot_frame',
                                       'pose': 'robot_pose',
                                       'delay': 'no_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'FACE_NEAR_END',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('FACE_NEAR_END', MoveToRelativePose(),
                                   remapping={
                                       'object_pose': 'robot_pose',
                                       'offset': 'face_end_offset',
                                       'offset_angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'FIND_END_POSE'
                                   })

            smach.StateMachine.add('FIND_END_POSE', FindObject(),
                                   remapping={
                                       'frame': 'robot_frame',
                                       'pose': 'robot_pose',
                                       'delay': 'no_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'TURN_AROUND',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('TURN_AROUND', MoveToRelativePose(),
                                   remapping={
                                       'object_pose': 'robot_pose',
                                       'offset': 'zero_offset',
                                       'offset_angle': 'half_turn'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'DRIVE_TO_FAR_END'
                                   })

            smach.StateMachine.add('DRIVE_TO_FAR_END', DriveToEdge(),
                                   transitions={
                                       Outcomes.SUCCESS: 'FIND_FAR_END_POSE'
                                   })

            smach.StateMachine.add('FIND_FAR_END_POSE', FindObject(),
                                   remapping={
                                       'frame': 'robot_frame',
                                       'pose': 'robot_pose',
                                       'delay': 'no_delay'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'FACE_FAR_END',
                                       Outcomes.FAILURE: 'WAIT_FOR_BUTTON'
                                   })

            smach.StateMachine.add('FACE_FAR_END', MoveToRelativePose(),
                                   remapping={
                                       'object_pose': 'robot_pose',
                                       'offset': 'face_end_offset',
                                       'offset_angle': 'zero_angle'
                                   },
                                   transitions={
                                       Outcomes.SUCCESS: 'WAIT_FOR_BUTTON'
                                   })

        result = sm.execute()
        rospy.loginfo('Mission result: %s', result)


    def send_goal(self, goal):
        self.action_client.wait_for_server()
        self.action_client.send_goal(goal)
        self.action_client.wait_for_result()
        result = self.action_client.get_result()
        return result.success

    def on_button_a(self, msg):
        self.button_a = msg.value

    def on_button_b(self, msg):
        self.button_b = msg.value

    def on_laser_scan(self, msg):
        self.ranges = msg.ranges

    def on_goal_achieved(self, msg):
        pass

    def get_frame_pose(self, frame):
        try:
            trans = self.tf_buffer.lookup_transform(self.reference_frame,
                                                    frame, rospy.Time(0))
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException,
                tf2_ros.ExtrapolationException) as e:
            rospy.logerr('Cannot get transform %s --> %s: %s',
                         self.reference_frame, frame, str(e))
            return

        q = [trans.transform.rotation.x,
             trans.transform.rotation.y,
             trans.transform.rotation.z,
             trans.transform.rotation.w]
        roll, pitch, yaw = euler_from_quaternion(q)

        pose = Pose()
        pose.x = trans.transform.translation.x
        pose.y = trans.transform.translation.y
        pose.theta = yaw

        return pose

    def get_quaternion_pose(self, angle_pose):
        q = quaternion_from_euler(0, 0, angle_pose.theta)

        pose = PoseStamped()
        pose.header.frame_id = self.reference_frame
        pose.pose.position.x = angle_pose.x
        pose.pose.position.y = angle_pose.y
        pose.pose.orientation.x = q[0]
        pose.pose.orientation.y = q[1]
        pose.pose.orientation.z = q[2]
        pose.pose.orientation.w = q[3]
        return pose

    def send_goal_pose(self, pose):
        self.goal_pub.publish(self.get_quaternion_pose(pose))

    def set_speed(self, x_vel, theta_vel):
        rospy.loginfo('setting speed: %f %f', x_vel, theta_vel)
        twist = Twist()
        twist.linear.x = x_vel
        twist.angular.z = theta_vel
        self.vel_pub.publish(twist)

    def normalize_angle(self, angle):
        while angle >= 2*pi:
            angle -= 2*pi
        while angle < 0:
            angle += 2*pi
        return angle


if __name__ == "__main__":
    try:
        Strategy().run()

    except rospy.ROSInterruptException:
        pass
