import os
from yade import pack

dSnorm = 5000.0 #norm of the stress increment
nbProbes = 36 #number of stress directions tested
rampIte = 20 #nb iterations to increase the stress state until the final desired stress value
stabIte = 5000 #nb iterations to stabilize sample after reaching the final stress value
key='_probing_envelope_'

stabilityThreshold=0.001 # leave it too low can have a super stable state, but will eat onion for waiting so long


triax=TriaxialStressController()
stickLoad=triax.strain[1]
REC=TriaxialStateRecorder()
# an array for saving stress increments and strain responses; arrays are in "numpy" extension
import numpy
probings=numpy.zeros((3,nbProbes))

def increment(dsr=0,dsa=1):
#	rampIte = 20
	for ite in range(rampIte):# progressivaly increase of stress state
		O.run(20, True)
		#incrementation of stress state
		triax.goal2 = initSa+dsa/rampIte*ite
		triax.goal1 = triax.goal3 = initSr+dsr/rampIte*ite
		print triax.goal1, triax.goal2, '(stress value)'

	# fix the stress value for stabilization at the final state
	#REC.file=('WallStresses'+key2)
	triax.goal2 = initSa+dsa
	triax.goal1 = triax.goal3 = initSr+dsr

	while 1:
		O.run(100, True)
		unb=unbalancedForce()
		print 'unbalanced force:',unb,' strain: ',triax.strain
		O.save('checkpoint.xml') # reload this file to check if we see mistakes
		if unb<stabilityThreshold: break



# loop over all the stress directions
for i in range(nbProbes):

	# computation of the stress direction of the current stress probe
	alphaS = 2*pi/nbProbes*(i-1)
	print 'stress probe nb:',i,' stress direction (deg): ',degrees(alphaS)

	# computation of the stress increment in the axial direction
	dSa = dSnorm*sin(alphaS) 

	# computation of the stress increment in the radial direction
	dSr = dSnorm*cos(alphaS)/sqrt(2.0)

	#Load the initial anisotropic state before running a new stress probe
	O.load('firstpoint.xml')
	#key2='_DIR%d'%i
	#REC.file=('WallStresses'+key2)
	triax.stressMask=7
	triax.goal2=stickLoad
	triax.goal1=triax.goal3=100000
	#We redefine the "triax" label, else it would point to inactive engine from previous simulation that is still in memory
	triax=O.engines[4]

	initSa=triax.goal2  #save of the initial axial stress
	initSr=triax.goal1  #save of the initial radial stress

	# define the final stress state to be reached
	finalSa = initSa+dSa
	finalSr = initSr+dSr

	#... need to active stress control in 3 directions if not yet done
#	triax.stressControl_1=triax.stressControl_2=triax.stressControl_3=True
	triax.stressMask=7  # Hien: the previous is 7, that means the stress control is on, but the following code is focus on strain control so I turn this off

	# fix a high value of maximum strain rate, the progressive loading will be done by progressively increasing the desired stress state at each iteration
#	triax.strainRate1=triax.strainRate2=triax.strainRate3=1000.0
#	triax.goal1=triax.goal2=triax.goal3=1000.0 #!!!!!!
	increment(dSr,dSa)
	O.wait()
#	O.save('Direction_%d.xml'%(i-1))
	os.rename('WallStresses_probing_envelope_','Dir_%d'%(i-1))