#=============================================================================
# Monolithic FE^2
# Nils Lange, Geralf Huetter, Bjoern Kiefer
#   Nils.Lange@imfd.tu-freiberg.de, Geralf.Huetter@imfd.tu-freiberg.de, 
#   Bjoern.Kiefer@imfd.tu-freiberg.de
# distributed under CC BY-NC-SA 4.0 license
# (https://creativecommons.org/licenses/by-nc-sa/4.0/)
# Reference: 
#   N. Lange, G. Huetter, B. Kiefer: "An efficient monolithic solution scheme for FE2 problems",
#   https://arxiv.org/abs/2101.01802
#
# Further information on the implementation, structure of the source code,
# examples and tutorials can be found in the file doc/documentation.pdf
#
#=============================================================================


import os
import sys
import numpy as np
import subprocess
from set_FE2_Analysisparameters import *

def generate_training_inputfile(strain_amp,stress_amp,n_variations,dimens,jobname,nlgeom,Method,
		                            simulation_time,dtime_data_dump,training_mode,start_simulations,
		                            program_directory,ncpus):
	
	#definition of trainig directions over angles
	
	##define the strain (large deformation: right stretch) training directions
	#values=[]
	#if dimens==2: #simple rotation around one axis
	#	alpha=np.linspace(float(alpha.split(',')[0]),float(alpha.split(',')[1]),int(alpha.split(',')[2]))
	#	phi1=np.linspace(float(phi1.split(',')[0]),float(phi1.split(',')[1]),int(phi1.split(',')[2]))
	#	hydrostatic_part=np.eye(2)/np.sqrt(2)
	#	for Alpha in alpha:
	#		if np.abs(np.cos(Alpha))>0.001:
	#			for Phi1 in phi1:
	#				N=np.array([[np.cos(Phi1),-np.sin(Phi1)],
	#							[np.sin(Phi1),np.cos(Phi1)]])
	#				E=np.array([1.0,-1.0])
	#				deviatoric_part=(E[0]*np.outer(N[:,0],N[:,0])+E[1]*np.outer(N[:,1],N[:,1]))/np.sqrt(2)
	#				values.append(hydrostatic_part*np.sin(Alpha)+deviatoric_part*np.cos(Alpha))
	#		else:
	#			values.append(hydrostatic_part*np.sin(Alpha)/np.abs(np.sin(Alpha)))
	#else: #complex rotation using quaternions (Euler-Rodrigues Formula)
	#	phi1=np.linspace(float(phi1.split(',')[0]),float(phi1.split(',')[1]),int(phi1.split(',')[2]))
	#	phi2=np.linspace(float(phi2.split(',')[0]),float(phi2.split(',')[1]),int(phi2.split(',')[2]))
	#	phi3=np.linspace(float(phi3.split(',')[0]),float(phi3.split(',')[1]),int(phi3.split(',')[2]))
	#	alpha=np.linspace(float(alpha.split(',')[0]),float(alpha.split(',')[1]),int(alpha.split(',')[2]))
	#	theta=np.linspace(float(theta.split(',')[0]),float(theta.split(',')[1]),int(theta.split(',')[2]))
	#	hydrostatic_part=np.eye(3)/np.sqrt(3)
	#	for Alpha in alpha:
	#		if np.abs(np.cos(Alpha))>0.001:
	#			for Theta in theta:
	#				for Phi1 in phi1:
	#					for Phi2 in phi2:
	#						for Phi3 in phi3:
	#							q=[np.sin(Phi1)*np.cos(Phi2)*np.sin(Phi3/2.0),np.sin(Phi1)*np.sin(Phi2)*np.sin(Phi3/2.0),np.cos(Phi1)*np.sin(Phi3/2.0),np.cos(Phi3/2.0)] #quaternion
	#							N=np.array([[1-2*(q[1]**2+q[2]**2),2*(q[0]*q[1]-q[3]*q[2]),2*(q[0]*q[2]+q[3]*q[1])],
	#										[2*(q[0]*q[1]+q[3]*q[2]),1-2*(q[0]**2+q[2]**2),2*(q[1]*q[2]-q[3]*q[0])],
	#										[2*(q[0]*q[2]-q[3]*q[1]),2*(q[1]*q[2]+q[3]*q[0]),1-2*(q[0]**2+q[1]**2)]])
	#							E=np.array([np.cos(Theta)-np.sin(Theta)/np.sqrt(3.0),2.0*np.sin(Theta)/np.sqrt(3.0),-np.cos(Theta)-np.sin(Theta)/np.sqrt(3.0)])
	#							deviatoric_part=(E[0]*np.outer(N[:,0],N[:,0])+E[1]*np.outer(N[:,1],N[:,1])+E[2]*np.outer(N[:,2],N[:,2]))/np.sqrt(2)
	#							values.append(hydrostatic_part*np.sin(Alpha)+deviatoric_part*np.cos(Alpha))
	#		else:
	#			values.append(hydrostatic_part*np.sin(Alpha)/np.abs(np.sin(Alpha)))
	
	
	#define the strain (large deformation: right stretch) training directions
	variations=np.linspace(-1.0,1.0,n_variations)
	values=[]
	if dimens==2: #simple rotation around one axis
		for e11 in variations:
			for e22 in variations:
				for e12 in variations:
					N=np.array([[e11,e12],[e12,e22]])
					N_norm=np.linalg.norm(N)
					if N_norm>0.00001:
						values.append(N/N_norm)
	else:
		for e11 in variations:
			for e22 in variations:
				for e33 in variations:
					for e12 in variations:
						for e23 in variations:
							for e13 in variations:
								N=np.array([[e11,e12,e13],[e12,e22,e23],[e13,e23,e33]])
								N_norm=np.linalg.norm(N)
								if N_norm>0.00001:
									values.append(N/N_norm)
	
	#Simulation of the Trainingdirections with Abaqus
	
	#P_total=len(values) #total number of directions
	#################################################################################
	##create abaqus basic input file (without training directions)
	#nodes={2:[[0.0,0.0],[1.0,0.0],[1.0,1.0],[0.0,1.0]],
	#       3:[[0.0,0.0,0.0],[1.0,0.0,0.0],[1.0,1.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0],[1.0,0.0,1.0],[1.0,1.0,1.0],[0.0,1.0,1.0]]}
	#f=open("basic.inp","w")
	#f.write("*Heading\n")
	#f.write("*Preprint, echo=NO, model=NO, history=NO, contact=NO\n")
	#f.write("*Part, name=RVE\n")
	#f.write("*Node\n")
	#for i in range(len(nodes[dimens])):
	#	f.write(str(i+1)+','+','.join(map(str,nodes[dimens][i][:]))+'\n')
	#if dimens==2:
	#	f.write("*Element, type=CPE4R\n1, 1, 2, 3, 4\n")
	#else:
	#	f.write("*Element, type=C3D8R\n1, 1, 2, 3, 4, 5, 6, 7, 8\n")
	#f.write("*Elset, elset=Set-1\n1,\n")
	#for n in range(len(nodes[dimens])):
	#	f.write("*Nset, nset=Set-"+str(n+1)+"\n")
	#	f.write(str(n+1)+",\n")
	#f.write("*Solid Section, elset=Set-1, controls=EC-1, material=Material-1\n ,\n*End Part\n")
	#f.write("*Assembly, name=Assembly\n")
	#f.write("*Instance, name=Sample, part=RVE\n*End Instance\n")
	#a=1
	#for i in range(dimens):
	#	for j in range(i,dimens):
	#		f.write("*Node\n"+str(a)+",0.0,0.0,0.0\n")#reference points for strain/stress
	#		a+=1
	#a=1
	#for i in range(dimens):
	#	for j in range(i,dimens):
	#		f.write("*Nset,nset=E"+str(i+1)+str(j+1)+"\n"+str(a)+",\n")#reference point sets for strain/stress
	#		a+=1
	##equations for impose strain/stress onto nodes
	#for n in range(1,len(nodes[dimens])):
	#	for i in range(dimens):
	#		f.write("*Equation\n"+str(dimens+2)+"\nSample.Set-"+str(n+1)+","+str(i+1)+",-1\nSample.Set-1,"+str(i+1)+",1\n")
	#		for j in range(dimens):
	#			f.write("E"+str(np.min([i+1,j+1]))+str(np.max([i+1,j+1]))+",1,"+str(nodes[dimens][n][j]-nodes[dimens][0][j])+"\n")
	#f.write("*End Assembly\n")
	#f.write("*Section Controls, name=EC-1, hourglass=ENHANCED\n")
	#f.write("1., 1., 1.\n")
	#f.write("*Material, name=Material-1\n")
	#f.write("*User Material, constants=1\n1\n")
	#f.write("*Step, name=Step-1, nlgeom="+nlgeom+", inc=1000\n*Static\n0.001, "+str(simulation_time)+", 1e-07, "+str(dtime_data_dump)+"\n")
	#f.write("*Controls,Parameters=Field,Field=Displacement\n0.0001,,,,,,,\n,,,\n")
	#f.close()
	##create file with the training directions
	#f=open("directions.inp","w")
	#for p in range(P_total):
	#	if (training_mode=="strain" or training_mode=="stress and strain"):
	#		f.write("*Boundary\n")
	#		for i in range(dimens):
	#			for j in range(i,dimens):
	#				f.write("E"+str(i+1)+str(j+1)+",1,1,"+str(values[p][i,j]*strain_amp)+"\n")
	#		for i in range(dimens):
	#			f.write("Sample.Set-1,"+str(i+1)+","+str(i+1)+"\n") #prevent rigid body motion
	#		f.write("*End Step\n")
	#	if (training_mode=="stress" or training_mode=="stress and strain"):
	#		f.write("*Cload\n")
	#		for i in range(dimens):
	#			for j in range(i,dimens):
	#				f.write("E"+str(i+1)+str(j+1)+",1,")
	#				if i==j:
	#					f.write(str(values[p][i,j]*stress_amp)+"\n")
	#				else:
	#					f.write(str(values[p][i,j]*stress_amp*2.0)+"\n")
	#		f.write("*Boundary\n") #prevent rigid body motion
	#		for i in range(dimens):
	#			f.write("Sample.Set-1,"+str(i+1)+","+str(i+1)+"\n")
	#		f.write("*End Step\n")
	#f.close()
	
	#Simulation of the Trainingdirections with UMAT Driver
	
	with open(jobname+'.inp',"w") as f:
		######################## basic informations ##################################
		f.write('*Material\nMonolithFE2\n1\n1\n')
		f.write('*Depvar\n1\n')
		if dimens==2:
			f.write('*NDI\n2\n*NSHR\n1\n')
		else:
			f.write('*NDI\n3\n*NSHR\n3\n')
		#now define the steps
		for i in range(len(values)):
			if (training_mode=="stress" or training_mode=="stress and strain"):
				f.write('*Step\n*name\nTrainingstep\n')
				f.write('*Nlgeom\n'+nlgeom+'\n')
				f.write('*Static\n0.01,'+str(simulation_time)+',0.00001,'+str(dtime_data_dump)+'\n')
				f.write('*Stress\n1.0\n')
				if dimens==2:
					f.write(str(values[i][0,0]*stress_amp)+','+str(values[i][1,1]*stress_amp)+','+str(values[i][0,1]*stress_amp)+"\n")
				else:
					f.write(str(values[i][0,0]*stress_amp)+','+str(values[i][1,1]*stress_amp)+','+str(values[i][2,2]*stress_amp)+','+str(values[i][0,1]*stress_amp)+','+str(values[i][0,2]*stress_amp)+','+str(values[i][1,2]*stress_amp)+"\n")
				f.write('*end_Step\n')
			if (training_mode=="strain" or training_mode=="stress and strain"):
				f.write('*Step\n*name\nTrainingstep\n')
				f.write('*Nlgeom\n'+nlgeom+'\n')
				f.write('*Static\n0.01,'+str(simulation_time)+',0.00001,'+str(dtime_data_dump)+'\n')
				if nlgeom.upper()=="NO":
					f.write('*STRAN\n1.0\n')
					if dimens==2:
						f.write(str(values[i][0,0]*strain_amp)+','+str(values[i][1,1]*strain_amp)+','+str(2.0*values[i][0,1]*strain_amp)+"\n")
					else:
						f.write(str(values[i][0,0]*strain_amp)+','+str(values[i][1,1]*strain_amp)+','+str(values[i][2,2]*strain_amp)+','+str(2.0*values[i][0,1]*strain_amp)+','+str(2.0*values[i][0,2]*strain_amp)+','+str(2.0*values[i][1,2]*strain_amp)+"\n")
				else:
					f.write('*DFGRD\n1.0\n')
					if dimens==2:
						f.write(str(values[i][0,0]*strain_amp)+','+str(values[i][1,0]*strain_amp)+",0.0,"+str(values[i][0,1]*strain_amp)+','+str(values[i][1,1]*strain_amp)+',0.0,0.0,0.0,0.0\n')
					else:
						f.write(str(values[i][0,0]*strain_amp)+','+str(values[i][1,0]*strain_amp)+','+str(values[i][2,0]*strain_amp)+','+str(values[i][0,1]*strain_amp)+','+str(values[i][1,1]*strain_amp)+','+str(values[i][2,1]*strain_amp)+','+str(values[i][0,2]*strain_amp)+','+str(values[i][1,2]*strain_amp)+','+str(values[i][2,2]*strain_amp)+'\n')
				f.write('*end_Step\n')
	
	if Method=="ROM":
		reduction='full'
	else:
		reduction='reduced'
	
	#set the Analysisparameters so that the training data is generated
	set_FE2_Analysisparameters('staggered',0.00001,12,'notsaved','indefinite','symm',reduction,
	                           training_data='ROM',simulation_time=simulation_time,dtime_data_dump=dtime_data_dump)
	#start the training simulations
	if start_simulations:
		os.chdir(os.sep.join(program_directory.split(os.sep)[:-1]))
		subprocess.Popen(['.'+os.sep+program_directory.split(os.sep)[-1],"cpus="+str(ncpus),"job="+jobname])
	
if __name__=='__main__':
	
	#if the script is started directly from terminal, call evaluate_training_data with the supplied arguments
	
	keyword_dictionary={"strain_amp":0.0,
	                    "stress_amp":0.0,
	                    "n_variations":1,
	                    "dimens":2,
	                    "nlgeom":"NO",
	                    "simulation_time":1.0,
	                    "dtime_data_dump":0.1,
	                    "training_mode":"strain",
	                    "jobname":None,
	                    "program_directory":None,
	                    "start_simulations":False,
	                    "ncpus":1,
	                    "Method":"ROM"}
	
	for i in range(1,len(sys.argv)):
		KW=sys.argv[i].split("=")
		if not(KW[0] in keyword_dictionary):
			raise Exception("Keyword"+KW[0]+" misspelled.")
		else:
			keyword_dictionary[KW[0]]=KW[1]
	
	generate_training_inputfile(float(keyword_dictionary["strain_amp"]),float(keyword_dictionary["stress_amp"]),int(keyword_dictionary["n_variations"]),
	                            int(keyword_dictionary["dimens"]),keyword_dictionary["jobname"],keyword_dictionary["nlgeom"],keyword_dictionary["Method"],
	                            float(keyword_dictionary["simulation_time"]),float(keyword_dictionary["dtime_data_dump"]),keyword_dictionary["training_mode"],
	                            bool(int(keyword_dictionary["start_simulations"])),keyword_dictionary["program_directory"],int(keyword_dictionary["ncpus"]))
