# -*- encoding=utf-8 -*-

# Read parameters from a table, or assign default values
readParamsFromTable(young=1.e8, friction=0.5, viscosity=10, roughness=7.e-4, numberOfGrains=10000, savedir=".", dv_xy=0, dv_xz=10, g=-50, dRho=1000)

from yade.params import table

from yade import pack, ymport, export, geom, bodiesHandling, plot
import math
import pylab
#from yade import qt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import datetime as DT
import os
import re

# Prepare python dict to save parameters in output file
m_simuParams = {k: yade.params.table.__dict__[k] for k in table.__dict__['__all__']}
m_simuParams['date_start'] = DT.datetime.now().isoformat()

# Prepare output destinations
m_savefile = table.savedir + '/results/data'

if not os.path.exists(table.savedir + '/results/'):
	os.makedirs(table.savedir + '/results/')

i = 0
while os.path.exists(m_savefile + '.' + str(i) + '.txt'):
	i = i + 1

m_savefile += '.' + str(i) + '.txt'
open(m_savefile, 'a').close()
# Create the empty file asap

m_vtkDir = table.savedir + '/results/vtk.' + str(i) + '/'

if not os.path.exists(m_vtkDir):
	os.makedirs(m_vtkDir)

sp = pack.SpherePack()
launchDT = DT.datetime.now().isoformat()

# Physics parameters
m_young = table.young
#Pa
m_friction = atan(table.friction)
#rad
m_viscosity = table.viscosity
#Pa.s
m_sphRadius = 1e-2
# m
m_epsilon = table.roughness

# Simulation
m_N = table.numberOfGrains
# Number of spheres
m_savePeriod = 1. / max(abs(table.dv_xy), max(abs(table.dv_xz), 1.))

m_simuParams['a'] = m_sphRadius

#define material for all bodies:
id_Mat = O.materials.append(FrictMat(young=m_young, poisson=0.3, density=table.dRho, frictionAngle=m_friction))
Mat = O.materials[id_Mat]

# Simulation cell
sq = math.floor((table.numberOfGrains * 20)**(1. / 3.))

width = sq * m_sphRadius
length = sq * m_sphRadius
height = 10. * sq * m_sphRadius

m_simuParams['A'] = width * length

O.periodic = True
O.cell.hSize = Matrix3(width, 0, 0, 0, height, 0, 0, 0, length)

# Shearing
O.cell.velGrad = Matrix3(0, table.dv_xy, table.dv_xz, 0, 0, 0, 0, 0, 0)

# Sedimentation time
tss = (-2. * height / table.g)**(1. / 2.)

# Lower plate (break periodicity)
sphs = [
        utils.sphere(radius=m_sphRadius, center=m_sphRadius * Vector3(x, random.uniform(-0.5, 0.5), z), fixed=True, mask=3)
        for x in range(0, sq + 2)
        for z in range(0, sq + 2)
]
O.bodies.append(sphs)

# Generate sphere cloud
No = sp.makeCloud((0, 5 * m_sphRadius, 0), O.cell.size - Vector3(0, 5 * m_sphRadius, 0), m_sphRadius, 0.1, m_N, periodic=True, seed=1)
spheres = [utils.sphere(s[0], s[1], mask=2) for s in sp]
O.bodies.append(spheres)

# Simulations engines
O.engines = [
        ForceResetter(),
        InsertionSortCollider([Bo1_Sphere_Aabb(aabbEnlargeFactor=2., label="aabb")], verletDist=-1, allowBiggerThanPeriod=False, label="collider"),
        InteractionLoop(
                [Ig2_Sphere_Sphere_ScGeom6D(interactionDetectionFactor=2., label="Ig2")],
                [Ip2_FrictMat_FrictMat_LubricationPhys(eta=m_viscosity, eps=m_epsilon)],
                [
                        Law2_ScGeom_ImplicitLubricationPhys(
                                activateTangencialLubrication=True,
                                activateTwistLubrication=True,
                                activateRollLubrication=True,
                                resolution=2,
                                theta=1,
                                SolutionTol=1.e-8,
                                MaxIter=50,
                                MaxDist=1.
                        )
                ]  # Use lubrication law (suspension)
        ),
        LinearDragEngine(nu=m_viscosity, dead=1, label="drag"),
        NewtonIntegrator(damping=0., gravity=(0, table.g, 0), label="newton"),
        GlobalStiffnessTimeStepper(active=1, timeStepUpdateInterval=100, timestepSafetyCoefficient=0.8, defaultDt=1e-6, label="TimeStepper", viscEl=False),
        VTKRecorder(fileName=m_vtkDir, recorders=['spheres', 'velocity', 'lubrication', 'mask'], virtPeriod=m_savePeriod, label="VtkRecorder"),
        PyRunner(command="upPlot()", virtPeriod=m_savePeriod),
        PyRunner(command="SavePlot()", realPeriod=600),
        PyRunner(command="CheckFlip()", virtPeriod=1. / (10 * max(table.dv_xy, max(table.dv_xz, 1.)))),  # Very important: Check if cell must be flipped
        PyRunner(command="drag.ids=[b.id for b in O.bodies if len(O.interactions.withBody(b.id)) < 2 and b.mask == 2]",
                 iterPeriod=1000),  # Maybe important: apply drag for particles with less than 2 interactions: prevent free fall on the interface
        PyRunner(command="drag.dead=0;", nDo=2, virtPeriod=2 * tss),
]

collider.avoidSelfInteractionMask = 1
# Do not collide spheres on the bottom plane

O.dt = 1e-5

# Plot will show height isovalues
plot.plots = {'time': ('y0', 'y50', 'y80', 'y95', 'ym')}
plot.plot()


# Enhanced flipping
def CheckFlip():

	f_xy = 0
	f_xz = 0

	if abs(O.cell.hSize[0, 2] / O.cell.hSize[0, 0]) > 1:
		f_xz = -ceil(abs(O.cell.hSize[0, 2] / O.cell.hSize[0, 0])) * O.cell.hSize[0, 2] / abs(O.cell.hSize[0, 2])

	if abs(O.cell.hSize[0, 1] / O.cell.hSize[0, 0]) > 1:
		f_xy = -ceil(abs(O.cell.hSize[0, 1] / O.cell.hSize[0, 0])) * O.cell.hSize[0, 1] / abs(O.cell.hSize[0, 1])

	if f_xy != 0 or f_xz != 0:
		flipCell(Matrix3(0, f_xy, f_xz, 0, 0, 0, 0, 0, 0))


# Update plot
def upPlot():
	yBodies = [b.state.pos[1] for b in O.bodies if b.mask == 2]
	yBodies.sort()

	plot.addData(
	        iter=O.iter,
	        time=O.time,
	        strain=O.cell.trsf,
	        y0=yBodies[0],
	        y50=yBodies[math.floor(0.5 * len(yBodies))],
	        y80=yBodies[math.floor(0.8 * len(yBodies))],
	        y95=yBodies[math.floor(0.95 * len(yBodies))],
	        ym=yBodies[-1]
	)


# Save plot (not too often to preserve performances)
def SavePlot():
	global m_savefile, m_simuParams
	plot.saveDataTxt(m_savefile, headers=m_simuParams)


# Call this function to quit
def SaveAndQuit():
	print("Quit condition reach.")
	SavePlot()
	O.stopAtIter = O.iter + 1


# Usefull ifbatch
if 'YADE_BATCH' in os.environ:
	O.run()
	waitIfBatch()
