#!/usr/bin/env python

from bwi_tools import loadMapFromFile, resolveResource
from geometry_msgs.msg import PoseWithCovarianceStamped
from multi_level_map_msgs.msg import LevelMetaData, MultiLevelMapData
from multi_level_map_msgs.srv import ChangeCurrentLevel, ChangeCurrentLevelRequest
import multi_level_map_utils.utils as utils
from nav_msgs.msg import MapMetaData, OccupancyGrid
from nav_msgs.srv import GetMap, SetMap
import rospy
from std_srvs.srv import Empty
import tf

class LevelMultiplexer:

    def __init__(self):

        rospy.init_node("map_mux")

        self.global_frame_id = rospy.get_param('~global_frame_id', 'level_mux/map')
        self.default_current_level = rospy.get_param('~default_current_level', None)

        # Subcribe to the multi level map data to get information about all the maps.
        self.multimap_subscriber = rospy.Subscriber("map_metadata", MultiLevelMapData, self.process_multimap)
        self.multimap_available = False
        multimap_topic_name = rospy.resolve_name("map_metadata")
        if '/' in multimap_topic_name:
            self.multimap_topic_prefix = multimap_topic_name[0:multimap_topic_name.rfind('/') + 1]
        else:
            self.multimap_topic_prefix = ""

        # We also need to publish information about the current level.
        self.current_level = None
        self.status_publisher = rospy.Publisher("~current_level", LevelMetaData, latch=True, queue_size=1)

        # Based on the current level id, publish map, static_map, map_metadata and initial_pose (i.e. the mux output).
        self.level_publisher = None
        self.level_metadata_publisher = None
        self.level_map_service = None
        self.initialpose_publisher = None

        # Finally, provide a service that allows the current level to be changed.
        # This service will get enabled once the multimap becomes available.
        self.change_level_service = None

    def process_multimap(self, msg):

        # Get information about all the levels.
        self.levels = msg.levels
        self.maps = {}
        for level in self.levels:
            resolved_map_file = resolveResource(level.map_file)
            if resolved_map_file is None:
                str_err = "Unable to resolve resource provided by multi map server: %s" %level.map_file
                rospy.logfatal(str_err)
                raise ValueError(str_err)
            self.maps[level.level_id] = loadMapFromFile(resolved_map_file)
            if self.maps[level.level_id] is None:
                str_err = "Unable to load yaml file for map: %s" %resolved_map_file
                rospy.logfatal(str_err)
                raise ValueError(str_err)
            self.maps[level.level_id].map.info.map_load_time = rospy.Time.now()
            self.maps[level.level_id].map.header.frame_id = utils.frameIdFromLevelId(level.level_id)

        # If the service to change the level hasn't been advertised, then advertise it now.
        if self.change_level_service is None:
            self.change_level_service = rospy.Service('~change_current_level', ChangeCurrentLevel, self.change_level)

        self.multimap_available = True

        # Perform some sanity checks on the current_level
        if len(self.levels) == 1:
            # If there is only one level, then change to that level without publishing initialpose.
            req = ChangeCurrentLevelRequest()
            req.new_level_id = self.levels[0].level_id
            req.publish_initial_pose = False
            self.change_level(req)
        elif self.current_level is not None:
            # This can only happen when the multimap changes from what it was previously. I'm not really sure if this
            # will ever be possible, but let's perform some sanity checks here to make sure we don't get screwed.
            current_level_found = False
            for level in self.levels:
                if self.current_level == level.level_id:
                    current_level_found = True
                    break
            if not current_level_found:
                # Of no! The multimap changed, and our current level is not longer available. Unregister everything!
                rospy.logerr("The multimap changed, and our selected current level " + self.current_level + " is no longer available. Navigation will probably not work properly!")
                if self.level_publisher is not None:
                    self.level_publisher.unregister()
                    self.level_publisher = None
                    self.level_metadata_publisher.unregister()
                    self.level_publisher = None
                    self.level_map_service.unregister()
                    self.level_map_service = None
                    self.initialpose_publisher.unregister()
                    self.initialpose_publisher = None
                self.current_level = None

        # Perform some sanity checks on the current_level
        if self.default_current_level is not None and self.current_level is None:
            default_level_found = False
            for level in self.levels:
                if self.default_current_level == level.level_id:
                    default_level_found = True
                    break
            if default_level_found:
                req = ChangeCurrentLevelRequest()
                req.new_level_id = self.default_current_level
                req.publish_initial_pose = False
                self.change_level(req)
            else:
                rospy.logwarn("Default level requested as " + self.default_current_level + ", but unavailable in multimap.")

    def change_level(self, req):
        success = True
        error_message = ""

        level_found = False
        available_level_names = []

        for level in self.levels:
            if level.level_id == req.new_level_id:

                # Make this the current level
                self.current_level = level.level_id

                # Initialize the mux output channels if they have not been initialized.
                if self.level_publisher is None:
                    self.level_publisher = rospy.Publisher("~map", OccupancyGrid, latch=True, queue_size=1)
                    self.level_metadata_publisher = rospy.Publisher("~map_metadata", MapMetaData, latch=True, queue_size=1)
                    self.level_map_service = rospy.Service('~static_map', GetMap, self.process_level_service)
                    # Note that initialpose is not latched, as the publication is only valid when the 
                    # change_current_level service is called. Soon after, the robot may have moved, and this estimate
                    # is incorrect.
                    self.initialpose_publisher = rospy.Publisher("initialpose", PoseWithCovarianceStamped, queue_size=1) 
                    self.set_map = rospy.ServiceProxy("set_map", SetMap)
                    try:
                        clear_costmap_service = rospy.ServiceProxy('move_base/clear_costmaps', Empty)
                        # Clear all costmaps.
                        clear_costmap_service()
                    except rospy.ROSInterruptException:
                        # shutting down.
                        return False, "Shutting down!"
                    except rospy.ServiceException:
                        # Service does not exist, level_mux being used without move_base OR a configuration snafu
                        rospy.logwarn("Couldn't find move_base clear costmap service. This is an error if navigation is running!")

                self.maps[self.current_level].map.header.stamp = rospy.Time.now()

                # Publish latched map messages
                self.level_metadata_publisher.publish(self.maps[self.current_level].map.info)
                self.level_publisher.publish(self.maps[self.current_level].map)

                if req.publish_initial_pose:
                    try: 
                        self.set_map(self.maps[self.current_level].map, req.initial_pose)
                    except rospy.ServiceException:
                        # Service does not exist, level_mux being used without move_base OR a configuration snafu
                        rospy.logwarn("MultiLevelMap - Multiplexer::changeLevel: " + 
                                      "Couldn't set initial pose via service call. Trying unlatched publication.")
                        self.initialpose_publisher.publish(req.initial_pose)
                self.status_publisher.publish(level)

                level_found = True
                break
            else:
                available_level_names.append(level.level_id)

        if not level_found:
            success = False
            error_message = "Level " + req.new_level_id + " not found in available levels " + str(available_level_names)

        return success, error_message

    def process_level_service(self, req):
        self.maps[self.current_level].map.header.stamp = rospy.Time.now()
        return self.maps[self.current_level]

    def spin(self):
        br = tf.TransformBroadcaster()
        r = rospy.Rate(10)
        while not rospy.is_shutdown():
            if self.current_level is not None:
                br.sendTransform((0, 0, 0),
                                 tf.transformations.quaternion_from_euler(0, 0, 0),
                                 rospy.Time.now(),
                                 self.global_frame_id,
                                 utils.frameIdFromLevelId(self.current_level))
            r.sleep()

if __name__ == '__main__':
    level_mux = LevelMultiplexer()
    try:
        level_mux.spin()
    except rospy.ROSInterruptException: 
        pass
