#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# osim-extract.py -- Extract information from .osim file
#
# Copyright (C) 2013 Tobias Klauser <tklauser@distanz.ch>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.

import os
import sys
import getopt
import xml.etree.ElementTree as ET

def usage():
    print("""usage: {} [OPTION...] FILE

options:
    -b  print only body names
    -j  print only joint names
    -m  print only muscle names
    -v  verbose mode
    -h  show this help and exit""".format(os.path.basename(sys.argv[0])))

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "bjmvh")
    except getopt.GetoptError as err:
        print(str(err))
        usage()
        sys.exit(-1)

    do_print = { 'b': False, 'j': False, 'm': False }
    verbose = False
    for o, a in opts:
        if o == '-b':
            do_print['b'] = True;
        elif o == '-j':
            do_print['j'] = True;
        elif o == '-m':
            do_print['m'] = True;
        elif o == '-v':
            verbose = True
        elif o == '-h':
            usage()
            sys.exit(0)
        else:
            assert False, "unhandled option"

    if len(args) < 1:
        print("Error: no .osim model file specified")
        usage()
        sys.exit(-1)

    if not True in do_print.values():
        for k in do_print.iterkeys():
            do_print[k] = True

    try:
        tree = ET.parse(args[0])
    except:
        print("Error: failed to parse .osim model file")
        sys.exit(-1)

    root = tree.getroot()
    if verbose:
        print("Found root with version {}".format(root.attrib['Version']))
    model = tree.find('Model')
    if verbose:
        print("Found model '{}'".format(model.attrib['name']))

    if do_print['b']:
        bodies = [ body.attrib['name'] for body in root.findall('Model/BodySet/objects/Body') ]
        if verbose:
            print("bodies:")
        print(','.join(bodies))
    if do_print['j']:
        joints = [ joint.attrib['name'] for joint in root.findall('Model/BodySet/objects/Body/Joint/*') ]
        if verbose:
            print("joints:")
        print(','.join(joints))
    if do_print['m']:
        muscles = [ force.attrib['name'] for force in root.findall('Model/ForceSet/objects/Thelen2003Muscle') ]
        if verbose:
            print("muscles:")
        print(','.join(muscles))

if __name__ == '__main__':
    main()