#!/usr/bin/env python

'''
-------------------
Copyright (c) 2015 Computational Biomechanics (CoBi) Core, Department of
Biomedical Engineering, Cleveland Clinic

Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to permit
persons to whom the Software is furnished to do so, subject to the
following conditions:

The above copyright notice and this permission notice shall be included
in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
USE OR OTHER DEALINGS IN THE SOFTWARE.
-------------------

prescribe_motion_4_FEBio.py

DESCRIPTION:
Python script to read experimental kinetics/kinematics data and apply the
desired curves to an FEBio model.

REQUIREMENTS:
Python (http://www.python.org)
SciPy (http://www.scipy.org)
NumPy (http://www.numpy.org)

WRITTEN BY:
Craig Bennetts, MS
Computational Biomodeling (CoBi) Core
Department of Biomedical Engineering
Lerner Research Institute
Cleveland Clinic
Cleveland, OH
bennect2@ccf.org


LEGEND:
-------
<.> = variable argument

USAGE:
------
./SLURM_FEBio_BATCH.py <FEBio_model.feb> <FEBio_parameters.cfg> <SLURM_FEBio_submit.sh>


INPUT FILES:
------------

<input>.cfg:

	DESCRIPTION:
	This is a user-defined input configuration file to specify a set of input
	parameters to automatically run a set of FEBio simulations on the 
	departments high performace computing (HPC) cluster.

	LEGEND:
	-------
	[#] = text file line number
	t = line triplet number, positive integer >= 0

	FILE FORMAT (EXAMPLE):
	----------------------
	[ 3*t ] <FEBio_element_Xpath,string>
	[3*t+1] <parameter_type,string>
	[3*t+2] <VALUE,float>, <MIN,float>:<STEP_SIZE,float>:<MAX.inclusive,float>, ...

	  -or-

	[3*t+2] (<VALUE,float>, ...), (<VALUE,float>, ...), ...

	FORMAT EXAMPLES:
	----------------
	<FEbio_element_Xpath,string> example:
		Step[@name='Step01']/Constraints/rigid_body/prescribed[@bc='Rx']
		in general: Element_name[@attribute='string']/...
		NOTE: Xpath XML input is only available in Python 2.7<
		      on the Lerner HPC, load Python 2.7.8 before running this script:
			     module load python/2.7.8

	<parameter_type,string> options:
		absolute = replaces existing value
		factor   = multiplies by existing value
		relative = adds to existing value

<input_FEBio_model>.feb:

	DESCRIPTION:
	Default FEBio model file which has pre-existing load curves and associated
	kinematic/kinetic constraints (so it is capable of running prior to
	the modifications applied in this script).


OUTPUT FILE:
------------

<input_FEBio_model>_MOTION.feb

	DESCRIPTION:
	Modified FEBio model, with specified kinematic/kinetic DOF data curves applied
 
'''


import sys
from re import sub
from numpy import arange,unique,ndarray
from itertools import product
from subprocess import Popen
#from xml.etree.ElementTree import ElementTree, parse, SubElement, dump, tostring
try:
    from xml.etree.cElementTree import ElementTree  # preferrably load cElementTree is much faster than ElementTree
except ImportError:
    from xml.etree.ElementTree import ElementTree


def USAGE(argv):

	print
	print 'USAGE: ' + argv[0] + ' <FEBio_model.feb> <FEBio_parameters.cfg> <SLURM_FEBio_submit.sh>'
	print
	print '        <FEBio_parameters.cfg>  = FEBio parameters configuration file (see script header for format)'
	print '        <SLURM_FEBio_submit.sh> = bash script to submit FEBio job to SLURM scheduler'
	print


# Function to reformat XML ElementTree
def indent_elements(elem, indent_string, level=0):
    i = "\n" + level*indent_string
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + indent_string
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            indent_elements(elem, indent_string, level+1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i


def SLURM_FEBio_BATCH(argv):

	# CHECK SCRIPT USAGE
	if len(argv)!=4:
		USAGE(argv)
		sys.exit(1)

	#--------------------------#
	# INPUT CONFIGURATION FILE #
	#--------------------------#

	# Set FEBio filename
	input_FEBio_filename = argv[1]	

	# Set configuration filename
	config_filename = argv[2]

	# Set SLURM 
	input_submit_filename = argv[3]

	# SET DEFAULTS
	VALID_PARAMETER_TYPES = ['absolute','factor','relative']

	# READ CONFIGURATION FILE
	ConfigFile = open(config_filename,'r')
	ConfigFileLines = ConfigFile.readlines()
	ConfigFile.close()


	#------------------------------#
	# CREATE FEBio PARAMETER LISTS #
	#------------------------------#
	# WARNING: no configuration input file format checking here!
	# ASSUMPTION:
	# triplets of lines as follows

	FEBio_parameter_tags = []
	FEBio_parameter_types = []
	FEBio_parameter_lists = []
	parameter_singletons = []
	

	line_index = 0

	# CONFIG INPUT LOOP
	while line_index<len(ConfigFileLines):

		# TODO: may want to check for other OS text file carriage returns
		if ConfigFileLines[line_index]!='\n':

			# get the next triplet of lines from the configuration file
			FEBio_tag = ConfigFileLines[line_index].strip('\r\n')  # strips carriage return for both linux and Windows input configuration files
			line_index += 1
	
			FEBio_parameter_tags.append( FEBio_tag )
	
			value_type = ConfigFileLines[line_index].strip('\r\n')
			line_index += 1
	
			if not value_type in VALID_PARAMETER_TYPES:
				print
				print 'ERROR: invalid parameter type specified in configuration file!'
				print '       allowed types: absolute, factor, relative'
				print
				sys.exit(1)
	
			FEBio_parameter_types.append( value_type )
	
			item_line = ConfigFileLines[line_index].strip('\r\n')
			line_index += 1
	
			# Create a list of parameter values from specified values and/or ranges, or tuples.
			# FORMAT:
			# - comma separated:
			#   - values, <float>
			#   - ranges, <float>:<float>:<float> = MIN.inclusive : STEP_SIZE : MAX.inclusive
			#     -or-
			#   - tuples, (<float>,...) 
			# - order doesn't matter
			# - duplicates will be removed
			#
			# FORMAT EXAMPLES:
			# 1,2,4:0.5:6  # mixture of values and ranges
			# (1,2),(3,4)  # must all be the same type of tuples with same # of parameters (#>1), no ranges!
	
			parameter_list = []
			
			# remove spaces if present
			item_line = sub(' ','',item_line)		
	
			# check to see if tuples were specified
			if ('(' in item_line) and (')' in item_line):
				item_line = item_line.lstrip('(').rstrip(')')
				item_list = item_line.split('),(')
			else:
				item_list = item_line.split(',')
			
			for item in item_list:
	
				# tuple
				if ',' in item:
					subitem_list = item.split(',')

					subitem_range_list = []
					for subitem in subitem_list:
						# tuple value range
						if ':' in subitem:
							range_args = map(float,subitem.split(':'))
							if len(range_args)!=3:
								print
								print 'ERROR: parameter range format must be MIN:STEP_SIZE:MAX!'
								print
								sys.exit(1)
							else:
								range_list = list( arange(range_args[0],range_args[2]+range_args[1],range_args[1]) )
								# correct length of range list if numerical precision throws this off
								range_len = int( round( (range_args[2]-range_args[0])/range_args[1] ) ) + 1
								subitem_range_list.append( range_list[:range_len] )
						# tuple value
						else:
							value = float(subitem)
							subitem_range_list.append( [value] )
					
					#print subitem_range_list					
					
					#subitem_range_permutations = map(list,list(product(*subitem_range_list)))
					subitem_range_permutations = list(product(*subitem_range_list))
					parameter_list += subitem_range_permutations

					#value_list = []
					#for value in subitem_list:
					#	value_list.append( float(value) )
					#parameter_list.append( value_list )
				# range
				elif ':' in item:
					range_args = map(float,item.split(':'))
					if len(range_args)!=3:
						print
						print 'ERROR: parameter range format must be MIN:STEP_SIZE:MAX!'
						print
						sys.exit(1)
					else:
						range_list = list( arange(range_args[0],range_args[2]+range_args[1],range_args[1]) )
						# correct length of range list if numerical precision throws this off
						range_len = int( round( (range_args[2]-range_args[0])/range_args[1] ) ) + 1
						#subitem_range_list.append( range_list[:range_len] )
						parameter_list += range_list[:range_len]
						#parameter_list += list( arange(range_list[0],range_list[2]+range_list[1],range_list[1]) )
				# single value
				else:
					value = float(item)
					parameter_list.append( value )

			#print parameter_list

			# create a sorted and unique subset of parameter values
			parameter_list = list(unique(parameter_list))

			if len(parameter_list)==1:
				parameter_singletons.append( True )
			else:
				parameter_singletons.append( False )
	
			FEBio_parameter_lists.append( parameter_list )

		else:
			line_index += 1

	# END OF CONFIG INPUT LOOP

	# NOTE: FEBio_parameter_lists contains floating values and/or arrays of floating values


	# WRITE SIMULATION INDEX KEY

	# Create a new filename (related to the config file, .cfg) with an explicit
	# list of sorted parameter values to match simulation indices with parameter
	# values in case the user specified ranges.
	sim_key_filename = config_filename[:config_filename.rfind('.')] + '.key'

	SimKeyFile = open(sim_key_filename,'w')
	for parameter_index in range(len(FEBio_parameter_tags)):
		SimKeyFile.write(FEBio_parameter_tags[parameter_index]+'\n')
		SimKeyFile.write(FEBio_parameter_types[parameter_index]+'\n')
		for item_index in range(len(FEBio_parameter_lists[parameter_index])):
			if type(FEBio_parameter_lists[parameter_index][item_index]) is ndarray:
				SimKeyFile.write( '(' + ','.join( map(str, FEBio_parameter_lists[parameter_index][item_index] ) ) + ')' )
			else:
				SimKeyFile.write( str( FEBio_parameter_lists[parameter_index][item_index] ) )
			if item_index<len(FEBio_parameter_lists[parameter_index])-1:
				SimKeyFile.write(',')
		SimKeyFile.write('\n')
	SimKeyFile.close()


	#----------------------------------#
	# DETERMINE PARAMETER PERMUTATIONS #
	#----------------------------------#

	# determine permutations of all values for all parameters
	parameter_permutations = map(list,list(product(*FEBio_parameter_lists)))
	
	# Create permutations of parameter indices
	# corresponds to order of values in parameter_permutations list
	parameter_index_lists = []
	for parameter_list in FEBio_parameter_lists:
		parameter_index_lists.append( range(1,len(parameter_list)+1) )

	index_permutations = map(list,list(product(*parameter_index_lists)))

	#print parameter_permutations
	#print index_permutations


	#---------------------------------------------------#
	# CREATE SIMULATIONS FOR EACH PARAMETER PERMUTATION #
	#---------------------------------------------------#

	# Input SLURM FEBio sumbission template
	SubmitFile = open(input_submit_filename,'r')
	SubmitFileLines = SubmitFile.readlines()
	SubmitFile.close()	

	# Set default FEBio XML output file option
	OUTPUT_XML_INDENT_STRING = "\t"

	# SIMULATION LOOP
	for (indices,parameters) in zip(index_permutations,parameter_permutations):
		
		# Create new directory for current set of simulation parameters
		# use parameter indices 		
		sim_dir_name = 'sim'
		for i in range(len(indices)):
			if not parameter_singletons[i]:
				sim_dir_name += '_' + str(indices[i])

		p = Popen(['mkdir',sim_dir_name])

		#-------------------#
		# UPDATE FEBio FILE #
		#-------------------#
	
		# Input FEBio XML datastructure from input .feb file
		# NOTE: could do this outside the loop and create a deepcopy
		#       but it seems faster inputing from file each time than copying
		tree = ElementTree(file=input_FEBio_filename)
		root = tree.getroot()

		#print root.tag
		#print root.attrib

		# Modify specified FEBio parameters for the current parameter permutation

		for tag_index in range(len(FEBio_parameter_tags)):
			
			# obtain reference to LoadData element
			element_list = root.findall(FEBio_parameter_tags[tag_index])  # WARNING: assumes only one FEBio_tag in input FEBio file

			for element in element_list:
			
				old_text = element.text
				# try splitting by comma to see if this is a multi-valued parameter
				old_text_list = old_text.split(',')

				if type(parameters[tag_index]) is ndarray:		
					if len(old_text_list)!=len(parameters[tag_index]):
						print
						print 'ERROR: mismatched number of existing and specified number of parameter!'
						print
						sys.exit(1)
				elif len(old_text_list)!=1:
					print
					print 'ERROR: mismatched number of existing and specified number of parameter!'
					print
					sys.exit(1)
	
	
				# NOTE: valid types are checked earlier upon input
				if FEBio_parameter_types[tag_index]=='absolute':
					if type(parameters[tag_index]) is ndarray:
						element.text = ','.join( map(str,parameters[tag_index]) )
					else:
						element.text = str( parameters[tag_index] )
				elif FEBio_parameter_types[tag_index]=='factor':  # multiply relative value by current value
					if type(parameters[tag_index]) is ndarray:
						old_value_list = map(float,old_text_list)
						new_text_list = []
						for (factor,old_value) in zip(parameters[tag_index],old_value_list):
							new_text_list += str( factor*old_value )
						element.text = ','.join( new_text_list )
					else:
						element.text = str( parameters[tag_index]*float(element.text) )
				elif FEBio_parameter_types[tag_index]=='relative':
					if type(parameters[tag_index]) is ndarray:
						old_value_list = map(float,old_text_list)
						new_text_list = []
						for (relative_value,old_value) in zip(parameters[tag_index],old_value_list):
							new_text_list += str( old_value + relative_value )
						element.text = ','.join( new_text_list )
					else:
						element.text = str( float(element.text) + parameters[tag_index] )				
				#else:


		#---------------------------#
		# WRITE MODIFIED FEBio FILE #
		#---------------------------#

		# Create new filename for current set of parameter indices
		#output_FEBio_filename = input_FEBio_filename[:input_FEBio_filename.rfind('.')] + '_' + sim_dir_name + '.feb'
		output_FEBio_filename = input_FEBio_filename
		
		# reformat FEBio XML tree to output hierarchical levels using TABS
		indent_elements(root,OUTPUT_XML_INDENT_STRING)
		
		# TODO?: get input XML file encoding to properly rewrite output file?
		tree.write(sim_dir_name + '/' + output_FEBio_filename, xml_declaration=True, encoding='ISO-8859-1', method="xml")
				

		#---------------------------------#
		# CREATE/MODIFY SUBMISSION SCRIPT #
		#---------------------------------#

		#output_submit_filename = input_submit_filename[:input_submit_filename.rfind('.')] + '_' + sim_dir_name + '.sh'
		output_submit_filename = input_submit_filename

		OutputSubmitFile = open(sim_dir_name + '/' + output_submit_filename,'w')

		for line_index in range(len(SubmitFileLines)):
			
			line = SubmitFileLines[line_index]
			
			if line_index==len(SubmitFileLines)-1:
				new_line = line[:line.rfind('-i')+2] + ' ' + output_FEBio_filename + '\n'
				OutputSubmitFile.write(new_line)
			elif 'job-name' in line:
				new_line = line[:line.rfind('=')+1] + sim_dir_name + '\n'
				OutputSubmitFile.write(new_line)
			elif 'error' in line:
				#new_line = line[:line.rfind('=')+1] + sim_dir_name + '.job_%J.err\n'
				new_line = line[:line.rfind('=')+1] + sim_dir_name + '.err\n'
				OutputSubmitFile.write(new_line)
			elif 'output' in line:
				#new_line = line[:line.rfind('=')+1] + sim_dir_name + '.job_%J.out\n'
				new_line = line[:line.rfind('=')+1] + sim_dir_name + '.out\n'
				OutputSubmitFile.write(new_line)
			else:
				OutputSubmitFile.write(line)

		OutputSubmitFile.close()

		#-------------------------------#
		# SUBMIT SLURM FEBio SIMULATION #
		#-------------------------------#

		print 'sbatch ' + sim_dir_name + '/' + output_submit_filename

		#p = Popen(['sbatch',output_submit_filename],cwd=sim_dir_name)

	# END OF SIMULATION LOOP

# END OF SLURM_FEBio_BATCH()


# Default calls SLURM_FEBio_BATCH() if this script is not called as a function
if __name__ == "__main__":
	SLURM_FEBio_BATCH(sys.argv)