# -*- encoding=utf-8 -*-
################################################################################
# 1) Generate a packing of particles inside an axis-aligned bounding box using 
# makeCloud(). 2) Use triaxialStressController() with internal confinement to
# achieve an average isotropic stress on the aabb. 3) Finally, confine the
# aggregate by applying an isotropic Cauchy stress tensor on the model's
# power diagram described by the alpha shape surface reconstruction algorithm 
################################################################################
from yade import pack
#######################################
###       FUNCTION DEFINITIONS      ###
#######################################
def averageVelGrad(xq,Aq,vol0): # Compute Velocity Gradient of the Aggregate
  velGrad = Matrix3.Zero
  for i, el in enumerate(Aq):
    velGrad = velGrad + Vector3.outer(xq[i],Aq[i]) / vol0
  return velGrad
  
def triStressMicro(stressPt,triThresh,creepThresh): # Apply Uniform Stress Tensor
  while 1:	
    O.run(1000, True) 
    meanStress = TW.calcAlphaStress(alpha,alphaShrinked,True)
    triDiff = stressPt.diagonal() - meanStress.diagonal()
    caps=TW.getAlphaCaps(alpha,alphaShrinked,True)
    velArr = [O.bodies[cap[0]].state.vel for cap in caps]
    areaArr = [cap[1] for cap in caps]
    lbar = averageVelGrad(velArr,areaArr,TW.alphaCapsVol)
    meanBVel = 0.5 * lbar.spectralDecomposition()[1].maxAbsCoeff()
    integrator.damping=max(0.2,min(0.99,(meanBVel/boundVel)**4))
    print('Boundary Velocity:',meanBVel,'Stress Diagonal',meanStress.diagonal())
    if meanBVel < abs(creepThresh * boundVel) and abs(triDiff[0]/stressPt[0][0]) < triThresh and abs(triDiff[1]/stressPt[1][1]) < triThresh and abs(triDiff[2]/stressPt[2][2]) < triThresh:
      break
    print('damping = ',integrator.damping)

def isoLoad(stabThresh,triThresh): # Use TriaxialStressController for Initial Confinement
  while 1:	
    O.run(1000, True) 
    unb=unbalancedForce()
    triDiff=abs((triax.goal1-triax.meanStress)/(triax.goal1))
    print('unb:',unb,' triDiff: ',triDiff,' meanS: ',triax.meanStress, ' porosity: ',triax.porosity)
    if unb<stabThresh and triDiff<triThresh:
      break
############################################
###   DEFINING VARIABLES AND MATERIALS   ###
############################################
confinement_p = -5.0e6
damp=0.2 # damping coefficient
fricDegree=35.0

young=160.0e9 # contact stiffness
poisson=0.17

sphereMat=FrictMat( young=young,
                    poisson=poisson,
                    frictionAngle=radians(fricDegree),
                    density=2720,
                    label='spheres')

wallMat=FrictMat( young=young,
                  poisson=poisson,
                  frictionAngle=0,
                  density=0,
                  label='walls')
#######################################
###            MAKE THE PACK       ###
#######################################
num_spheres=1000   # number of spheres
mn,mx=Vector3(0,0,0),Vector3.Ones # corners of the packing

sp=pack.SpherePack()
sp.makeCloud(mn,mx,num=num_spheres,rMean = 0.04,rRelFuzz=0.5,
				distributeMass=True,seed=1)
#######################################
###          DEFINE THE ENGINES     ###
#######################################
engineList = [
  ForceResetter(),
  InsertionSortCollider([ Bo1_Sphere_Aabb(),
                          Bo1_Box_Aabb()],
                          label="sorter"),
  InteractionLoop(
    [ Ig2_Sphere_Sphere_ScGeom(),
      Ig2_Facet_Sphere_ScGeom(),
      Ig2_Box_Sphere_ScGeom()],
    [Ip2_FrictMat_FrictMat_MindlinPhys( label='mindlinContact')],
    [Law2_ScGeom_MindlinPhys_Mindlin( label="mindlinLaw2")],),
  GlobalStiffnessTimeStepper( active=1,timeStepUpdateInterval=100,
                              timestepSafetyCoefficient=0.2, label='stepper'),
  NewtonIntegrator(damping=damp, label='integrator'),
  TriaxialStressController(label='triax'),
  PyRunner(iterPeriod=0,
            command="TW.applyAlphaForces(stressM,alpha,alphaShrinked,True)",
            label='loader')]
##########################################
### INTERNAL COMPACTION TO CONFINEMENT ###
##########################################
O.materials.append(sphereMat)
O.materials.append(wallMat)
walls=aabbWalls([mn,mx],thickness=0,material='walls')
wallIds=O.bodies.append(walls)
sp.toSimulation(material='spheres')
O.engines=engineList
triax.goal1=triax.goal2=triax.goal3=confinement_p
isoLoad(0.01,0.01)
#######################################
###   APPLY UNIFORM STRESS TENSOR   ###
#######################################
sp.fromSimulation()
O.resetThisScene()
O.materials.append(sphereMat)
sp.toSimulation(material='spheres')

engineList.remove(triax)
O.engines=engineList

stepper.timestepSafetyCoefficient=0.6
stressM = Matrix3.Identity * (confinement_p)
minRad = min(b.shape.radius for b in O.bodies)
meanDia = 2. * sum([b.shape.radius for b in O.bodies]) / len(O.bodies)
alpha = (22. * minRad) ** 2
alphaShrinked = (21. * minRad) ** 2
loader.iterPeriod=2
TW=TesselationWrapper()
integrator.damping=0.5
iNum = 7.9e-5       # Inertial Number Target
boundVel = 0.45 * iNum * sqrt(5.e6 /O.bodies[0].mat.density) / meanDia
print('Quasistatic Strain Rate', boundVel)
triStressMicro(stressM,0.01,0.1)
#######################################
###        SAVE CONFINED GZIP       ###
#######################################
O.save('confinedState.yade.gz')
print("###      Confining Pressure Achieved      ###")
#######################################
### VTK SPHERES AND POWER DIAGRAM   ###
#######################################
import yade.export
vtkExporter = yade.export.VTKExporter('vtkConfined')
vtkExporter.exportSpheres()

caps=TW.getAlphaCaps(alpha,alphaShrinked,True)
segments=TW.getAlphaGraph(alpha,alphaShrinked,True)
f=open('LVdiagram.vtk','w')
f.write("# vtk DataFile Version 2.0\nGenerated from YADE\n")
f.write("ASCII\nDATASET POLYDATA\n")
f.write("POINTS %d float \n" %(len(segments)))
[f.write("%f %f %f\n" %(point[0],point[1],point[2])) for point in segments]
f.write("LINES %d %d \n" %(len(segments)/2, 3*len(segments)/2))
[f.write("2 %d %d\n" %(2*k,2*k+1)) for k in range(int(len(segments)/2))]
f.close()
