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

# This script perform a simple shear experiment with lubrication law.
# It shows the use of
# - Lubrication law
# - PDFEngine
# - VTK Lubrication recorder

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

sp = pack.SpherePack()

m_savefile = "data.txt"
m_vtkDir = "vtk"
if not os.path.exists(m_vtkDir):
	os.makedirs(m_vtkDir)

# Physics parameters
m_young = 1e9
#Pa
m_friction = atan(0.5)
#rad
m_viscosity = 100
#Pa.s
m_sphRadius = 0.1
# m
m_epsilon = 1e-3

# Simulation
m_shearRate = 10
m_N = 500
# Number of spheres
m_targetVolfrac = 0.5
m_stopOnStrain = 100

# Saving
m_sampling = 100.
# Number of sampling while 1 deformation
m_vtkSampling = 10.
# Number of sampling while 1 deformation

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

# Simulation cell
sq = (4. / 3. * pi * m_N / m_targetVolfrac)**(1. / 3.)

width = sq * m_sphRadius
length = sq * m_sphRadius
height = sq * m_sphRadius

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

# Sphere pack
No = sp.makeCloud((0, 0, 0), O.cell.size, m_sphRadius, 0.05, m_N, periodic=True, seed=1)
spheres = [utils.sphere(s[0], s[1]) for s in sp]
O.bodies.append(spheres)

# Setup interaction law
law = Law2_ScGeom_ImplicitLubricationPhys(
        activateTangencialLubrication=True, activateTwistLubrication=True, activateRollLubrication=True, resolution=2, theta=1, SolutionTol=1.e-8, MaxIter=50
)

# Setup engines
O.engines = [
        ForceResetter(),
        InsertionSortCollider([Bo1_Sphere_Aabb(aabbEnlargeFactor=2., label="aabb")], verletDist=-0.2, allowBiggerThanPeriod=False),
        InteractionLoop(
                [Ig2_Sphere_Sphere_ScGeom6D(interactionDetectionFactor=2., label="Ig2")],
                [Ip2_FrictMat_FrictMat_LubricationPhys(eta=m_viscosity, eps=m_epsilon)], [law]
        ),
        NewtonIntegrator(damping=0.),
        GlobalStiffnessTimeStepper(active=1, timeStepUpdateInterval=100, timestepSafetyCoefficient=0.8, defaultDt=1e-6, label="TimeStepper", viscEl=False),
        PDFEngine(filename="PDF.txt", virtPeriod=1. / (m_vtkSampling * m_shearRate), numDiscretizeAnglePhi=9, numDiscretizeAngleTheta=13),
        VTKRecorder(
                fileName=m_vtkDir + '/', recorders=['spheres', 'velocity', 'lubrication'], virtPeriod=1. / (m_vtkSampling * m_shearRate), label="VtkRecorder"
        ),
        PyRunner(command="UpPlot()", virtPeriod=min(1. / (m_sampling * m_shearRate), 0.1), label="UpdatePlot"),
        PyRunner(command="SavePlot()", realPeriod=600, label="SaveDataPlot"),
        PyRunner(command="checkStartShear()", iterPeriod=10, label="beginCheck")
]

plot.plots = {
        'time': ('totalStress_yy', 'normalContactStress_yy', 'shearContactStress_yy', 'normalLubrifStress_yy', 'shearLubrifStress_yy', 'kineticStress_yy'),
        'time2': ('phi'),
        'time3': ('totalStress_xy', 'normalContactStress_xy', 'shearContactStress_xy', 'normalLubrifStress_xy', 'shearLubrifStress_xy', 'kineticStress_xy')
}

plot.plot(subPlots=True)

O.dt = 1e-6

# Firstly, compress to target volumic fraction
O.cell.velGrad = Matrix3(0, 0, 0, 0, -10, 0, 0, 0, 0)


def SavePlot():
	global m_savefile
	plot.saveDataTxt(m_savefile)


def UpPlot():
	global m_stopOnStrain
	[normalContactStress, shearContactStress, normalLubrifStress, shearLubrifStress,
	 potentialStress] = Law2_ScGeom_ImplicitLubricationPhys.getTotalStresses()
	kineticStress = getTotalDynamicStress()

	totalStress = normalContactStress + shearContactStress + normalLubrifStress + shearLubrifStress + potentialStress + kineticStress

	phi = 1. - porosity()

	if abs(O.cell.hSize[0, 1] / O.cell.hSize[0, 0]) > 1:
		flipCell()

	plot.addData(
	        totalStress=totalStress,
	        totalStress2=getStress(),
	        kineticStress=kineticStress,
	        normalContactStress=normalContactStress,
	        shearContactStress=shearContactStress,
	        normalLubrifStress=normalLubrifStress,
	        shearLubrifStress=shearLubrifStress,
	        potentialStress=potentialStress,
	        phi=phi,
	        iter=O.iter,
	        strain=O.cell.trsf,
	        time=O.time,
	        time2=O.time,
	        time3=O.time,
	        velGrad=O.cell.velGrad
	)

	if (m_stopOnStrain > 0) & (O.cell.trsf[0, 1] > m_stopOnStrain):
		SaveAndQuit()


def checkStartShear():
	global m_shearRate
	phi = 1. - porosity()
	start = m_targetVolfrac < phi

	if start:
		print("Start shear.")
		O.cell.velGrad = Matrix3(0, m_shearRate, 0, 0, 0, 0, 0, 0, 0)
		O.cell.trsf = Matrix3(1, 0, 0, 0, 1, 0, 0, 0, 1)
		beginCheck.dead = 1


def SaveAndQuit():
	print("Quit condition reach.")
	SavePlot()
	O.stopAtIter = O.iter + 1
