#!/usr/bin/env python
# encoding: utf-8

from __future__ import print_function
from builtins import range
from yade import utils, plot, qt
o = Omega()

# Physical parameters
fr = 0.5;
rho= 1000.0

k = 1000.0
mu = 10.0
tc = 0.01; en = 0.7; et = 0.7;
vel = 0.0

#Rad = 0.006
#h = 0.011

Rad = 0.015
h = 0.03

#Rad = 0.02
#h = 0.04

o.dt = 0.0005

X = 4.0
Z = 3.0
yCoeff = 4.0

SpheresX = 1.0
SpheresZ = 2.0

# Add material
mat1 = O.materials.append(ViscElMat(frictionAngle=fr,density=rho, SPHmode=True, h=h, mu=mu,tc=tc, en=en, et=et, KernFunctionPressure = 1, KernFunctionVisco = 1))
idBox = pack.regularHexa(pack.inAlignedBox((0.0,-Rad*yCoeff,0.0),(X,Rad*yCoeff,Z)),radius=Rad,gap=0.0,color=(0,1,1), material=mat1, mask = 1, fixed=True)

idBoxAdd = []
idSpheresAdd = []

for i in range(len(idBox)):
  if (((idBox[i].state.pos[0])<yCoeff*Rad) or
      ((idBox[i].state.pos[0])>(X-yCoeff*Rad - 2.0*Rad)) or
      ((idBox[i].state.pos[2])<(yCoeff*Rad)) or
      ((idBox[i].state.pos[2])>(Z - 2.0*Rad))):
    idBoxAdd.append(idBox[i])
  elif (((idBox[i].state.pos[0])<SpheresX) and
        ((idBox[i].state.pos[2])<SpheresZ) and
        ((idBox[i].state.pos[1])>-Rad) and
        ((idBox[i].state.pos[1])<2.0*Rad)):
          idBox[i].shape.color = Vector3(1,0,0)
          idBox[i].state.fixed = False
          idBox[i].state.blockedDOFs = 'y'
          idSpheresAdd.append(idBox[i])
    


O.bodies.append(idBoxAdd)
idSpheres = O.bodies.append(idSpheresAdd)

    
# Add engines
o.engines = [
  ForceResetter(),
  InsertionSortCollider([Bo1_Sphere_Aabb(label='is2aabb')],ompThreads=1),
  InteractionLoop(
    [Ig2_Sphere_Sphere_ScGeom(label='ss2sc')],
    [Ip2_ViscElMat_ViscElMat_ViscElPhys()],
    [Law2_ScGeom_ViscElPhys_Basic()],
  ),
  NewtonIntegrator(damping=0.01,gravity=[0,0,-9.81]),
  SPHEngine(mask=1, k=k, rho0 = rho, h=h, KernFunctionDensity= 1),
  PyRunner(command='addPlotData()',iterPeriod=10,initRun=True,dead=False),
]

enlargeF = h/Rad*1.1
print("enlargeF = %g"%enlargeF)
is2aabb.aabbEnlargeFactor = enlargeF
ss2sc.interactionDetectionFactor = enlargeF


# Function to add data to plot
def addPlotData(): 
  plot.addData(t = O.time, eKin = utils.kineticEnergy())
  
plot.plots={'t':('eKin')};
plot.plot()

qt.View()
