#*************************************************************************
#  Copyright (C) 2022 by Bruno Chareyre                                  *
#  bruno.chareyre_at_grenoble-inp.fr                                     *
#*************************************************************************/

## This script is a variant of examples/triax-tutorial/session1.py showing how to include capillary forces in a simulation using CapillarityEngine.
## For more documentation (beyhond capillarity), find documentation in the original script session1.py.

from yade import pack

nRead = readParamsFromTable(
        num_spheres=1000,  # number of spheres
        compFricDegree=6,  # initial contact friction for the confining phase
        finalFricDegree=30,  # true contact friction for actual shear phase,
        targetPorosity=0.40,
        key='_triax_base_',  # put you simulation's name here
        unknownOk=True
)
from yade.params import table

num_spheres = table.num_spheres  # number of spheres
key = table.key
targetPorosity = table.targetPorosity  #the porosity we want for the packing
compFricDegree = table.compFricDegree  # initial contact friction during the confining phase (will be decreased during the REFD compaction process)
finalFricDegree = table.finalFricDegree  # contact friction during the deviatoric loading
rate = -0.01  # loading rate (strain rate)
damp = 0.2  # damping coefficient
stabilityThreshold = 0.01  # we test unbalancedForce against this value in different loops (see below)
young = 5e6  # contact stiffnessO.bodies
mn, mx = Vector3(0, 0, 0), Vector3(1, 1, 1)  # corners of the initial packing

## create materials for spheres and plates
O.materials.append(FrictMat(young=young, poisson=0.5, frictionAngle=radians(compFricDegree), density=2600, label='spheres'))
O.materials.append(FrictMat(young=young, poisson=0.5, frictionAngle=0, density=0, label='walls'))

## create walls around the packing
walls = aabbWalls([mn, mx], thickness=0, material='walls')
wallIds = O.bodies.append(walls)

## use a SpherePack object to generate a random loose particles packing
sp = pack.SpherePack()
sp.makeCloud(mn, mx, -1, 0.5, num_spheres, False, 0.95, seed=1)  #"seed" make the "random" generation always the same
sp.toSimulation(material='spheres')

triax = TriaxialStressController(
        maxMultiplier=1. + 2e4 / young,  # spheres growing factor (fast growth)
        finalMaxMultiplier=1. + 2e3 / young,  # spheres growing factor (slow growth)
        thickness=0,
        stressMask=7,
        internalCompaction=True,  # If true the confining pressure is generated by growing particles
)

newton = NewtonIntegrator(damping=damp)

O.engines = [
        ForceResetter(),
        InsertionSortCollider([Bo1_Sphere_Aabb(aabbEnlargeFactor=1.2), Bo1_Box_Aabb()]),
        InteractionLoop(
                [Ig2_Sphere_Sphere_ScGeom(interactionDetectionFactor=1.2),
                 Ig2_Box_Sphere_ScGeom(interactionDetectionFactor=1.2)], [Ip2_FrictMat_FrictMat_CapillaryPhysDelaunay()],
                [Law2_ScGeom_FrictPhys_CundallStrack(neverErase=True)]
        ),
        GlobalStiffnessTimeStepper(active=1, timeStepUpdateInterval=100, timestepSafetyCoefficient=0.5),

        ######################     INSERT CAPILLARITY ENGINE HERE ############################################
        CapillarityEngine(capillaryPressure=10000, liquidTension=100, dead=True, label="capillarity"),
        ######################################################################################################
        triax,
        newton,
        PyRunner(iterPeriod=20, command='history()', label='recorder'),
        PyRunner(iterPeriod=500, command='transitions()')
]

Gl1_Sphere.stripes = 1
#if nRead==0: yade.qt.Controller(), yade.qt.View()

triax.goal1 = triax.goal2 = triax.goal3 = -10000

phase = 0
epsv_0 = 0
e22_0 = 0
doShear = True


def transitions():
	global phase, compFricDegree, epsv_0, e22_0
	if phase == 0:
		if unbalancedForce() < stabilityThreshold and abs(-10000 - triax.meanStress) / 10000 < 0.001:
			phase += 1
	if phase == 1:
		if triax.porosity > targetPorosity:  # reduce friction to reach target porosity
			compFricDegree = 0.97 * compFricDegree
			setContactFriction(radians(compFricDegree))
			triax.internalCompaction = False
			print(
			        "\r Friction: ",
			        compFricDegree,
			        " porosity:",
			        triax.porosity,
			)
		else:
			if unbalancedForce() < (0.5 * stabilityThreshold):
				print("start capillary forces")
				capillarity.dead = False
				setContactFriction(radians(finalFricDegree))
				phase += 1
			else:
				print("continue stabilizing:", unbalancedForce())
	if phase == 2 and doShear:
		if unbalancedForce() < 0.25 * stabilityThreshold:
			print("start shear")
			triax.stressMask = 5
			triax.goal2 = rate
			triax.goal1 = -10000
			triax.goal3 = -10000
			newton.damping = 0.1
			epsv_0 = -triax.strain[0] - triax.strain[1] - triax.strain[2]
			e22_0 = -triax.strain[1] - e22_0
			phase += 1
			O.timingEnabled = True
		else:
			print("stabilize with bridges")


from yade import plot


def history():
	plot.addData(
	        e11=-triax.strain[0],
	        e22=-triax.strain[1] - e22_0,
	        e33=-triax.strain[2],
	        ev=-triax.strain[0] - triax.strain[1] - triax.strain[2] - epsv_0,
	        s11=-triax.stress(triax.wall_right_id)[0],
	        s22=-triax.stress(triax.wall_top_id)[1],
	        s33=-triax.stress(triax.wall_front_id)[2],
	        i=O.iter
	)


plot.plots = {'i': ('s11', 's22', 's33', None, 'ev'), 'e22': ('s11', 's22', 's33', None, 'ev')}

O.run(40000, False)
plot.plot()
#from yade import timing
#timing.stats()
