import sys
import time
import os

try:				from lxml import etree as et		# not in my standard Salome's python but better
except ImportError:	import xml.etree.cElementTree as et	# functionally equavalent, but no pretty printing

# the built-in Element constructor doesn't allow including text as a non-tag item, so we make a function that does
def myElement(tag, parent=None, text=None, attrib={}, **extra):
	e = et.Element(tag, attrib, **extra)
	if text:	e.text = text
	if parent is not None:	# because "if parent" doesn't work as I think is should.
		parent.append(e)
	return e

def MainSofa(parent):
	"""Define all the "standard" sofa elements
	 	that will reminan unchanged by mesh changes"""
	myElement('VisualStyle', parent, displayFlags='showVisualModels')
	myElement('BackgroundSetting', parent, image = "Meshes/Multis.png")
	myElement('DefaultPipeline', parent, name="DefaultCollisionPipeline",verbose="0",draw="0",depth="100")
	myElement('BruteForceDetection', parent, name = "Detection")
	myElement('MinProximityIntersection', parent, name="Proximity",alarmDistance="0.03",contactDistance="0.01")
	myElement('DefaultContactManager', parent, name="Response", response="default")
	myElement('DefaultCollisionGroupManager', parent, name="Group")
	myElement('MinProximityIntersection', parent, alarmDistance="8", contactDistance="5")	# This one might actually need to be changed.
	myElement('EulerImplicit', parent, name="cg_odesolver", printLog="false")
	myElement('CGLinearSolver', parent, iterations="25",name="linear solver",tolerance="1.0e-5",threshold="1.0e-5")

def SofaObjectNode(parent, part):
	"""Creates the xml Node for an object in the Sofa display"""
	name = part.name
	object = myElement('Node', parent, name=name)
	if part.mat == 'rigid':
		myElement('MechanicalObject', object, template="Rigid3d", name="{}Object".format(name), translation="0 0 0", rotation="0 0 0")
		visu=myElement('Node', object, name="Visu")
		myElement('MeshGmshLoader', visu, filename='{}.msh'.format(name), name="loader")
		myElement('TriangleSetTopologyContainer', visu, name="Container", src="@loader")
		myElement('VisualModel', visu, template="ExtVec3f", name="Visual", material="Default Diffuse 1 0.7 0.7 0.7 1 Ambient 1 0.8 0.8 0.8 1 Specular 0 1 0 0 1 Emissive 0 1 0 0 1 Shininess 0 45")
		myElement('RigidMapping', visu, input="@..", output="@Visual")
		for tie in part.Ties:
			t=myElement('Node', object, name=tie)
			myElement('MeshGmshLoader', t, filename='{}_@_{}.msh'.format(tie,name), name="loader")
			myElement('MechanicalObject', t, src="@loader", name="Attach")
			myElement('RigidMapping', t, input="@..",output="@Attach")
	else:
		myElement('MeshGmshLoader', object, filename='{}.msh'.format(name), name="loader", translation="0 0 0", rotation="0 0 0")
		myElement('MechanicalObject', object, src="@loader", name="Volume")
		myElement('include', object, href="Objects/TetrahedronSetTopology.xml", src="@loader")
		myElement('UniformMass', object, totalmass="100")				#hmmm
		myElement('TetrahedralCorotationalFEMForceField', object, youngModulus="300", poissonRatio="0.4" ,method="large")	# more logic needed!
		surf = myElement('Node', object, name="Surface")
		myElement('TriangleSetTopologyContainer', surf, name="Container")
		myElement('TriangleSetTopologyModifier', surf, name="Modifier")
		myElement('TriangleSetTopologyAlgorithms', surf, name="TopoAlgo", template="Vec3d")
		myElement('TriangleSetGeometryAlgorithms', surf, name="GeomAlgo", template="Vec3d")
		myElement('Tetra2TriangleTopologicalMapping', surf, input="@../Container", output="@Container")
		myElement('TriangleSet', surf, contactStiffness='3000')
		myElement('LineModel',   surf, contactStiffness='3000')
		myElement('PointModel',  surf, contactStiffness='3000')
		visu=myElement('Node', surf ,name="Visu")
		myElement('OglModel', visu, name="Visual", material="Default Diffuse 1 1 0.6 0.6 1 Ambient 1 1 0.6 0.6 1 Specular 0 1 0 0 1 Emissive 0 1 0.9 0.8 1 Shininess 0 10")
		myElement('IdentityMapping', visu)

def WriteAsGmsh(salomeDat):
	""""Takes a SALOME dat file of a mesh and write a gmsh format v1.0 mesh"""
	outMesh = salomeDat.replace('.dat', '.msh')
	print 'Writing {}'.format(outMesh)
	gmsh=open(outMesh, 'w')
	data=open(salomeDat, 'r')
	nodes, elements = [int(s) for s in data.readline().split()]
	gmsh.write('$NOD\n')
	gmsh.write('{}\n'.format(nodes))
	#
	nodeConversion={}
	for i in range(1, nodes+1):
		line=data.readline().split()
		nodeConversion[line[0]]=str(i)
		line[0]=str(i)
		line = ' '.join(line) + '\n'
		gmsh.write(line)
	#
	gmsh.write('$ENDNOD\n')
	gmsh.write('$ELM\n')
	gmsh.write('{}\n'.format(elements))
	for i in range(1,elements+1):
		line = data.readline().split()
		#
		# line contains		Salome elem #, Salome elem type, node list
		# becomes			Gmsh elem #, Gmsh elem type, PG#, EE#, #ofnodes, modified node list
		#
		numNodes=int(line[1][2:])		# line[1] == Salome elem type --> dimension of elem, 0,  # of nodes
		if line[1]=='203':			line[1] = '2 0 1 3'  # triangle,    PG#, EE#, # of nodes
		elif line[1]=='304':		line[1] = '4 1 1 4'  # tetrahedron, PG#, EE#, #nodes
		else:					print line[1], ' not supported by this script\n', 'See documentation.';	sys.exit()
		line[-numNodes:] = [nodeConversion[i] for i in line[-numNodes:]]  # change node list to renumbered nodes
		line = ' '.join(line) + '\n'
		gmsh.write(line)
	gmsh.write('$ENDELM\n')
	gmsh.close()
	data.close()
	os.remove(salomeDat)

def _testSofa(sofaFile='Sofa.scn', connectivityFile='Connectivity.xml'):
	import os
	assert os.path.isfile(connectivityFile)
	import ConnectivityXml as xml
	import StlToMed
	if os.path.dirname(connectivityFile) != os.getcwd() and os.path.dirname(connectivityFile):
		os.chdir(os.path.dirname(connectivityFile))  # in case connectivity file doesn't give file path information.  Because it likely won't.
	print '\nUsing {} at\n{}'.format(os.path.basename(connectivityFile),os.getcwd())
	#
	Assembly = xml.ReadConnectivity(connectivityFile)
	#
	for part in Assembly.itervalues():
		part.mesh = StlToMed.MakeMesh(part.file, part.mat, part.name)
	#
	sofa = et.Element('Node', {'name':'root','dt':'0.02','gravity':'0 0 0'})
	tree = et.ElementTree(sofa)
	#
	import inspect
	comment = et.Comment(
		'This document was written by\n	{}\n	on {}'.format(inspect.getfile(inspect.currentframe()),
																 time.strftime('%b %d %Y at %H:%m:%S %p')))
	sofa.append(comment)
	#
	MainSofa(sofa)
	#
	for part in Assembly.itervalues():
		SofaObjectNode(sofa, part)
		#
		file = '{}.dat'.format(part.name)
		if part.mat == 'rigid':
			part.mesh.ExportDAT(file)
			for tiedName in part.Ties:
				StlToMed.MakeProximityGroups(part, Assembly[tiedName])
				#
				tiedFile = '{}_@_{}'.format(tiedName, file)
				group, = Assembly[tiedName].mesh.GetGroupByName('{}_@_{}Faces'.format(tiedName, part.name))
				Assembly[tiedName].mesh.ExportDAT(tiedFile, group)
				WriteAsGmsh(tiedFile)
				#
				elems = group.GetIDs()
				nodes2 = set()
				for elem in elems:
					nodes = Assembly[tiedName].mesh.GetElemNodes(elem)
					for n in nodes:
						nodes2.add(n)
				#
				nodes2 = sorted([i - 1 for i in nodes2])  # Sofa nodes start at zero, gmsh at 1
				nodes1 = range(len(nodes2))  # We use a renumbered subset of the Ties mesh as our connection
				#
				myElement('AttachConstraint', sofa, object1="@{}/{}/Attach".format(part.name, tiedName),
						  object2="@{}".format(tiedName),
						  twoWay="False", radius="1", indices1=str(nodes1)[1:-1], indices2=str(nodes2)[1:-1])
		else:
			group, = part.mesh.GetGroupByName('{}_All'.format(part.name))  # Makes list of one
			part.mesh.ExportDAT(file, group)
		WriteAsGmsh(file)
	try:				tree.write(sofaFile, xml_declaration=True, pretty_print=True)
	except TypeError:	tree.write(sofaFile, xml_declaration=True)
	print '\nWrote Sofa.scn'


if __name__ == '__main__':
	print 'Usage: salome -t StlToSofa.py args:Sofa.scn, ConnectivityFile'
	_testSofa(*sys.argv[1:])