""" tools for generating launch files for ROS """
import yaml

from lxml import etree
from lib.util import convert_ros, iterate_chain


def generate_leg_graph_launch(key, leg):
    """ generates leg graph generators launch files for schpin_stg package """

    cfg = {
        "chain_root": leg['chain_root'],
        "segments": dict(),
        "pos": leg['pos'],
        "rot": leg['rot']
    }

    for seg_name, seg_cfg in leg['segments'].items():

        seg_cfg = convert_ros(seg_cfg)

        cfg['segments'][seg_name] = {
            'mass': seg_cfg['mass'],
            'pos': seg_cfg['pos'],
            'mass_center': seg_cfg['mass_center'],
            'rot': seg_cfg['rot'],
        }

        if 'joint' in seg_cfg:
            cfg['segments'][seg_name]['joint'] = {
                'graph': {
                    'steps': seg_cfg['joint']['graph']['steps']
                },
                'angle': {
                    'min': seg_cfg['joint']['angle']['min'],
                    'max': seg_cfg['joint']['angle']['max']
                },
                'max_velocity': seg_cfg['joint']['max_velocity'],
                'max_torque': seg_cfg['joint']['max_torque'],
                'stall_current': seg_cfg['joint']['stall_current'],
                'running_current': seg_cfg['joint']['running_current'],
                'axis': seg_cfg['joint']['axis']
            }

        if 'next_segment' in seg_cfg:
            cfg['segments'][seg_name]['next_segment'] = seg_cfg['next_segment']

    launch = etree.Element('launch')

    etree.SubElement(
        launch,
        'node',
        pkg='schpin_stg',
        type='leg_graph_generator',
        name='{}_leg_graph_generator'.format(key)
    )
    rosparam = etree.SubElement(launch, 'rosparam')
    rosparam.text = yaml.dump(cfg, default_flow_style=False)

    return launch


def generate_robot_description():
    """ genereates launch file for setting robot description from this package """

    launch = etree.Element("launch")

    etree.SubElement(
        launch,
        "param",
        name="robot_description",
        textfile="$(find schpin_tote)/urdf/tote.urdf"
    )

    return launch

def master_conditional_node(*args, **kwargs):
    kwargs['if'] = "$(eval env('ROBOT_ADDRESS') == 'localhost')"
    etree.SubElement(
            *args,
            **kwargs
            )
    kwargs['if'] = "$(eval env('ROBOT_ADDRESS') != 'localhost')"
    etree.SubElement(
            *args,
            machine='master',
            **kwargs
            )

def generate_pwm_joints_node(xml_root, name, segments_map):

    joints = etree.SubElement(
        xml_root,
        "node",
        pkg="schpin_control",
        type="pwm_joints",
        name=name,
        respawn="true",
    )

    controllers = {
        "joint_state_controller": {
            "type": "joint_state_controller/JointStateController",
            "publish_rate": 50
        },
    }

    joints_list = []

    for chain in segments_map:
        for segment_name, segment in chain['segments'].items():

            if 'joint' not in segment:
                continue

            urdf_name = chain['name_prefix'] + segment_name

            joints_list.append(urdf_name)

            params = {
                'angle_min': segment['joint']['angle']['min'],
                'angle_center': segment['joint']['angle']['center'],
                'angle_max': segment['joint']['angle']['max'],
                'pulse_min': segment['joint']['pulse']['min'],
                'pulse_center': segment['joint']['pulse']['center'],
                'pulse_max': segment['joint']['pulse']['max'],
                "velocity": segment['joint']['max_velocity'],
                "update_freq": 50,  # TODO: hardcoded value is far from ideal
                "reverse": segment['joint']['reverse'],
                'channel':
                chain['channel_offset'] + segment['joint']['channel'],
                'default_pos': segment['joint']['default_pos'],
            }

            param_element = etree.SubElement(joints, "rosparam", ns=urdf_name)
            param_element.text = yaml.dump(
                convert_ros(params), default_flow_style=False
            )

            controllers[urdf_name + "_joint_position_controller"] = {
                "type": "position_controllers/JointPositionController",
                "joint": urdf_name
            }

    jointslist_param = etree.SubElement(joints, "rosparam")
    jointslist_param.text = yaml.dump(
        {
            "joints": joints_list
        }, default_flow_style=False
    )

    controllers_element = etree.SubElement(joints, "rosparam")
    controllers_element.text = yaml.dump(controllers, default_flow_style=False)

    etree.SubElement(
        xml_root,
        "node",
        name="controller_spawner",
        pkg="controller_manager",
        type="spawner",
        respawn="false",
        output="screen",
        ns=name,
        args=" ".join(controllers.keys()),
    )


def generate_state_publishers(element, joint_source_list):

    etree.SubElement(
        element,
        "node",
        name="robot_state_publisher",
        pkg="robot_state_publisher",
        type="robot_state_publisher",
    )
    js_pub = etree.SubElement(
        element,
        "node",
        name="joint_state_publisher",
        pkg="joint_state_publisher",
        type="joint_state_publisher",
    )

    js_params = etree.SubElement(js_pub, "rosparam")
    js_params.text = yaml.dump(
        {
            "source_list": joint_source_list,
            "rate": 50
        },
        default_flow_style=False
    )


def generate_machines(element):
    master = etree.SubElement(
        element, "machine", name="master", address="$(env MASTER_ADDRESS)",
    )
    master.set('if', "$(eval env('ROBOT_ADDRESS') != 'localhost')")

    robot = etree.SubElement(
        element,
        "machine",
        name="robot",
        address="$(env ROBOT_ADDRESS)",
        default="true",
        user="$(env ROBOT_USER)",
    )
    robot.set('if',"$(eval env('ROBOT_ADDRESS') != 'localhost')")
    robot.set("env-loader", "~/setup.bash")


def generate_control_layer(cfg):
    """ generate control layer of schpin system, meaning all recuired devices AND
        robot_state_publisher/joint_state_publisher combo """

    launch = etree.Element('launch')

    generate_state_publishers(
        launch,
        ["/control/joint_states"]
    )

    segments_map = list()

    for key, l_cfg in cfg['legs'].items():

        segments_map.append(
            {
                "channel_offset": l_cfg["channel_offset"],
                "segments": l_cfg['segments'],
                "name_prefix": "leg/" + key + "/"
            }
        )

    segments_map.append(
        {
            "channel_offset": cfg['gimbal']['channel_offset'],
            "segments": cfg['gimbal']['segments'],
            "name_prefix": "gimbal/"
        }
    )

    generate_pwm_joints_node(
        xml_root=launch, name='control', segments_map=segments_map
    )

    return launch


def generate_general(cfg):

    launch = etree.Element('launch')

    generate_machines(launch)

    etree.SubElement(
        launch,
        "include",
        file="$(find schpin_tote)/launch/robot_description.launch"
    )

    return launch


def generate_user_layer(cfg):

    launch = etree.Element('launch')

    master_conditional_node(
        launch,
        "node",
        name="$(anon rviz)",
        pkg="rviz",
        type="rviz"
    )
    return launch


def generate_main(cfg):
    launch = etree.Element('launch')

    to_include = [
        "general.launch", "control_layer.launch", "user_layer.launch"
    ]

    for filename in to_include:
        etree.SubElement(
            launch, "include", file="$(find schpin_tote)/launch/" + filename
        )

    return launch

def generate_debug_joints(cfg):
    launch = etree.Element('launch')

    etree.SubElement( launch, "node", name="joint_marker",  pkg="schpin_user", type="joint_markers.py")

    return launch

def generate_launch_files(cfg):
    """ generates all required launch files """

    files = {
        "robot_description.launch": generate_robot_description(),
        "control_layer.launch": generate_control_layer(cfg),
        "user_layer.launch": generate_user_layer(cfg),
        "general.launch": generate_general(cfg),
        "debug/joints.launch": generate_debug_joints(cfg),
        "main.launch": generate_main(cfg),
    }

    for key in cfg['legs'].keys():
        files["{}_leg_graph_generator.launch".format(key)
              ] = generate_leg_graph_launch(key, cfg['legs'][key])

    for filename, element in files.items():
        tree = etree.ElementTree(element)

        tree.write("../launch/" + filename, pretty_print=True)
