#!/usr/bin/env python

import os
import sys
import string
import time
from stl import mesh

# This file could use the library xml.etree.ElementTree but doesn't.  The xml is explicitly written so its ordered and pretty.


class FebioElement:
	def __init__(self, material):
		self.mat = material
		self.type = None
		self.connect = {}		# contains matched pairs of id# and element connectivity
		self.nodes=set()

	def __repr__(self):
		return "<FebioElement type:{} mat:{} connect:{} elements>".format(self.type, self.mat, len(self.connect))

nodeData={}
elemData={}
physGroups={}	# physGroups{PG#] = "name in mesh"
previousNodes = 0	# number of nodes in all previous meshes
contactGroups = {}

inputDirectory = '/home/landisb/workspace/IndentationTests/Leg/FeBio/BasicLeg'

os.chdir(inputDirectory)
for file in os.listdir(inputDirectory):
	if file.endswith('.msh'):
		gmshMapping = {'1': 'line', '2': 'tri3', '3': 'quad4', '4': 'tet4',
					   '5': 'hex8', '8': 'line3', '9': 'tri6', '11': 'tet10',
					   '12': 'hex27', '15': 'point', '16': 'quad8', '17': 'hex20'}
		meshFile = open(file, 'r')
		print 'Parsing Mesh {}'.format(file)
		pN=previousNodes
		#
		line = meshFile.readline()
		while line:
			if line.startswith('$Nodes') or line.startswith('$NOD'):
				#~ print 'Reading Nodes...'
				nodes = int( meshFile.readline() )
				for i in xrange(nodes):
					line=meshFile.readline().split()	#readline and split by whitespace
					nodeData[int(line[0])+pN]= tuple(float(i) for i in line[-3:])		# mesh node number + previousNodes
				previousNodes+=nodes
			elif line.upper().startswith('$EL'):
				#~ print 'Reading Elements...'
				if line.startswith('$Elements'):	pg = 3	# pg is the index of the physical group line
				elif line.startswith('$ELM'):		pg = 2
					#~ try:	elemData[line[2]].connect[int(line[0])] = [int(j)+pN for j in line[-nodeNum:]]
					#~ except KeyError:
						#~ f='"'+file.split('.')[0]
						#~ if nodeNum == 3:
							#~ f = f+'_surf_'+line[pg]
						#~ elemData[line[2]] = FebioElement(line[2], f+'"' )
						#~ elemData[line[2]].connect[int(line[0])] = [int(j)+pN for j in line[-nodeNum:]]
				elements = int( meshFile.readline() )
				for i in xrange(elements):
					line=meshFile.readline().split()	#readline and split by whitespace
					mapping = gmshMapping[line[1]]
					name = physGroups[line[pg]]
					#
					# #		Good spot for checking elemData is initialized
					#
					nodeNum = int(mapping.strip(string.ascii_letters))  # number of nodes in the element
					if   mapping.startswith('line'):	print 'Lines are currently ignored.'
					elif mapping.startswith('point'):	print 'Points are currently ignored.'
					elif (mapping.startswith('quad') or mapping.startswith('hex') or
						mapping.startswith('tri') or mapping.startswith('tet') ):
						elemData[name].connect[int(line[0])] = [int(j)+pN for j in line[-nodeNum:]]
						elemData[name].type=mapping
						elemData[name].nodes.update( int(j)+pN for j in line[-nodeNum:] )
						if mapping is'tet10':	# gmsh and febio second order tetrahedra switch the ordering of nodes 9 and 10
							c = elemData[name].connect[int(line[0])]
							c[-2],c[-1] = c[-1],c[-2]		# swap last two elements in connectivity for tet10
					else:				print 'Element type {} unaccounted for'.format(line[1])
			elif line.startswith('$PhysicalNames'):
				names = int( meshFile.readline() )
				for i in xrange(names):
					line=meshFile.readline().split()
					# line[0]	dimensionality of group as a string
					# line[1]	physical group # as a string
					# line[2]	name as a string already in quotes
					physGroups[line[1]] = line[2]
					#   Populate FebioElementData here
					elemData[line[2]] = FebioElement(line[1])
			line=meshFile.readline()
		meshFile.close()
	elif file.endswith('.stl'):
		stlMaterialNumber = '4'			# for material assignment in FeBio
		pN = previousNodes + 1		# +1 beacuse python list start at 0 and the index will be added
		myMesh = mesh.Mesh.from_file(file)
		pts = set()
		for p in myMesh.points:
			pts.add(tuple(p[:3]))
			pts.add(tuple(p[6:]))
			pts.add(tuple(p[3:6]))
		#
		points = list(pts)
		previousNodes += len(points)
		#
		for id, p in enumerate(points):
			nodeData[pN+id] = p  	# mesh node number + previousNodes
		#
		name = '"{}"'.format(file.split('/')[-1])
		elemData[name] = FebioElement(stlMaterialNumber)
		physGroups[stlMaterialNumber] = name
		for e, p in enumerate(myMesh.points):
			v1, v2, v3 = tuple(p[:3]), tuple(p[6:]), tuple(p[3:6])  # stl vertices
			connect = (pN+points.index(v1), pN+points.index(v2), pN+points.index(v3))
			elemData[name].connect[pN+e] = connect
			elemData[name].type = 'tri3'
	elif os.path.isdir(file):		pass
	elif file.endswith('._msh' ):	pass
	elif file.endswith('.py'  ):	pass
	elif file.endswith('.feb' ):	pass
	elif file.endswith('.xplt'):	pass
	elif file.endswith('.png' ):	pass
	elif file.endswith('.geo' ):	pass
	elif file.endswith('.txt' ):	pass
	elif file.endswith('.log' ):	pass
	elif file.endswith('.prv' ):	pass
	elif file.endswith('.dat' ):	pass
	else:	print 'This file does not handle {}'.format(file)


outFile = myDirectory.split('/')[-1]+'.feb'
output = open(outFile, 'w')
print '\n', 'Converting...'


for name in sorted(elemData):
	if 'contact' in name:
		contactName = name.replace('surf_', '')
		otherSurf = name.rsplit('_', 1)[0]+'"'
		if len(elemData[name].connect) <= len(elemData[otherSurf].connect):
			contactGroups[contactName]=(name, otherSurf)
		else:contactGroups[contactName]=(otherSurf, name)
	elif '.stl' in name:
		newName = name.split('.')[0] + '_surf"'
		contactName = '"contact_'+name[1:]
		otherSurf  = '"skin_surf"'
		if len(elemData[name].connect) <= len(elemData[otherSurf].connect):
			contactGroups[contactName]=(newName, otherSurf)
		else:contactGroups[contactName]=(otherSurf, newName)


output.write('<?xml version="1.0" encoding="ISO-8859-1"?>\n')
output.write('<!-- This document was written by\n	{}\n	on {}-->\n'.format(sys.argv[0], time.strftime('%b %d %Y at %H:%m:%S %p') ))
output.write('<febio_spec version="2.0">\n')

# Material
output.write('	<Material>\n')
for name,e in sorted(elemData.iteritems(), key=lambda eD: eD[1].mat):	# sort by material
	if 'surf' in name	or	'contact' in name	or	'Fixed' in name:	pass
	elif 'Probe' in name or 	'Rigid' in name:
		output.write('		<material id="{}" name={} type="rigid body">\n'.format(e.mat, name))
		output.write('			<density>2.8e-09</density>\n')
		output.write('		</material>\n')
	else:
		output.write('		<material id="{}" name={} type="neo-Hookean">\n'.format(e.mat ,name))
		# How should I handle setting the material parameters?
		output.write('			<density>1e-09</density>\n')
		output.write('			<E>100</E>\n')
		output.write('			<v>0.3</v>\n')
		output.write('		</material>\n')
output.write('	</Material>\n')

# Geometry
output.write('	<Geometry>\n')
output.write('		<Nodes>\n')
for n in nodeData:	output.write('			<node id="{}"> {: 10.7e}, {: 10.7e}, {: 10.7e} </node>\n'.format(n, *nodeData[n]) )
output.write('		</Nodes>\n')
for name, e in sorted(elemData.iteritems()):		# NodeSets,	seperate loops for pretty xml
	if 'bone' in name or 'Symmetry' in name or 'Fixed' in name:
		output.write('		<NodeSet name={}>\n'.format(name))
		for id in sorted(e.nodes):	output.write('			<node id="{}"/>\n'.format(id) )
		output.write('		</NodeSet>\n')
for name,e in sorted(elemData.iteritems(), key=lambda eD: eD[1].mat):	# sort by material number
	# FeBio does not require different element to have unique elem ids
	if not e.type:	raw_input('None type found')
	if (    e.type.startswith('tet') or
		 e.type.startswith('hex') or
		(e.type.startswith('tri') and '.stl' in name) ):
			output.write('		<Elements type="{}" mat="{}" name={}>\n'.format(e.type, e.mat, name))
			for id, c in sorted(e.connect.iteritems()):    output.write('			<elem id="{}"> {} </elem>\n'.format(id, str(c)[1:-1]))
			output.write('		</Elements>\n')
for name, e in sorted(elemData.iteritems()):		# Surfaces,	seperate loops for pretty xml
	if not e.type:raw_input('None type found')
	if e.type.startswith('tri') or e.type.startswith('quad'):
		if '_surf' in name:
			output.write('		<Surface name={}>\n'.format(name))
			for id, c in sorted(e.connect.iteritems()):	output.write('			<{} id="{}"> {} </{}>\n'.format( e.type, id, str(c)[1:-1], e.type ))
			output.write('		</Surface>\n')
		elif '.stl' in name:
			output.write('		<Surface name={}>\n'.format(name.split('.')[0]+'_surf"'))
			for id, c in sorted(e.connect.iteritems()):	output.write('			<{} id="{}"> {} </{}>\n'.format(e.type, id, str(c)[1:-1], e.type))
			output.write('		</Surface>\n')
# ElementData?
output.write('	</Geometry>\n')

# Boundary
#	Repeat fix as necessary (options: x, y, z displacement and u, v, w rotation on x, y, z)
output.write('	<Boundary>\n')
for name in elemData:
	if 'bone' in name:			output.write('		<fix bc="xyz" set={}/>\n'.format(name) )
	elif 'Symmetry' in name:		output.write('		<fix bc="{}" set={}/>\n'.format(name[1], name) )
	elif 'Fixed' in name:		output.write('		<fix bc={}" set={}/>\n'.format(name.split('Fixed')[0], name) )
	# elif 'Fixed' in name:		output.write('		<fix bc={} set={}/>\n'.format(name.replace('Fixed', ''), name))
output.write('	</Boundary>\n')

#Initial Constraints
output.write('	<Constraints>\n')
for name,e in sorted(elemData.iteritems(), key=lambda eD: eD[1].mat):	# sort by material
	if 'Probe' in name or 	'Rigid' in name:
		output.write('		<rigid_body mat="{}">\n'.format(e.mat))
		output.write('			<fixed bc="x"/>\n')
		output.write('			<fixed bc="y"/>\n')
		output.write('			<fixed bc="z"/>\n')
		output.write('			<fixed bc="Rx"/>\n')
		output.write('			<fixed bc="Ry"/>\n')
		output.write('			<fixed bc="Rz"/>\n')
		output.write('		</rigid_body>\n')
output.write('	</Constraints>\n')

# Contact
output.write('	<Contact>\n')
for name, cG in contactGroups.iteritems():
	if 'stl' in name:
		output.write('		<contact type="facet-to-facet sliding" name={}>\n'.format(name) )
		output.write('			<laugon>1</laugon>\n')
		output.write('			<penalty>10</penalty>\n')
		output.write('			<maxaug>10</maxaug>\n')
		output.write('			<two_pass>0</two_pass>\n')
		output.write('			<surface type="master" set={}/>\n'.format(cG[0]) )
		output.write('			<surface type="slave"  set={}/>\n'.format(cG[1]) )
		output.write('		</contact>\n')
	elif 'contact' in name:
		output.write('		<contact type="sticky" name={}>\n'.format(name) )		# sticky is tied but initial contact is not necessary!
		output.write('			<laugon>1</laugon>\n')
		output.write('			<penalty>10</penalty>\n')
		output.write('			<maxaug>10</maxaug>\n')
		output.write('			<surface type="master" set={}/>\n'.format(cG[0]) )
		output.write('			<surface type="slave"  set={}/>\n'.format(cG[1]) )
		output.write('		</contact>\n')
output.write('	</Contact>\n')

# # Loads (i.e. gravity)
output.write('	<Loads>\n')
output.write('		<body_load type = "const">\n')
output.write('			<z lc="1"> 1000 </z>\n')
output.write('		</body_load >\n')
output.write('	</Loads >\n')

#Load Data
output.write('	<LoadData>\n')
output.write('		<loadcurve id="1" type="smooth">\n')
output.write('			<point>0,0</point>\n')
output.write('			<point>1,1</point>\n')
output.write('		</loadcurve>\n')
#
output.write('		<loadcurve id="2" type="step">\n')
steps=50
for p in xrange(steps+1):
	output.write('			<point>{:4.2f},{}</point>\n'.format(float(p)/steps,1.0/steps) )
output.write('		</loadcurve>\n')
#
output.write('	</LoadData>\n')

#Output requirements
output.write('	<Output>\n')
output.write('		<plotfile type="febio">\n')
output.write('			<var type="contact gap"/>\n')
output.write('			<var type="contact pressure"/>\n')
output.write('			<var type="contact traction"/>\n')
output.write('			<var type="displacement"/>\n')
output.write('			<var type="reaction forces"/>\n')
output.write('			<var type="stress"/>\n')
output.write('		</plotfile>\n')
output.write('	</Output>\n')


# Step
output.write('	<Step name="Step01">\n')
#
output.write('		<Module type="solid"/>\n')
#
output.write('		<Control>\n')
output.write('			<time_steps>{}</time_steps>\n'.format(steps))		# required tag
output.write('			<step_size>{}</step_size>\n'.format(1.0/steps))		# required tag
output.write('			<time_stepper>\n')
output.write('				<dtmin>1e-10</dtmin>\n')
output.write('				<dtmax lc="2">{}</dtmax>\n'.format(1.0/steps))	# Load curve necessary for MustPoint printing
output.write('				<max_retries>25</max_retries>\n')
output.write('				<opt_iter>10</opt_iter>\n')
output.write('				<aggressiveness>1</aggressiveness>\n')
output.write('			</time_stepper>\n')
output.write('			<optimize_bw>1</optimize_bw>\n')
output.write('			<plot_level>PLOT_MUST_POINTS</plot_level>\n')
output.write('			<analysis type="static"/>\n')
output.write('		</Control>\n')
#
output.write('		<Constraints>\n')
for name,e in sorted(elemData.iteritems(), key=lambda eD: eD[1].mat):	# sort by material
	if 'Probe' in name or 	'Rigid' in name:
		output.write('		<rigid_body mat="{}">\n'.format(e.mat))
		#~ output.write('				<prescribed bc="x" lc="1">-20</prescribed>\n')
		output.write('				<prescribed bc="y" lc="1">-30</prescribed>\n')
		#~ output.write('				<prescribed bc="z" lc="1">-35</prescribed>\n')
		output.write('			</rigid_body>\n')
output.write('		</Constraints>\n')
#
output.write('	</Step>\n')


output.write('</febio_spec>\n')
output.close()

print 'FEBio XML written to', outFile