# -*- coding: utf-8

from __future__ import print_function
from builtins import range
from yade import pack,export,qt,geom
import itertools
from numpy import *


kinEnergyMax = 100000
# Parameters
tc=0.001# collision time 
en=.3  # normal restitution coefficient
es=.3  # tangential restitution coefficient
frictionAngle=radians(35)# 
density=2700
# facets material
facetMat=O.materials.append(ViscElMat(frictionAngle=frictionAngle,tc=tc,en=en,et=es))
# default spheres material
dfltSpheresMat=O.materials.append(ViscElMat(density=density,frictionAngle=frictionAngle,tc=tc,en=en,et=es)) 

O.dt=.05*tc # time step

Rs=0.05 # particle radius

# Create geometry

x0=0.; y0=0.; z0=0.; ab=.7; at=2.; h=1.; hl=h; al=at*3

zb=z0;   x0b=x0-ab/2.; y0b=y0-ab/2.; x1b=x0+ab/2.; y1b=y0+ab/2.
zt=z0+h; x0t=x0-at/2.; y0t=y0-at/2.; x1t=x0+at/2.; y1t=y0+at/2.
zl=z0-hl;x0l=x0-al/2.; y0l=y0-al/2.; x1l=x0+al/2.; y1l=y0+al/2.

vibrationPlate = O.bodies.append(geom.facetBunker((x0,y0,z0),
  dBunker=ab*4.0,dOutput=ab*1.7,hBunker=ab*3,hOutput=ab,hPipe=ab/3.0,
  wallMask=4,segmentsNumber=10,material=facetMat))

# Create engines
O.engines=[
	ForceResetter(),
	InsertionSortCollider([Bo1_Sphere_Aabb(),Bo1_Facet_Aabb()]),
	InteractionLoop(
		[Ig2_Sphere_Sphere_ScGeom(), Ig2_Facet_Sphere_ScGeom()],
		[Ip2_ViscElMat_ViscElMat_ViscElPhys()],
		[Law2_ScGeom_ViscElPhys_Basic()],
	),
	DomainLimiter(lo=(-4,-4,-1),hi=(4,4,4),iterPeriod=5000,label='domLim'),
	NewtonIntegrator(damping=0, gravity=[0,0,-9.81]),
	PyRunner(iterPeriod=6500,command='addBodies()',nDo=7,label='addb'),
	PyRunner(iterPeriod=1000,command='state()',label='state'),
]

numSphereGen = 0
def addBodies():
	global numSphereGen
	# Create clumps...
	clumpColor=(0.0, 0.5, 0.5)
	for k,l in itertools.product(arange(0,10),arange(0,10)):
		clpId,sphId=O.bodies.appendClumped([sphere(Vector3(x0t+Rs*(k*4+2),y0t+Rs*(l*4+2),i*Rs*2+zt+ab*3),Rs,color=clumpColor,material=dfltSpheresMat) for i in range(4)])
		numSphereGen += len(sphId)

	# ... and spheres
	spheresColor=(0.4, 0.4, 0.4)
	for k,l in itertools.product(arange(0,9),arange(0,9)):
		sphAloneId=O.bodies.append( [sphere( Vector3(x0t+Rs*(k*4+4),y0t+Rs*(l*4+4),i*Rs*2.3+zt+ab*3),Rs,color=spheresColor,material=dfltSpheresMat) for i in range(4) ] )
		numSphereGen += len(sphAloneId)

def state():
	global numSphereGen
	print("Iter %d: Total number of generated spheres %d, removed particles %d, current particles %d, kinEnergy %g"%(O.iter, numSphereGen, domLim.nDeleted, numSphereGen-domLim.nDeleted, utils.kineticEnergy()))
	if (utils.kineticEnergy() > kinEnergyMax):
		print("Kinetic energy is over a threshold value! Error!")

addBodies()


from yade import qt
qt.View()
O.saveTmp()

#O.run(50001)
