#!/usr/bin/env python
#
#   moveskeleton.py
#
#   The is the skeleton for the move node.  This execute movements,
#   first simple movements, and ultimately via a planner.  This will
#   also use/update the map with obstacles, as appropriate/needed.
#
#   Node:       /move
#
#   Params:     none yet?
#
#   Subscribe:  /pose                   geometry_msgs/PoseStamped
#               /move_base_simple/goal  geometry_msgs/PoseStamped
#               /scan                   sensor_msgs/LaserScan
#               /map                    nav_msgs/OccupancyGrid
#
#   Publish:    /vel_cmd                geometry_msgs/Twist
#               /obstacles              nav_msgs/OccupancyGrid
#
#   TF Required:  map -> laser
#
import math
import numpy as np
import rospy
import sys
import tf2_ros

from geometry_msgs.msg      import PoseStamped
from geometry_msgs.msg      import TransformStamped
from geometry_msgs.msg      import Twist
from nav_msgs.msg           import OccupancyGrid
from sensor_msgs.msg        import LaserScan

from PlanarTransform import PlanarTransform


#
#   Constants
#
#   Perhaps these should be parameters?
#
VMAX = 0.20             # Max forward speed (m/s)
WMAX = 1.00             # Max rotational speed (rad/sec)

# Other constants?

TPUBMAP = 3.0           # Time between publishing the obstacle map


#
#   Global Variables
#
# The obstancle map info.
map2grid  = None        # Map bottom/left corner (origin) in map frame
grid2map  = None        # Map frame relative to the bottom/left corner
w         = 0           # Width
h         = 0           # Height
res       = 1.0         # Resolution (meters/pixel)
obstacles = None        # The actual map

# Desired - these will likely become a waypoint list...
xdes     = 2.0
ydes     = 0.0
thetades = 0.0

# Flag whether to start moving
drive = False


######################################################################
#   Driving
######################################################################
#
#   Goal Pose Callback
#
#   This comes from the RVIZ "2D Nav Goal" button, telling us where we
#   want the robot to go.
#
def callback_goal(posemsg):
    # Make sure the goal pose is published w.r.t. the map frame.
    if not (posemsg.header.frame_id == 'map'):
        rospy.logerr("Goal Pose is not in 'map' frame!")
        return

    # Extract the bot's goal pose w.r.t. the map frame and report.
    map2goal = PlanarTransform.fromPose(posemsg.pose)
    rospy.loginfo("New target: %s" % map2goal)

    #
    #   CALL A PLANNER HERE...  For now, we simply drive directly to
    #   the goal position, without consideration of the map!
    #
    # Use the goal x/y/theta as the desired value for the current motion.
    global xdes, ydes, thetades
    xdes     = map2goal.x()
    ydes     = map2goal.y()
    thetades = map2goal.theta()

    # Start the motion by turning the drive flag on.
    global drive
    drive = True


#
#   Actual Pose Callback
#
#   When we receive an actual pose, update the velocity command to
#   head to the next desired waypoint.
#
def callback_pose(posemsg):
    # Use/change the global drive flag.
    global drive

    # Make sure the actual pose is published w.r.t. the map frame.
    if not (posemsg.header.frame_id == 'map'):
        rospy.logerr("Actual pose is not in 'map' frame!")
        return

    # Grab the actual x/y/theta.
    map2pose = PlanarTransform.fromPose(posemsg.pose)
    x     = map2pose.x()
    y     = map2pose.y()
    theta = map2pose.theta()

    #
    #   DETERMINE VELOCITY COMMANDS!  You may want to turn off the
    #   drive flag when you want to stop?
    #
    # For fun, drive in a CCW circle...
    vx = 0.08
    wz = 1.0


    # Make sure we are supposed to be driving!  This means the bot
    # doesn't drive until the first goal is received or after
    # something turns if off.
    if not drive:
        wz = 0.0
        vx = 0.0

    # Send the commands.
    cmdmsg = Twist()
    cmdmsg.linear.x  = vx
    cmdmsg.angular.z = wz
    cmdpub.publish(cmdmsg)


######################################################################
#   Mapping
######################################################################
#
#   Grab and Pre-Process the Hand-Created Map
#
#   We could do this repeatedly, i.e. set up a subscriber.  But I'm
#   assuming the map will remain constant?
#
def prepareMap():
    # Grab a single message of the map topic.  Give it 10 seconds.
    rospy.loginfo("Move: waiting for a map...")
    mapmsg = rospy.wait_for_message('/map', OccupancyGrid, 10.0)

    # Extract the info.
    global map2grid, grid2map, w, h, res
    map2grid = PlanarTransform.fromPose(mapmsg.info.origin)
    grid2map = map2grid.inv()
    w        = mapmsg.info.width
    h        = mapmsg.info.height
    res      = mapmsg.info.resolution

    # Report
    rospy.loginfo("Basic map: %dx%d = %5.3fm x %5.3fm (res = %4.3fm)"
                  % (w, h, res * w, res * h, res))
    rospy.loginfo("Map grid bottom/left at %s" % map2grid)

    # Set up the initial obstacle map.  From the map message doc:
    # The is row-major order, starting with (0,0).  Occupancy
    # probabilities are in the range [0,100].  Unknown is -1.
    global obstacles
    obstacles = np.array(mapmsg.data, dtype=np.int8).reshape(h,w)

    # Leave 0 (presumably free) and 100 (unchangably obstacle).
    # Change -1 to 50 (mid way between the two).
    obstacles[np.where(obstacles < 0)] = 50


#
#   Publish the Obstacle Map
#
def callback_obsmap(event):
    # Create the obstacle map message.
    obsmapmsg = OccupancyGrid()
    obsmapmsg.header.stamp    = rospy.Time.now()
    obsmapmsg.header.frame_id = 'map'
    obsmapmsg.info.resolution = res
    obsmapmsg.info.width      = w
    obsmapmsg.info.height     = h
    obsmapmsg.info.origin     = map2grid.toPose()

    # Add the data, row-wise.
    obsmapmsg.data = obstacles.flatten().tolist()

    # Publish.
    obspub.publish(obsmapmsg)


#
#   Process the Laser Scan
#
def callback_scan(scanmsg):
    # For debugging report the message.
    #rospy.loginfo("Scan #%4d at %10d secs"
    #              % (scanmsg.header.seq, scanmsg.header.stamp.secs))

    # Extract the angles and ranges from the scan information.
    ranges = np.array(scanmsg.ranges)
    alphas = (scanmsg.angle_min +
              scanmsg.angle_increment * np.arange(len(ranges)))

    # Also grab the transform from the map to the laser scan data's
    # frame.  In particular, ask for the transform at the time the
    # scan data was captured - this makes things as self-consistent as
    # possible.  Give it up to 200ms, in case Linux has (temporarily)
    # swapped out the other threads.
    tfmsg = tfBuffer.lookup_transform('map', scanmsg.header.frame_id,
                                      scanmsg.header.stamp,
                                      rospy.Duration(0.200))
    map2laser = PlanarTransform.fromTransform(tfmsg.transform)

    #
    #   PROCESS THE MAP... For now we do nothing.
    #

    # For the fun of it, make the pixel at the map frame (0,0) "grey".
    # First map the (0,0) coordinates into the grid frame.
    xg = 0 - map2grid.x()
    yg = 0 - map2grid.y()
    # or more generally, especially if the map is rotated.
    (xg, yg) = grid2map.inParent(0, 0)
    # And then convert to u/v coordinataes.
    u = max(0, min(w-1, int(xg/res)))
    v = max(0, min(h-1, int(yg/res)))
    # And set the map value.
    global obstacles
    obstacles[v][u] = 50


######################################################################
#   Main code
######################################################################
if __name__ == "__main__":
    # Initialize the ROS node.
    rospy.init_node('move')

    # Any ROS parameters to grab?
    # updatefrac = rospy.get_param('~update_fraction', 0.2)

    # Grab and prepare the map.
    prepareMap()


    # First create a TF2 listener.  This implicily fills a local
    # buffer, so we can always retrive the transforms we want.
    # Allow 0.25sec for the buffer to accumulate past transforms.
    tfBuffer = tf2_ros.Buffer()
    tflisten = tf2_ros.TransformListener(tfBuffer)
    rospy.sleep(0.25)

    # Create a publisher to send the velocity command and obstacle map
    # messages.  We pick a queue size of 1 for the velocity command,
    # as using older commands is pointless.  We do the same for the
    # obstacle map, though latch so you can also see the last
    # published.
    cmdpub = rospy.Publisher('/vel_cmd', Twist, queue_size=1)
    obspub = rospy.Publisher('/obstacles',
                             OccupancyGrid, queue_size=1, latch=True)

    # Then create subscribers to listen to the actual pose, scan, and
    # goal pose.  We'll grab the initial map once (not subscribe).
    # For the scans (potentially incuring processing time) set the
    # queue size to 1.  This means the incoming queue will not keep
    # more than the newest message, and drop/skip unanswered messages.
    # Also the buffer size (defaulting to 64KB) must be big enough
    # (>20 msgs) to accomodate any bursts that might arrive.  Else
    # messages might be buffered upstream, defeating the local queue
    # limit and not getting dropped, causing things to lag behind.
    rospy.Subscriber('/move_base_simple/goal', PoseStamped, callback_goal)
    rospy.Subscriber('/pose', PoseStamped, callback_pose)
    rospy.Subscriber('/scan', LaserScan,   callback_scan, queue_size=1)

    # And finally, set up a timer to force publication of the obstacle
    # map, for us to view in RVIZ.
    timer = rospy.Timer(rospy.Duration(TPUBMAP), callback_obsmap)

    
    # Report and spin (waiting while callbacks do their work).
    rospy.loginfo("Move node spinning...")
    rospy.spin()
    rospy.loginfo("Move node stopped.")