#!/usr/bin/env python3

# Imports
import argparse
import os
import sys
from colorama import Style, Fore

# ROS imports
import rospy
from interactive_markers.interactive_marker_server import InteractiveMarkerServer
from interactive_markers.menu_handler import MenuHandler
from visualization_msgs.msg import InteractiveMarkerControl, Marker, InteractiveMarker

# ATOM imports
import atom_core.config_io
import atom_core.ros_utils
import atom_msgs.srv
from atom_calibration.collect.data_collector import DataCollector
from atom_core.system import resolvePath

server = None
menu_handler = MenuHandler()


def menuFeedback(feedback):
    handle = feedback.menu_entry_id
    if handle == 1:  # collect snapshot
        print('Save collection selected ')
        data_collector.saveCollection()
    elif handle == 2:
        print('Automatic data collection selected ')
        rospy.wait_for_service('/automatic_data_collector/toggle_automatic_data_collector')
        toggle_automatic_data_collector = rospy.ServiceProxy(
            '/automatic_data_collector/toggle_automatic_data_collector', atom_msgs.srv.ToggleAutomaticDataCollector)
        _ = toggle_automatic_data_collector()


def initMenu():
    menu_handler.insert("Save collection", callback=menuFeedback)
    menu_handler.insert("Automatic data collection", callback=menuFeedback)


def createInteractiveMarker(world_link):
    marker = InteractiveMarker()
    marker.header.frame_id = world_link
    trans = (1, 0, 1)
    marker.pose.position.x = trans[0]
    marker.pose.position.y = trans[1]
    marker.pose.position.z = trans[2]
    quat = (0, 0, 0, 1)
    marker.pose.orientation.x = quat[0]
    marker.pose.orientation.y = quat[1]
    marker.pose.orientation.z = quat[2]
    marker.pose.orientation.w = quat[3]
    marker.scale = 0.2

    marker.name = 'menu'
    marker.description = 'menu'

    # insert a box
    control = InteractiveMarkerControl()
    control.always_visible = True

    marker_box = Marker()
    marker_box.type = Marker.SPHERE
    marker_box.scale.x = marker.scale * 0.7
    marker_box.scale.y = marker.scale * 0.7
    marker_box.scale.z = marker.scale * 0.7
    marker_box.color.r = 0
    marker_box.color.g = 1
    marker_box.color.b = 0
    marker_box.color.a = 0.2

    control.markers.append(marker_box)
    marker.controls.append(control)

    marker.controls[0].interaction_mode = InteractiveMarkerControl.MOVE_3D

    control = InteractiveMarkerControl()
    control.orientation.w = 1
    control.orientation.x = 1
    control.orientation.y = 0
    control.orientation.z = 0
    control.name = "move_x"
    control.interaction_mode = InteractiveMarkerControl.MOVE_AXIS
    control.orientation_mode = InteractiveMarkerControl.FIXED
    marker.controls.append(control)

    control = InteractiveMarkerControl()
    control.orientation.w = 1
    control.orientation.x = 0
    control.orientation.y = 1
    control.orientation.z = 0
    control.name = "move_z"
    control.interaction_mode = InteractiveMarkerControl.MOVE_AXIS
    control.orientation_mode = InteractiveMarkerControl.FIXED
    marker.controls.append(control)

    control = InteractiveMarkerControl()
    control.orientation.w = 1
    control.orientation.x = 0
    control.orientation.y = 0
    control.orientation.z = 1
    control.name = "move_y"
    control.interaction_mode = InteractiveMarkerControl.MOVE_AXIS
    control.orientation_mode = InteractiveMarkerControl.FIXED
    marker.controls.append(control)

    server.insert(marker, markerFeedback)
    menu_handler.apply(server, marker.name)


def markerFeedback(feedback):
    pass
    # print('Received feedback')


if __name__ == "__main__":
    # Parse command line arguments
    ap = argparse.ArgumentParser()
    ap.add_argument("-s", "--marker_size", help='Scale of the interactive markers.', type=float, default=0.5)
    ap.add_argument('-o', '--output_folder', help='Output folder to where the collected data will be stored.', type=str,
                    required=True)
    ap.add_argument("-c", "--calibration_file", help='full path to calibration file.', type=str, required=True)
    ap.add_argument("-ov", "--overwrite", help="Overwrite dataset folder if it exists.", action='store_true',
                    default=False)
    ap.add_argument("-ssl", "--skip_sensor_labeling", default=None, type=lambda s: eval(s, globals()),
                    help='A string to be evaluated into a lambda function that receives a sensor name as input and '
                         'returns True or False to indicate if the sensor should be labelled.'
                         ' The Syntax is lambda name: f(x), where f(x) is the function in python '
                         'language. Example: lambda name: name in ["lidar1", "camera2"] , to avoid labeling of sensors'
                         ' lidar1 and camera2.')

    args = vars(ap.parse_args(args=atom_core.ros_utils.filterLaunchArguments(sys.argv)))

    output_folder = resolvePath(args['output_folder'])

    if output_folder == os.getenv("HOME"):
        print('\nError: Dataset ' + Fore.RED + output_folder + Style.RESET_ALL +
              ' is the same as $HOME path. Are you sure you know what your are doing?\n' + Style.RESET_ALL)
        rospy.signal_shutdown()

    # Initialize ROS stuff
    rospy.init_node("collect_data")
    # rospack = rospkg.RosPack()  # get an instance of RosPack with the default search paths
    server = InteractiveMarkerServer("data_labeler")
    robot_description = rospy.get_param('/robot_description')
    # rospy.sleep(0.5)

    # Process robot description and create an instance of class Sensor for each sensor
    data_collector = DataCollector(args, server, menu_handler)

    createInteractiveMarker(data_collector.world_link)
    initMenu()
    menu_handler.reApply(server)
    server.applyChanges()
    print('Changes applied ...')

    rospy.spin()