#! /usr/bin/env python
import sys
import xml.etree.ElementTree as ET

def edit_world(file_name, poses):
    """Edit poses in a world file.
    
    Arguments:
    file_name (str): world file path
    poses (dict): pose of the objects {object_name: pose}
    """

    # Open world file
    tree = ET.parse(file_name)

    root = tree.getroot()
    root_state = root.findall('world')[0].findall('state')[0]

    objects_pose = poses
    models = {} 

    for model in root_state.iter('model'):
        if model.attrib['name'] in objects_pose:
            models[model.attrib['name']] = model

    # Edit pose elements
    for model_name, element in models.iteritems():
        element.findall('pose')[0].text = objects_pose[model_name]
    
    # Write world file
    tree.write(file_name)

if __name__ == '__main__':
    """Edit gazebo world file. """
    args_number = len(sys.argv) - 2

    package_path = sys.argv[1]
    
    poses = {}

    world_name = sys.argv[2]

    poses['my_box'] = sys.argv[3]

    for i in range(4, args_number):
        if i == 4:
            poses['iRobot'] = sys.argv[i]
        else:
            poses['iRobot_' + str(i-5)] = sys.argv[i]


    edit_world(package_path + '/worlds/' + world_name + '.world', poses)
