# -*- encoding=utf-8 -*-
# 2020 © Vasileios Angelidakis <v.angelidakis2@ncl.ac.uk>
# 2020 © Bruno Chareyre <bruno.chareyre@grenoble-inp.fr> 

# Benchmark of basic performance of open-source DEM simulation systems, 
# Case 1: Silo flow
# this version uses custom positions for spheres, and custom contact properties, change 

particleRadius = 8 # 2mm in the benchmark, bigger for less particles and quicker simulations
numThreads=4
reportTiming=False

# Configure MPI module if needed
mpi = 'MPI' in yade.config.features
if mpi:
    from yade import mpy as mp
else:
    print("yade is compiled without MPI support, numThreads>1 ignored")
    numThreads = 4

# The stl file is in mm -> We can keep the units below or scale the stl file to meters (ymport.stl does not have a scale parameter, but I can edit the .stl file manually)

# -------------------------------------------------------------------- #
# Units: N/mm/MPa/ton
# 	force:   1 N
# 	length:  1 mm = 1e-3 m
# 	stress:  1 MPa = 1 N/mm^2 = 1e6 Pa = 1e-3 GPa
#	mass:    1 ton = 1e3 kg
#	density: 1 ton/mm^3 = 1e-12 kg/m^3

# -------------------------------------------------------------------- #
# Input Data -> Define Material and Orifice size. Uncomment the prefered choice

#FIXME: most of what follows don't have to be guarded by rank==0
fileName='SiloLargeOrifice'
#fileName='SiloSmallOrifice'
if fileName=='SiloLargeOrifice':
    z=70	# This is the height of the lowest point of the funnel (at the orifice), measuring from the lowest cylindrical cross section of the silo
elif fileName=='SiloSmallOrifice':
    z=80




granularMaterial='M1'
#granularMaterial='M2'

# -------------------------------------------------------------------- #
# Materials
Steel = O.materials.append(FrictMat(young=210e-3,poisson=0.2,density=7200e-12,label='Steel'))

# -------------------------------------------------------------------- #
# Asign values based on the Material and Orifice size
# Coeff of restitution (e) / Coeff of friction (f)
e_M1_M2=0.45;	f_M1_M2=0.2
e_M1_M1=0.5;	f_M1_M1=0.3
e_M1_St=0.4;	f_M1_St=0.2
e_M2_M2=0.4;	f_M2_M2=0.4
e_M2_St=0.4;	f_M2_St=0.2

if granularMaterial=='M1':
    M1=O.materials.append(FrictMat(young=1.0e-3,poisson=0.2,density=2500e-12,label='M1'))
    e_gg=e_M1_M1	# Coefficient of restitution (e) between granular material (g) and granular material (g)
    f_gg=f_M1_M1	# Coefficient of friction (f)...
    e_gs=e_M1_St	# Coefficient of restitution (e) between granular material (g) and steel (s)
    f_gs=f_M1_St	# Coefficient of friction (f)...
elif granularMaterial=='M2':
    M2=O.materials.append(FrictMat(young=0.5e-3,poisson=0.2,density=2000e-12,label='M2'))
    e_gg=e_M2_M2
    f_gg=f_M2_M2
    e_gs=e_M2_St
    f_gs=f_M2_St

F_gg=atan(f_gg) # Friction Angle between granular material (g) and granular material (g)
F_gs=atan(f_gs) # Friction Angle between granular material (g) and steel (s)


# -------------------------------------------------------------------- #
## Engines 
O.engines=[
    ForceResetter(),
    InsertionSortCollider([Bo1_Sphere_Aabb(),Bo1_Facet_Aabb()],label="collider"),
    InteractionLoop(
        [Ig2_Sphere_Sphere_ScGeom(), Ig2_Facet_Sphere_ScGeom()],
        [Ip2_FrictMat_FrictMat_MindlinPhys(
            frictAngle = MatchMaker(matches=((1,1,F_gg),(0,1,F_gs))), # 0 being the id of Steel and
            en         = MatchMaker(matches=((1,1,e_gg),(0,1,e_gs)))  # 1 being the id of granularMaterial
        )],
        [Law2_ScGeom_MindlinPhys_Mindlin()],
    ),
    NewtonIntegrator(damping=0,gravity=[0,0,-9.810],label="newton"),
    #GlobalStiffnessTimeStepper(active=1,timestepSafetyCoefficient=0.8, timeStepUpdateInterval=100, parallelMode=False, label = "ts",defaultDt=PWaveTimeStep()), #FIXME Remember to reinstate parallelMode=True when we use MPI
    #VTKRecorder(virtPeriod=0.04,fileName='/tmp/Silo-',recorders=['spheres','facets']),
]


# This condition is not abolutely necessary but it would be inelegant to
# download *.stl and generate densePack N times when we need it done only on master (centralized scene method)
if not mpi or mp.rank==0:

    # -------------------------------------------------------------------- #
    # Generate initial packing. Choose among regularOrtho, regularHexa or randomDensePack (which I think is best).

    ## Using regularOrtho
    #sp=pack.regularOrtho(pack.inCylinder(Vector3(0,0,0),Vector3(0,0,305),radius),radius=2,gap=r*1/10.,material=granularMaterial)

    ## Using regularHexa
    #sp=pack.regularHexa(pack.inCylinder(Vector3(0,0,0),Vector3(0,0,215),radius),radius=2,gap=r*1/10.,material=granularMaterial)  

    # Using randomDensePack
    sp=pack.randomDensePack(pack.inCylinder((0,0,0),(0,0,254),100),radius=particleRadius,spheresInCell=500,returnSpherePack=False,material=granularMaterial,seed=1) 

    # -------------------------------------------------------------------- #
    # Sort packing in ascending Z coordinates and delete excess particles to achieve sample size of 122k
    zValues=[]
    for s in sp:
        zValues.append(s.state.pos[2])

    from operator import itemgetter
    indices, zValues_sorted = zip(*sorted(enumerate(zValues), key=itemgetter(1)))
    list(zValues)
    list(indices)

    sp_new=[]
    for i in range(0,len(sp)):
        sp_new.append(sp[indices[i]])

    Nspheres=122000
    sp_new=sp_new[0:Nspheres]

    from yade import ymport
    if not os.path.exists(fileName+'.stl'):
        print("Downloading mesh file")
        try:
            os.system('wget http://yade-dem.org/publi/data/DEM8/'+fileName+'.stl')
        except:
            print("** probably no internet connection, grab the *.stl files by yourself **")
    facets = ymport.stl(fileName+'.stl',color=(0,1,0),material=Steel)
    fctIds = range(len(facets))
    
    O.bodies.append(facets)
    O.bodies.append(sp_new)
        
    # -------------------------------------------------------------------- #
    # Count the number of spherical particles to verify sample size. We can comment this out later on.
    numSpheres=0
    for b in O.bodies:
        if isinstance(b.shape,Sphere):
            numSpheres=numSpheres+1
    print('The total number of spheres is: ',numSpheres)
        
    collider.verletDist = 0.6*particleRadius
    O.dt=0.6 * PWaveTimeStep()
    O.dynDt=False

# -------------------------------------------------------------------- #
# Erase particles flowing out of the silo

def eraseEscapedParticles():
    global numErased
    count=0
    ts=time.time()
    ers=[]
    for b in O.bodies:
        if isinstance(b.shape,Sphere) and b.state.pos[2]<-z-20: # I do not delete the particles right after they pass the orifice, to disturb the simulation as little as possible
            ers.append(b.id)
            count+=1            
    if mpi:
         mp.bodyErase(ers)
    else:
        for b in ers: O.bodies.erase(b.id)
    numErased+=count
    
#-------------------------------------------------------------------- #
#Record time-dependent number of retained particles and vtk export

from yade import plot
plot.plots={'time':(('retained','bo--'),None,('Cu',"kx--"))}

numErased = 0
def addPlotData(Cu): 
	plot.addData(retained=numSpheres-numErased, time=O.time, Cu=Cu)

from yade import export
vtk = export.VTKExporter("spheresFinal")


# -------------------------------------------------------------------- #
# Run iterations

if mpi: # import and tune MPI module    
    mp.DOMAIN_DECOMPOSITION =True
    mp.ACCUMULATE_FORCES=False
    mp.MERGE_W_INTERACTIONS=False
    mp.REALLOCATE_MINIMAL = False # if true, intersections are minimized before reallocations, hence minimizing the number of reallocated bodies
    mp.REALLOCATE_FREQUENCY=20
    mp.USE_CPP_REALLOC=True
    mp.MINIMAL_INTERSECTIONS = True
    mp.MASTER_UPDATE_STATES = False
    mp.ERASE_REMOTE_MASTER = True
    mp.YADE_TIMING=reportTiming
else:
    O.timingEnabled=reportTiming

substeps=500
#while len(O.bodies)-numErased>0 or len(O.bodies)==0:
for k in range(100):
    t1=time.time()
    if mpi:
        mp.mpirun(substeps,numThreads,withMerge=True) # if numThreads=1 this will fall-back to normal O.run() and mp.rank=0
    else:
        O.run(substeps,True)
    tbf=time.time()
    eraseEscapedParticles()
    if mpi and mp.rank>0: continue # mpi workers do not record
    t2=time.time()
    addPlotData((numSpheres-numErased)*substeps/(t2-t1))
    vtk.exportSpheres(what=dict(particleVelocity='b.state.vel',domain='b.subdomain'))
    plot.plot(noShow=True).savefig(fileName+'_'+granularMaterial+'_np'+str(numThreads)+'.png')
    plot.saveDataTxt(fileName+'_'+granularMaterial+'.txt')
    print("iter=",O.iter,", last substep erased", numErased,"in",t2-t1,"s")

# -------------------------------------------------------------------- #
# GUI
if opts.nogui==False:
	from yade import qt
	v=qt.View()

	v.eyePosition = Vector3(0,-600,100)
	v.upVector    = Vector3(0,0,1)
	v.viewDir     = Vector3(0,1,0)
#	v.grid=(False,True,False)

	rndr=yade.qt.Renderer()
	#rndr.shape=False
	#rndr.bound=True

## To play interactively with mpi execution:
## mp.mpirun(100,numThreads,True) #'True' so we see merged scene after the run
## eraseEscapedParticles()
## mp.mpirun(100,numThreads,True)
## etc.
