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

# This example allows to test different kernel functions
from __future__ import print_function
from yade import utils, plot, qt
o = Omega()

# Physical parameters
fr = 0.5;
rho=1000.0

k = 1.0
mu = 100.0
tc = 0.0001; en = 0.7; et = 0.7;
vel = 0.05
Rad = 15.0e-3
o.dt = 0.0001

shift = 1.0
# Add material
mat1 = O.materials.append(ViscElMat(frictionAngle=fr,density=rho, SPHmode=True,mu=mu,tc=tc, en=en, et=et, KernFunctionPressure = 1, KernFunctionVisco = 1))
mat2 = O.materials.append(ViscElMat(frictionAngle=fr,density=rho, SPHmode=True,mu=mu,tc=tc, en=en, et=et, KernFunctionPressure = 1, KernFunctionVisco = 2))
mat3 = O.materials.append(ViscElMat(frictionAngle=fr,density=rho, SPHmode=True,mu=mu,tc=tc, en=en, et=et, KernFunctionPressure = 1, KernFunctionVisco = 3))
mat4 = O.materials.append(ViscElMat(frictionAngle=fr,density=rho, SPHmode=True,mu=mu,tc=tc, en=en, et=et, KernFunctionPressure = 1, KernFunctionVisco = 4))
mat5 = O.materials.append(ViscElMat(frictionAngle=fr,density=rho, SPHmode=True,mu=mu,tc=tc, en=en, et=et, KernFunctionPressure = 1, KernFunctionVisco = 5))

# Add spheres
#inert = True
inert = True
id11 = O.bodies.append(sphere(center=[0,0,0],radius=Rad,                       material=mat1, mask = 1, fixed=True))
id12 = O.bodies.append(sphere(center=[0,0,(Rad*2.0*shift)],radius=Rad,          material=mat1, mask = 1, fixed=inert))

id21 = O.bodies.append(sphere(center=[3.0*Rad,0,0],radius=Rad,                 material=mat2, mask = 1, fixed=True))
id22 = O.bodies.append(sphere(center=[3.0*Rad,0,(Rad*2.0*shift)],radius=Rad,    material=mat2, mask = 1, fixed=inert))

id31 = O.bodies.append(sphere(center=[6.0*Rad,0,0],radius=Rad,                 material=mat3, mask = 1, fixed=True))
id32 = O.bodies.append(sphere(center=[6.0*Rad,0,(Rad*2.0*shift)],radius=Rad,    material=mat3, mask = 1, fixed=inert))

id41 = O.bodies.append(sphere(center=[9.0*Rad,0,0],radius=Rad,                 material=mat4, mask = 1, fixed=True))
id42 = O.bodies.append(sphere(center=[9.0*Rad,0,(Rad*2.0*shift)],radius=Rad,    material=mat4, mask = 1, fixed=inert))

id51 = O.bodies.append(sphere(center=[12.0*Rad,0,0],radius=Rad,                material=mat5, mask = 1, fixed=True))
id52 = O.bodies.append(sphere(center=[12.0*Rad,0,(Rad*2.0*shift)],radius=Rad,   material=mat5, mask = 1, fixed=inert))

vel = 0.1
O.bodies[id12].state.vel=Vector3(0,0,-vel)
O.bodies[id22].state.vel=Vector3(0,0,-vel)
O.bodies[id32].state.vel=Vector3(0,0,-vel)
O.bodies[id42].state.vel=Vector3(0,0,-vel)
O.bodies[id52].state.vel=Vector3(0,0,-vel)

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

print("Time\tX\tRho\tP\tFpr ")

# Function to add data to plot
def addPlotData(): 
  #print "%.2f\t%.5f\t%.5f\t%.5f\t%.5f" % (O.time+O.dt, O.bodies[id2].state.pos[2], O.bodies[id2].rho, O.bodies[id2].press, O.forces.f(id2)[2])
  s1 = (O.bodies[id12].state.pos[2]-O.bodies[id11].state.pos[2])-(Rad*2.0)
  s2 = (O.bodies[id22].state.pos[2]-O.bodies[id21].state.pos[2])-(Rad*2.0)
  s3 = (O.bodies[id32].state.pos[2]-O.bodies[id31].state.pos[2])-(Rad*2.0)
  s4 = (O.bodies[id42].state.pos[2]-O.bodies[id41].state.pos[2])-(Rad*2.0)
  s5 = (O.bodies[id52].state.pos[2]-O.bodies[id51].state.pos[2])-(Rad*2.0)
  
  f1 = O.forces.f(id12)[2]
  f2 = O.forces.f(id22)[2]
  f3 = O.forces.f(id32)[2]
  f4 = O.forces.f(id42)[2]
  f5 = O.forces.f(id52)[2]
  
  plot.addData(sc=O.iter, s1 = s1, s2 = s2, s3 = s3, s4 = s4, s5 = s5,
               fc=O.iter, f1 = f1, f2 = f2, f3 = f3, f4 = f4, f5 = f5)

plot.plots={'sc':('s1', 's2', 's3', 's4', 's5'), 'fc':('f1', 'f2', 'f3', 'f4', 'f5')}; plot.plot()

O.run(1, True)

qt.View()
O.run(1500)
