import numpy as np
import pyvista as pv
from scipy.spatial.transform import Rotation as R

# Provide user inputs
# vtk file options
fn                     = '1matPEvbin' # File name and path for the .vtk to read from
scale_factor           = 1. # Multiplicative scaling factor that converts voxel dimension to mm
EulerAngleKey          = 'EulerAngles' # Variable name in vtk file corresponding to Euler angle
GrainKey               = 'FeatureIds' # Variable name in vtk file corresponding to grain ID

# Ansys expects Euler angles in degrees and in Z-X-Z order. Valid angle ranges are:
# Z1: -180 ~ 180
# X: 0 ~ 180
# Z2: -180 ~ 180
# If the Euler angles in the vtk file are in any other different formats, use the
# following function to convert to the Ansys format
def FormatEA( angles ):
	tmp = angles / np.pi * 180.
	tmp[0] -= 180
	tmp[-1] -= 180
	return tmp


# Elastic behavior
# Available options are: 
# 'CUBIC', Cubic elasticity, needs C11, C12 and C44
Elastic_type = 'CUBIC'
Elastic_constants = [ 1.067e5 , 0.604e5 , 0.283e5 ]

# Thermomechanical behavior
CTE                    = 4.473e-5 # Coefficient of thermal expansion

# Crystal plasticity parameters
# Nonuniform hardness distribution
NONUNIF = True
L_c = 2 # Characteristic length
delta_g = 50.
N_bins = 8


# Detailed explanation of the material parameters can be found in Section 4.8. Crystal Plasticity, of the Ansys help document
# Crystal type, available options are:
# 'FCC' , 'HCP' , 'BCC'
CrystalType            = 'FCC'
ini_hardness           = 220. - 0.38 * delta_g # Initial slip resistance
hardening_modulus      = 1750. # Linear hardening modulus
sat_hardness           = 400. - 0.38 * delta_g # Saturation hardness
r                      = 2 # Athermal hardness modulus sensitivity, Eq (4-73)
n                      = 0 # Saturation hardness sensitivity, Eq (4-74)
T_ref                  = 298. # Reference temperature, K
latent_hardening_ratio = 1.4 # Latent hardening ratio, Eq (4-71)
gamma_dot_0            = 1732000.0 # Reference slip rate
Q_s                    = 2.5e-19 # Activation energy, J
p                      = 0.131 # Flow rule exponent 1, Eq (4-67)
q                      = 1.1 # Flow rule exponent 2, Eq (4-67)
m                      = 0.005 # Rate dependence parameter, for HCP material only, Eq (4-66)
thermal_athermal_ratio = 0.7 # Initial ratio of thermal to athermal slip resistance
C_2_A_ratio            = 1 # Lattice C by A ratio



# End of user inputs
##############################################################################################

# Sub routines for non-uniform hardness distribution
def dis2ratio( dist , characteristic_length ):
	return 1. - np.tanh( dist / characteristic_length )
# def toBins( y_vals , characteristic_length ):
# 	return np.arctanh( 1 - y_vals ) * characteristic_length

def GetHardness( idx_3d , dims , my_gid , Grains , my_ea , slip_directions , L_character ):
	# Convert EA to rotation mtx
	Q = R.from_euler( 'ZXZ' , my_ea , degrees=True ).as_matrix()

	# Get voxel increments
	vox_increment = np.rint( np.einsum( 'ij,kj->ki' , Q , slip_directions ) ).astype(int)	
	inc_length = np.linalg.norm( vox_increment , axis = 1 )

	dist = np.zeros( [ len(vox_increment) , 2 ] )
	for did , direction in enumerate( [ -1 , 1 ] ): # Go in positive and negative directions
		for sid , step in enumerate( vox_increment ): # March in all slip directions
			start = idx_3d.copy(); gb_dis = 0.
			while 1:
				start += direction * step
				old_dis = gb_dis
				try:
					new_voxel = Grains[ np.ravel_multi_index( start , dims , order='F' ) ]
					gb_dis += inc_length[ sid ]
				except:
					pass
				if gb_dis == old_dis or np.any( start < 0 ):
					# OOB, hit free surface
					gb_dis = 1e6
					break
				if new_voxel != gid:
					gb_dis -= inc_length[ sid ] # Back track 1 step
					break
			dist[ sid , did ] = gb_dis

	# Take closer of the two directions
	min_gb_dis = np.min( dist , axis = 1 )
	# Return average over all slip directions
	return np.mean( dis2ratio( min_gb_dis , L_character ) ) 


# Unpack user inputs
CrystalMap = {'FCC':1 , 'HCP':2 , 'BCC':3 } 
CryType = CrystalMap[ CrystalType ]
if CryType == 1: # FCC
	nSS = 12 # Number of slip systems
	nSF = 1 # Number of slip families
	flow_paras = [ gamma_dot_0 , p , q , thermal_athermal_ratio , Q_s , C_2_A_ratio ]
	formulation = 2 # Exponential, Eq (4-67)

	slip_directions = np.array([[1,-1,0],\
								[-1,0,1],\
								[0,1,-1],\
								[1,0,1],\
								[-1,-1,0],\
								[0,-1,-1]]) / np.sqrt(2) # Normalize

elif CryType == 2: # HCP
	nSS = 30
	nSF = 5
	flow_paras = [ gamma_dot_0 , m , C_2_A_ratio ]
	formulation = 1 # Power law, Eq (4-66)

elif CryType == 3: # BCC
	nSS = 48
	nSF = 3
	flow_paras = [ gamma_dot_0 , p , q , thermal_athermal_ratio , Q_s , C_2_A_ratio ]
	formulation = 2 # Exponential, Eq (4-67)



# Step 1: Read vtk, write out inp file
vtk = pv.read( fn + '.vtk' )

# Basic geometry info
dims = np.array(vtk.dimensions) - 1 # Voxel dimension
length = dims * scale_factor # Physical length unit in WB is mm

# Reset origin
vtk.origin = np.zeros(3)
# Change vtk unit to mm
vtk_um = vtk.scale([scale_factor, scale_factor, scale_factor], inplace=False)
pv.save_meshio( 'vtk_bin_analysis.inp' , vtk_um )

# Add a placeholder material to the Abaqus inp
f = open('vtk_bin_analysis.inp','a')
f.write('*Elset, elset=M, internal, generate\n1,' + str(np.prod(dims)) + \
	', 1\n*Solid Section, elset=M, material=m1\n*Material, name=m1\n*Elastic\n1,0.3')



# Step 2: Read Euler angles in the vtk file, store and write element set
# Sanity check
if np.max( vtk[EulerAngleKey] ) <= 2.*np.pi:
	print('Euler angles in vtk are probably in radians! Please check if appropriate conversion function has been defined!')
element_set , Eulers , Hardness_dist = {} , {} , {}
y_bins = np.linspace( 0 , 1 , N_bins+2 )[1:-1]
# tmp = np.zeros(50**3)
for eidx , gid in enumerate( vtk[GrainKey] ):
	idx_3d = np.array(np.unravel_index( eidx , dims , 'F' ))
	eid = idx_3d[2] * dims[2] * dims[1] + idx_3d[1] * dims[0] + idx_3d[0] + 1

	if not NONUNIF: # If uniform initial hardness is used, group by grains
		if gid not in element_set.keys():
			element_set[ gid ] = [ eid ]
			# Orientation
			Eulers[ gid ] = FormatEA( vtk[EulerAngleKey][ eidx ] )
		else:
			element_set[ gid ].append( eid )
	else: # Non-uniform hardness is to be used, group by element
		# Orientation
		curr_EA = FormatEA( vtk[EulerAngleKey][ eidx ] )

		# Calculate normalized hardness distribution
		curr_hardness = GetHardness( idx_3d , dims , gid , vtk[GrainKey] , curr_EA , slip_directions , L_c )

		# Binning
		bin_idx = np.digitize( curr_hardness , y_bins )

		# Modify gid to include bin index
		gid = int(str(gid)+str(bin_idx))
		if gid not in element_set.keys():
			element_set[ gid ] = [ eid ]
			Hardness_dist[ gid ] = [ curr_hardness ]
			Eulers[ gid ] = curr_EA
		else:
			element_set[ gid ].append( eid )
			Hardness_dist[ gid ].append( curr_hardness )
# vtk['Hardness'] = tmp
# vtk.save(fn +'_Hardness.vtk')
# print( 'Mean hardness increment = ' , np.mean(vtk['Hardness']) )



# Write named selections
content = "\n! Begin writing named selections\n\n"
for curr_gid , grain_elems in element_set.items():
	curr_gid = int(curr_gid)

	content += "CMBLOCK,GRAIN" + str(curr_gid) + "_SET_EXTERNAL_MODEL_,ELEM,	   " + str(len(grain_elems)) + "\n(8i10)\n"
	for j , eid in enumerate( grain_elems ):
		eid = str(eid)
		fortran_formatted = " "*( 10 - len(eid) ) + eid
		content += fortran_formatted
		if (j+1) % 8 == 0 :
			content += "\n"
	content += "\n"
n_grains = len( Eulers.keys() )
content += "! Model contains " + str(n_grains) + " grains\n"



# Step 3: Write material assignments
# Convert elastic properties
if Elastic_type == 'CUBIC':
	assert( len(Elastic_constants) == 3 )
	C11 , C12 , C44 = Elastic_constants
	E_ = ( C11**2 + C12*C11 - 2*C12**2 ) / ( C11 + C12 )
	v_ = C12 / ( C11 + C12 )
	G_ = C44
elif Elastic_type == 'ISO':
	print('Please use the cubic elastic option to achieve isotropic elasticity')
	raise NotImplementedError

# Write APDL script for material assignment
content += "\n\n! Begin writing grain materials\n"
for matid in Eulers.keys():
	matid = int(matid)
	if matid == 0: # Voids
		content += "! Define void \n*SET,matid," + str(len(element_set)+1) + "\n"
		content += "MP,EX,matid,1e-6,\nMP,NUXY,matid,0.3,\n"
		content += "! Assign to void\nCMSEL,S,GRAIN0_SET_EXTERNAL_MODEL_\n"
		content += "EMODIF,ALL,MAT,matid\nALLSEL,ALL\n\n\n"
		continue
	content += "*SET,matid," + str(matid) + "\n"

	# 4.8.2.1. Crystal Elastic Properties
	if Elastic_type == 'ISO':	
		content += "TB,ELAS,matid,,2,\nTBDATA,1," + str(Elastic_constants[0]) + "," + str(Elastic_constants[1]) + "\n"
	else:
		content += "TB,ELAS,matid,,9,\nTBDATA,1," + str(E_) + "," + str(E_) + "," + str(E_) +\
			"," + str(2*G_) + "," + str(2*G_) + "," + str(2*G_) +\
			"\nTBDATA,7," + str(v_) + "," + str(v_) + "," + str(v_) + "\n"

	# 4.8.2.2. Orthogonal Coefficients
	content += "TB,CTE,matid,\nTBDATA,1," + str(CTE) + "," + str(CTE) + "," + str(CTE) + "\n"

	# Plasticity line, placeholder, does not affect results
	content += "TB,PLAS,matid\n" 
	
	# 4.8.2.3.1. Crystal Orientation Properties
	my_ea = Eulers[ matid ]
	content += "TB,XTAL,matid,,3,ORIE\nTBDATA,1," + str(my_ea[0]) + "," + str(my_ea[1]) + "," + str(my_ea[2]) + "\n"

	# 4.8.2.3.2. Number of Slip Families
	content += "TB,XTAL,matid,,1,NSLFAM\nTBDATA,1," + str(nSF) + "\n"

	# 4.8.2.3.3. Formulation
	content += "TB,XTAL,matid,,1,FORM\nTBDATA,1," + str(formulation) + "\n"

	# 4.8.2.3.4. Crystal Characteristics
	content += "TB,XTAL,matid,," + str( 5 + nSF ) + ",XPARAM\nTBDATA,1," + str(CryType) +\
		",0," + str(matid) + "," + str(nSS) + ",0"
	for sf in range(nSF):
		if sf + 5 + 1 == 7:
			content += "\nTBDATA,7" # Switch to a new line
		content += ",1" # Activate all slip families
	content += "\n"
	
	# 4.8.2.3.5. Hardness Characteristics
	content += "TB,XTAL,matid,1," + str(3*(nSF+1)) + ",HARD\nTBTEMP, " + str(T_ref) + ",\nTBDATA,1"
	if NONUNIF:
		avg_increase = np.mean( Hardness_dist[ matid ] )
		my_ini_hardness = ini_hardness + delta_g * avg_increase
		my_sat_hardness = sat_hardness + delta_g * avg_increase
	else:
		my_ini_hardness = ini_hardness
		my_sat_hardness = sat_hardness

	dat = [my_ini_hardness]*nSF + [hardening_modulus]*nSF + [my_sat_hardness]*nSF + [ r , n , latent_hardening_ratio ]
	for i , val in enumerate( dat ):
		if ( i ) % 6 == 0 and i > 0:
			mult = ( i ) // 6
			content += "\nTBDATA," + str(6*mult+1) # Switch to a new line
		content += "," + str( val )
	content += "\n"

	# 4.8.2.3.6. Flow Characteristics
	content += "TB,XTAL,matid,1," + str(len(flow_paras)) + ",FL"+CrystalType+"\nTBTEMP, " + str(T_ref) + ",\nTBDATA,1"
	for fp in flow_paras:
		content += "," + str(fp)
	content += "\n"

	# Material assignment
	content += "CMSEL,S,GRAIN" + str(matid) + "_SET_EXTERNAL_MODEL_\n"
	content += "EMODIF,ALL,MAT,matid\nALLSEL,ALL\n\n\n"

# Set reference temperature
content += 'allsel,all\nTREF, '+str(T_ref)+'\nTUNIF, '+str(T_ref)+'\n'

# Write file
f = open('./mCP_8bin.dat' , 'w')
f.write( content )
f.close() 
