""" Author: Ning Guo <ceguo@connect.ust.hk>
    run `mv footing.yade.gz 0.yade.gz`
    to generate initial RVE packing

How to run this script:
    sudo apt install python-escript
    cd examples/FEMxDEM
    export PYTHONPATH="/usr/lib/python-escript:../../py/FEMxDEM"
    export LD_LIBRARY_PATH=/usr/lib/python-escript/lib
    ln -s /path/to/yade ../../py/FEMxDEM/yadeimport.py
    /path/to/yade ./footing.py
Please amend these instructions if you find that they do not work.
"""

from esys.escript import *
from esys.finley import ReadGmsh
from esys.weipa import saveVTK
from esys.escript.pdetools import Projector
# FIXME: which one of those should it be: msFEM2D, msFEM3D, msFEMup ?
from msFEM import MultiScale
from saveGauss import saveGauss2D
import time
import os
import errno

try:
	os.mkdir('./result/')
	os.mkdir('./result/gauss')
	os.mkdir('./result/vtk')
	os.mkdir('./result/packing')
except OSError as exc:
	if exc.errno != errno.EEXIST:
		raise
	pass

vel = -0.00025
surcharge = -20.e3
# surcharge equals to the initial vertical stress of the RVE packing
dim = 2
B = 0.05
L = 0.6
H = 0.4
mydomain = ReadGmsh('footing.msh', numDim=dim, integrationOrder=2)  # read Gmsh mesh with 6-node triangle element (2500 tri6); each element has 3 Gauss points
k = kronecker(mydomain)
numg = 3 * 2500
# number of Gauss points
nump = 16
# number of processes in multiprocessing
packNo = list(range(0, numg, 50))

prob = MultiScale(domain=mydomain, ng=numg, np=nump, random=False, rtol=1e-2, usePert=False, pert=-2.e-5, verbose=False)

disp = Vector(0., Solution(mydomain))

t = 0
stress = prob.getCurrentStress()
proj = Projector(mydomain)
sig = proj(stress)
sig_bounda = interpolate(sig, FunctionOnBoundary(mydomain))
traction = matrix_mult(sig_bounda, mydomain.getNormal())
x = mydomain.getX()
bx = FunctionOnBoundary(mydomain).getX()
footingBase = whereZero(bx[1] - sup(bx[1])) * whereNonPositive(bx[0] - B)
tractFoot = traction * footingBase
forceFoot = integrate(tractFoot, where=FunctionOnBoundary(mydomain))
lengthFoot = integrate(footingBase, where=FunctionOnBoundary(mydomain))
fout = file('./result/bearing.dat', 'w')
fout.write('0 ' + str(forceFoot[0]) + ' ' + str(forceFoot[1]) + ' ' + str(lengthFoot) + '\n')

# Dirichlet BC, rollers at left and right, fixties at bottom, rigid and rough footing
Dbc = whereZero(x[0]) * [1, 0] + whereZero(x[0] - sup(x[0])) * [1, 0] + whereZero(x[1] - inf(x[1])) * [1, 1] + whereZero(x[1] - sup(x[1])
                                                                                                                        ) * whereNonPositive(x[0] - B) * [1, 1]
Vbc = whereZero(x[0]) * [0, 0] + whereZero(x[0] - sup(x[0])) * [0, 0] + whereZero(x[1] - inf(x[1])) * [
        0, 0
] + whereZero(x[1] - sup(x[1])) * whereNonPositive(x[0] - B) * [0, vel]
# Neumann BC, surcharge at the rest area of the top surface
Nbc = whereZero(bx[1] - sup(bx[1])) * wherePositive(bx[0] - B) * [0, surcharge]

time_start = time.time()
while t < 58:  # apply 58 loading step; further loading would abort the program due to severe mesh distortion

	prob.initialize(f=Nbc, specified_u_mask=Dbc, specified_u_val=Vbc)
	t += 1
	du = prob.solve(iter_max=100)

	disp += du
	stress = prob.getCurrentStress()

	dom = prob.getDomain()  # domain updated (Lagrangian)
	proj = Projector(dom)
	sig = proj(stress)

	sig_bounda = interpolate(sig, FunctionOnBoundary(dom))
	traction = matrix_mult(sig_bounda, dom.getNormal())
	tractFoot = traction * footingBase
	forceFoot = integrate(tractFoot, where=FunctionOnBoundary(dom))
	lengthFoot = integrate(footingBase, where=FunctionOnBoundary(dom))
	fout.write(str(t * vel) + ' ' + str(forceFoot[0]) + ' ' + str(forceFoot[1]) + ' ' + str(lengthFoot) + '\n')

	vR = prob.getLocalVoidRatio()
	rotation = prob.getLocalAvgRotation()
	fabric = prob.getLocalFabric()
	strain = prob.getCurrentStrain()
	saveGauss2D(name='./result/gauss/time_' + str(t) + '.dat', strain=strain, stress=stress, fabric=fabric)
	volume_strain = trace(strain)
	dev_strain = symmetric(strain) - volume_strain * k / dim
	shear = sqrt(2 * inner(dev_strain, dev_strain))
	saveVTK("./result/vtk/footing_%d.vtu" % t, disp=disp, stress=stress, shear=shear, e=vR, rot=rotation)

prob.getCurrentPacking(pos=packNo, time=t, prefix='./result/packing/')  # output packing
time_elapse = time.time() - time_start
fout.write("#Elapsed time in hours: " + str(time_elapse / 3600.) + '\n')
fout.close()
prob.exitSimulation()
