import sys
import time
import os
try:
	import salome
	import salome_notebook
	import SMESH
	from salome.smesh import smeshBuilder
except ImportError:
	import inspect
	print '\n\n\nUsage: salome {} args:ConnectivityFile,Abaqus.inp  \n\n\n'.format( inspect.getfile(inspect.currentframe()) )
	sys.exit()
from MEDLoader import MEDLoader

from collections import Counter
import itertools
chain = itertools.chain.from_iterable

import StlToMed

salome.salome_init()
theStudy = salome.myStudy

notebook = salome_notebook.NoteBook(theStudy)
smesh = smeshBuilder.New(theStudy)

def CloseSalome():
	"""Try to close out of Salome gracefully"""
	# there is nothing graceful about this
	try:
		sys.stdout = 'redirect to nowhere...'  # so I don't see the junk it prints
		from killSalomeWithPort import killMyPort
		killMyPort(2810)
	except:		pass

def WritePart(part, abaqus, formatting=False):
	"""Output an Abaqus Part"""
	abaqus.write('*Part, Name = {}\n'.format(part.name))
	mesh = part.mesh
	medMesh = MEDLoader.ReadUMeshFromFile(part.med, 0)
	dimension = mesh.MeshDimension()
	nodeCount = medMesh.getNumberOfNodesInCell(0)
	print "Nodes in Element: ", nodeCount
	# Nodes
	abaqus.write('	*Node\n')


	if formatting:	writeString ='	{:6d}, {: 12.7f}, {: 12.7f}, {: 12.7f}\n'
	else: 			writeString = '	{}, {},	{}, {}\n'
	nodes = MEDLoader.ReadUMeshFromFile(part.med, 0).getCoords()
	#print medMesh.getNumberOfNodesInCell
	#print *nodes
	for node in xrange(len(nodes)):
		#print nodes.getTuple(node)
		abaqus.write(writeString.format(node+1, *nodes.getTuple(node)))
	#
	if part.mat == 'rigid':
		abaqus.write(writeString.format(node+2, *StlToMed.FindCenter(mesh)))
		abaqus.write('	*NSet, NSet=RigidNode\n\t{}\n'.format(node+2))
	# Node Sets
	groups = mesh.GetGroups()
	for group in groups:
		name = group.GetName()
		if name.endswith('Nodes'):
			abaqus.write('	*NSet, NSet={}\n'.format(name))
			WriteListOfItems(sorted(group.GetListOfID()), abaqus,formatting)
	# Elements and Surfaces

	if medMesh.getMeshDimension() == 2:
		# Elements
		abaqus.write('	*Element, Type=S3\n')
		meshConnectivity = medMesh.getNodalConnectivity()
		numElems = medMesh.getNumberOfCells()
		if formatting:		writeString = '	{:6d}, {:6d}, {:6d}, {:6d}\n'
		else:				writeString = '	{}, '* nodeCount +' {}\n'
		for elem in xrange(numElems):
			index = elem * (nodeCount+1)
			connectivity = [meshConnectivity[index + c] + 1 for c in range(1, nodeCount + 1)]	# sliceing does't work quite right as this is a Salome structure
			abaqus.write(writeString.format(elem+1, *connectivity))
		# Elements Sets and Surfaces(built from Element Sets)
		for group in groups:
			name = group.GetName()
			if name.endswith('Faces'):
				abaqus.write('	*ElSet, ElSet=_{}, internal\n'.format(name))
				WriteListOfItems(sorted(group.GetListOfID()), abaqus, formatting)
				#
				abaqus.write('	*Surface, type=Element, name={}\n'.format(name))
				abaqus.write('	_{},SPOS\n'.format(name))
			elif name.endswith('Nodes'):			pass
			else:
				print 'Unknown Surface type', name
				sys.exit()
		abaqus.write('	*Shell Section, elset=_{}_All_Faces, material={}\n'.format(part.name, part.name))
		abaqus.write('	1., 5\n'.format(part.name, part.name))
	elif medMesh.getMeshDimension() == 3:

		# Elements
		abaqus.write('	*Element, TYPE={}\n'.format('C3D'+str(nodeCount)))
		meshConnectivity = medMesh.getNodalConnectivity()
		numElems = medMesh.getNumberOfCells()
		#####################
		#####################
		#####################
		if formatting:	writeString = '	{:7d}, {:6d}, ' + '{:6d}, '*(nodeCount-2) +'{:6d}\n'
		else:			writeString = '	{}, ' *nodeCount +'{}\n'
		# variable length is slightly slower than fixed  #	writeString = ('	{:6d}'+', {:6d}'*nodeCount+'\n')
		for elem in xrange(numElems):

			index = elem * (nodeCount + 1)
			connectivity = [meshConnectivity[index+c]+1 for c in range(1, nodeCount+1)]
			#swap nodes as needed to match between MED and abaqus
			if nodeCount == 4: #tet4
				connectivity[1], connectivity[2] = connectivity[2], connectivity[1]  # swap nodes n2 and n3
			elif nodeCount == 6: #penta6
				connectivity[1], connectivity[2] = connectivity[2], connectivity[1]  # swap nodes n2 and n3
				connectivity[4], connectivity[5] = connectivity[5], connectivity[4]  # swap nodes n5 and n6
			elif nodeCount == 8: #hex8
				connectivity[1], connectivity[3] = connectivity[3], connectivity[1]  # swap nodes n2 and n4
				connectivity[5], connectivity[7] = connectivity[7], connectivity[5]  # swap nodes n6 and n8
			elif nodeCount == 10: #tet10
				connectivity[1], connectivity[2] = connectivity[2], connectivity[1]  # swap nodes n2 and n3
				connectivity[4], connectivity[6] = connectivity[6], connectivity[4]  # swap nodes n5 and n7
				connectivity[8], connectivity[9] = connectivity[9], connectivity[8]  # swap nodes n9 and n10
			elif nodeCount == 15: #penta15
				connectivity[1], connectivity[2] = connectivity[2], connectivity[1]  # swap nodes n2 and n3
				connectivity[4], connectivity[5] = connectivity[5], connectivity[4]  # swap nodes n5 and n6
				connectivity[6], connectivity[8] = connectivity[8], connectivity[6]  # swap nodes n5 and n6
				connectivity[9], connectivity[11] = connectivity[11], connectivity[9]  # swap nodes n5 and n6
				connectivity[13], connectivity[14] = connectivity[14], connectivity[13]  # swap nodes n5 and n6
			elif nodeCount == 20: #hex20
				connectivity[1], connectivity[3] = connectivity[3], connectivity[1]  # swap nodes n2 and n4
				connectivity[5], connectivity[7] = connectivity[7], connectivity[5]  # swap nodes n6 and n8
				connectivity[8], connectivity[11] = connectivity[11], connectivity[8]  # swap nodes n9  and n12
				connectivity[9], connectivity[10] = connectivity[10], connectivity[9]  # swap nodes n10 and n11
				connectivity[12], connectivity[15] = connectivity[15], connectivity[12]  # swap nodes n13 and n16
				connectivity[13], connectivity[14] = connectivity[14], connectivity[13]  # swap nodes n14 and n15
				connectivity[17], connectivity[19] = connectivity[19], connectivity[17]  # swap nodes n18 and n20
			else:
				raw_input("Unexpected node count, Update script or input files. Press enter to break.")
				break


			#print writeString.format(elem+1, *connectivity)
			abaqus.write(writeString.format(elem+1, *connectivity))
		######
		######
		######

		# Elements Sets and Surfaces(built from Element Sets)
		for group in groups:
			name = group.GetName()
			print "Group Name: ", name
			numFaces = mesh.NbFaces()
			#print numFaces
			meshMatrix=mesh.GetElementsByType(SMESH.VOLUME)
			maxMatrix = max(meshMatrix)
			minMatrix = min(meshMatrix)
			#print  "min/max: ", minMatrix, maxMatrix
			#print mesh.GetMeshInfo()
			#map the number of nodes to the number of faces. I dont think a quad shell will ever show as a tri6 after StlToMed, so just map 6 nodes to the linear pentahedron
			faceDict = {3:1, 4:4, 6:5, 10:4, 15:5, 8:6, 20:6 }
			if   name.endswith('Faces'):
				faceSets = {i: [] for i in xrange(faceDict.get(nodeCount))}
				for i, elem in enumerate(sorted(group.GetListOfID())):
					oneFaceNodes = set(mesh.GetElemNodes(elem))  # a list of all nodes in the face, a set so order doesn't matter
					eleFromNodes=[]
					for node in oneFaceNodes:
						eleFromNodes.append( mesh.GetNodeInverseElements(node) )	# make a list of element in lists
						#raw_input('press something')
					joinedTetraList = chain(eleFromNodes)	# flatten the lists
					count = Counter(joinedTetraList)

					#remove any entry that isnt a volume from the counter
					for entry in list(count):
						if  minMatrix>(entry) or (entry)>maxMatrix:
							del count[entry]

					#count[elem] = 0	#this is no longer needed with the above loop
					#if 'Skin' and 'Probe' in name:
						#print list(count.most_common())
						#raw_input('...')
					mostCommonEle = (count.most_common(1)[0][0]) 		# this will be the 3d element

					for j in xrange(nodeCount):
						faceNodes = set(mesh.GetElemFaceNodes(mostCommonEle, j))

						if faceNodes == oneFaceNodes:				# match node is tetra face and triangle group face
							faceSets[j].append(mostCommonEle-(minMatrix-1)) # subtract minmatrix instead of numFaces to make code able to handle some of the stranger element numbering i saw in penta15 elements
							break			# only one face of the element can be the matching face
				#

				#print nodeCount
				# Abaqus Face numbers different than Salome faces numbers
				if nodeCount in [4, 10]:
					#print "Tet Swap"
					faceSets[1],faceSets[3] = faceSets[3],faceSets[1] #swap faceset 2 with faceset 4
				elif nodeCount in [6, 15]:
					#print "Penta Swap"

					faceSets[2], faceSets[4] = faceSets[4], faceSets[2] #swap faceset 3 with faceset 5
				elif nodeCount in [8, 20]:
					#print "Hex Swap"
					faceSets[3], faceSets[4] = faceSets[4], faceSets[3] #swap faceset 4 with faceset 5
					faceSets[2], faceSets[5] = faceSets[5], faceSets[2]	#swap faceset 3 with faceset 6
				#	print len(faceSets[0]),len(faceSets[1]),len(faceSets[2]),len(faceSets[3]),len(faceSets[4])

				#
				# #
				#
				# print faceSets
				# actualSetsFaces = {}
				# for i, fSet in faceSets.iteritems():  # TESTING for INTERSETION
				# 	actualSetsFaces[i+1] = set(fSet)
				# 	print 'length of set ', i+1, len(actualSetsFaces[i+1])
				# Faces12 = actualSetsFaces[1].intersection(actualSetsFaces[2])
				# Faces13 = actualSetsFaces[1].intersection(actualSetsFaces[3])
				# Faces14 = actualSetsFaces[1].intersection(actualSetsFaces[4])
				# Faces23 = actualSetsFaces[2].intersection(actualSetsFaces[3])
				# Faces24 = actualSetsFaces[2].intersection(actualSetsFaces[4])
				# Faces34 = actualSetsFaces[3].intersection(actualSetsFaces[4])
				# #
				# print 'Intersection of face 1 and 2', Faces12
				# print 'Intersection of face 1 and 3', Faces13
				# print 'Intersection of face 1 and 4', Faces14
				# print 'Intersection of face 2 and 3', Faces23
				# print 'Intersection of face 2 and 4', Faces24
				# print 'Intersection of face 3 and 4', Faces34
				# raw_input('...')
				#
				# #
				#
				for i, fSet in faceSets.iteritems():  # Build internal (hidden) element sets used to build surface
					abaqus.write('	*ElSet, ElSet=_{}_{}, internal\n'.format(name, i+1))
					WriteListOfItems(sorted(fSet), abaqus, formatting)
				#
				abaqus.write('	*Surface, type=Element, name={}\n'.format(name))
				for i in faceSets:		abaqus.write('	_{}_{}, S{}\n'.format(name, i+1, i+1))
			elif name.endswith('Volumes'):
				abaqus.write('	*ElSet, ElSet={}\n'.format(name))
				elemList = [e-(minMatrix-1) for e in sorted(group.GetListOfID())]
				WriteListOfItems(elemList, abaqus, formatting)
			elif name.endswith('Nodes'):	pass
			else:
				print 'Unknown Surface type', name
				sys.exit()
		abaqus.write('	*Solid Section, elset={}_ALL_VOLUMES, material={}\n\t,\n'.format(part.name, part.name))
	else:
		print 'Bad mesh Dimension:', part.mesh.MeshDimension()
		sys.exit()
	#
	abaqus.write('*End Part\n'.format(part.name))

def WriteListOfItems(listOfItems, abaqus, formatting=False):
	"""Writes the set of nodes or elements...reusing code is good
		one line is a good balance between efficiency and memory use"""
	n = 10		# number of elements on a line.  I prefer 10 to the 16 allowable in Abaqus.
	#
	if formatting:	writeString = '	{:7d},'*n +'\n'
	else:			writeString = '	{},'   *n +'\n'  # no formatting is much faster, but harder to read output
	#
	for i in xrange( len(listOfItems)/n ):
		abaqus.write(writeString.format(*listOfItems[i*n:(i+1)*n]))
	# And the last partial line
	i=len(listOfItems)/n
	if formatting:	abaqus.write(('	{:7d},'*len(listOfItems[i*n:(i+1)*n]) + '\n').format(*listOfItems[i*n:(i+1)*n]))
	else:			abaqus.write(('	{},'*len(listOfItems[i*n:(i+1)*n]) + '\n').format(*listOfItems[i*n:(i+1)*n]))

def WriteAssembly(assembly, abaqus, name='Assembly', formatting=False):
	"""assembly is from xml reading
		abaqus is a file already opened for writing"""
	for part in assembly.itervalues():
		print 'Writing Part {}'.format(part.name)
		WritePart(part, abaqus, formatting)
	print 'Creating Assembly'
	abaqus.write('*Assembly, Name = {}\n'.format(name))
	for part in assembly.itervalues():
		abaqus.write('	*Instance, Name={}, Part={}\n'.format(part.name, part.name))
		# possible part positioning data could go here, but that not relevant for the current project
		abaqus.write('	*End Instance\n')
	# Make Rigid bodies
	for name, part in assembly.iteritems():
		if part.mat == 'rigid':
			abaqus.write('	*Rigid Body, ref node={}.RigidNode, elset={}._{}_All_Faces\n'.format(name, name, name))
	# Set ties
	for master, part in assembly.iteritems():
		switched = False
		for slave in part.Ties:
			minSize, maxSize = assembly[master].mesh.GetMinMax(SMESH.FT_Length2D)  # approximate 'characteristic length' of the mesh
			try:				multiplier = part.Ties[slave]['proximity']
			except KeyError:	multiplier=1.0
			positionTol = (maxSize + minSize) * multiplier
			#
			if assembly[slave].mat == 'rigid':	# rigid body need to be the master surface
				master, slave = slave, master
				switched = True
			name = '{}_Tied_{}'.format(master, slave)
			abaqus.write('	*Tie, Name={}, Adjust=no, Position Tolerance={}\n'.format(name, positionTol))
			abaqus.write('	{}.{}_@_{}_TiesFaces, {}.{}_@_{}_TiesFaces\n'.format(slave,slave,master,	master,master,slave))
			if switched:						# return master to original value if changed
				master = slave					#	no need to change slave as loop is ending
	#
	abaqus.write('*End Assembly\n')

def WriteContact(assembly,abaqus, name='Assembly'):
	print 'Creating new Contact conditions\n'
	#
	#  first iterate for surface interactions and behavior
	for master, part in assembly.iteritems():
		for slave in part.Contact:
			switched = False
			if assembly[slave].mat == 'rigid':
				master, slave = slave, master
				switched = True
			contact = '{}_Contact_{}'.format(master, slave)
			abaqus.write('*Surface Interaction, Name={}\n1.,\n'.format(contact))
			abaqus.write('*Surface Behavior, direct\n')
			#
			abaqus.write('*Contact Pair, interaction={}, type=SURFACE TO SURFACE\n'.format(contact))
			# abaqus.write('*Contact Pair, interaction={}, type=NODE TO SURFACE\n'.format(contact))
			abaqus.write('	{}.{}.{}_@_{}_ContactFaces,' .format(name, slave,  slave,  master))
			abaqus.write('	{}.{}.{}_@_{}_ContactFaces\n'.format(name, master, master, slave ))
			#
			if switched: 				# return master to original value if changed
				master = slave			#	no need to change slave as loop is ending

def WriteMaterials(assembly,abaqus):
	"""Only writes a generic elastic material for each part
		The user should alter the material definitions as needed."""
	print 'Creating basic Material Definitions'
	for part in assembly.itervalues():
		abaqus.write('*Material, Name={}\n'.format(part.name))
		abaqus.write('	*Hyperelastic, Mooney-Rivlin\n')
		abaqus.write('	0.001,0.,2.\n')

def WriteBoundaries(assembly,abaqus):
	"""Creates a Step and sets boundary conditions for rigid objects"""
	print 'Creating basic rigid body Boundary conditions.'
	abaqus.write('*Step, name = Step - 1, nlgeom = YES\n')
	abaqus.write('	*Static\n')
	abaqus.write('	0.1, 1., 1e-05, 0.1\n')		# Initial time increment, Total Time period, Min increment, Max increment
	DoF = [i+1 for i in xrange(6)]
	for part in assembly.itervalues():
		if part.mat == 'rigid':
			abaqus.write('	*Boundary\n'.format(part.name))
			for d in DoF:
				abaqus.write('		{}.RigidNode, {}, {}\n'.format(part.name, d, d))
	#
	abaqus.write('*Restart, write, frequency=0\n')
	abaqus.write('*Output, field, variable=PRESELECT, time interval=0.1\n')
	abaqus.write('*Output, history, variable=PRESELECT, time interval=0.1\n')
	abaqus.write('*End Step\n')

def ReadConnectivity(connectivityFile):
	"""Read the connectivity file and asserts that all files it references exist."""
	import ConnectivityXml as xml
	print '\n'  # just make some white space for easier reading of the output
	print 'Reading from connectivity file {}'.format(connectivityFile)
	if os.path.dirname(connectivityFile) != os.getcwd() and os.path.dirname(connectivityFile):
		os.chdir(os.path.dirname(connectivityFile))  # connectivity file may not give full file path information.
	#
	try:
		Assembly = xml.ReadConnectivity(connectivityFile)  # make sure
		xml.RemoveDoubleCounting(Assembly)
	except IOError:
		print '\n\nNo connectivity file given.\nUsage: salome MedToAbaqus.py args:ConnectivityFile,AbaqusResult\n'
		sys.exit()
	#
	# make sure the MED directory exists
	medDir = os.path.join(os.path.dirname(connectivityFile), 'MED')
	if not os.path.isdir(medDir):
		# do something about running StlToMed
		pass
	#
	for part in Assembly.itervalues():
		file = os.path.join(os.path.dirname(part.file),'MED',os.path.basename(part.file).replace('.stl','.med'))
		if not os.path.isfile(file):
			other = os.path.join(os.path.dirname(connectivityFile),'MED',os.path.basename(part.file).replace('.stl','.med'))
			if os.path.isfile(other):	file = other
			else:
				print '\n\nScript cannot find file {}\n	associated with {}'.format(part.file, part.name)
				print part.file
				print part.name
				print 'Please check connectivity file.\n'
				CloseSalome()
				sys.exit()
		part.file = file
		part.med  = file
		part.mesh,=StlToMed.OpenMesh(file)	# comma necessary to unpack the list (of one) returned by opening MED files
	print ''
	return Assembly

def MakeAbaqus(connectivityFile='Connectivity.xml', abaqusFile='Abaqus.inp'):
	"""Read connectivity file, get MEDs and write Abaqus input file"""
	Assembly = ReadConnectivity(connectivityFile)
	#
	if not os.path.dirname(abaqusFile):
		abaqusFile = os.path.join('Abaqus', abaqusFile)
		if not os.path.exists('Abaqus'):    os.makedirs('Abaqus')
	abaqus = open(abaqusFile, 'w')
	#
	WriteAssembly(Assembly, abaqus, formatting =True)#, 'OpenKnee')
	# WriteContact(Assembly, abaqus)#, 'OpenKnee')
	WriteContact(Assembly, abaqus)#, 'OpenKnee')
	WriteMaterials(Assembly, abaqus)
	WriteBoundaries(Assembly, abaqus)
	#
	abaqus.close()
	print 'Writing to {} file'.format(abaqusFile)
	#
	CloseSalome()

def runTest(connectivityFile='Connectivity.xml', abaqusFile='MEDCoupling.inp'):
	print '\n\n\n  WARNING - Testing MEDCouplng speed up\n\n\n'
	import ConnectivityXml as xml
	print 'Reading from connectivity file {}'.format(connectivityFile)
	Assembly = xml.ReadConnectivity(connectivityFile)
	xml.RemoveDoubleCounting(Assembly)
	#
	if not os.path.dirname(abaqusFile):
		abaqusFile = os.path.join('Abaqus', abaqusFile)
		if not os.path.exists('Abaqus'):    os.makedirs('Abaqus')
	abaqus = open(abaqusFile, 'w')
	#
	for part in Assembly.itervalues():
		file = os.path.join(os.path.dirname(part.file),'MED',os.path.basename(part.file).replace('.stl','.med'))
		part.mesh,=StlToMed.OpenMesh(file)	# comma necessary to unpack the list (of one) returned by opening MED files
	#
	for part in Assembly.itervalues():
		print '\nWriting Part {}'.format(part.name)
		WritePart(part, abaqus)
	#
	abaqus.close()
	#
	print 'Writing to {} file'.format(abaqusFile)
	print '\n\n\n  WARNING - Testing MEDCouplng speed up\n\n\n'
	# try to close out of salome
	try:
		from killSalomeWithPort import killMyPort
		killMyPort(2810)
	except:	pass

if __name__ == '__main__':
	# runTest(*sys.argv[1:])
	MakeAbaqus(*sys.argv[1:])