""" Author: Ning Guo <ceguo@connect.ust.hk>
    run `mv undrain.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 ./undrain.py
Please amend these instructions if you find that they do not work.
"""

from esys.escript import *
from esys.finley import Rectangle
from esys.weipa import saveVTK
from esys.escript.pdetools import Projector
from msFEMup 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

confining = -2.e5
pore = 1.e5  # initial pore pressure
perm = 0.001**2 / (180. * 8.9e-4)
# unscaled permeability, using KC equation
kf = 2.2e9  # fluid bulk modulus
dt = .1
vel = -0.0001  # time step and loading speed
lx = 0.05
ly = 0.1  # sample dimension
nx = 8
ny = 16  # discretization
mydomain = Rectangle(l0=lx, l1=ly, n0=nx, n1=ny, order=2, integrationOrder=2)
k = kronecker(mydomain)
dim = 2.
numg = 4 * nx * ny
# no. of Gauss points
mpi = True  # use MPI

prob = MultiScale(domain=mydomain, pore0=pore, perm=perm, kf=kf, dt=dt, ng=numg, useMPI=mpi, rtol=1.e-2)
disp = Vector(0., Solution(mydomain))
t = 0
ux = mydomain.getX()  # disp. node coordinate
px = ReducedSolution(mydomain).getX()  # press. node coordinate
bx = FunctionOnBoundary(mydomain).getX()
topSurf = whereZero(bx[1] - ly)
uDbc = whereZero(ux[1]) * [0, 1] + whereZero(ux[0] - lx / 2.) * whereZero(ux[1]) * [1, 1] + whereZero(ux[1] - ly) * [0, 1]  # disp. Dirichlet BC mask
vDbc = whereZero(ux[1]) * [0, 0] + whereZero(ux[0] - lx / 2.) * whereZero(ux[1]) * [0, 0] + whereZero(ux[1] - ly) * [0, vel * dt]  # disp. Dirichlet BC value
uNbc = whereZero(bx[0]) * [-confining, 0] + whereZero(bx[0] - lx) * [confining, 0]  # disp. Neumann BC

stress = prob.getCurrentStress()  # effective stress at GP
proj = Projector(mydomain)
sig = proj(stress)  # effective stress at node (reduced)
sig_bound = interpolate(sig, FunctionOnBoundary(mydomain))
traction = matrix_mult(sig_bound, mydomain.getNormal())
tractTop = traction * topSurf
forceTop = integrate(tractTop, where=FunctionOnBoundary(mydomain))
lengthTop = integrate(topSurf, where=FunctionOnBoundary(mydomain))
fout = file('./result/resultantForce.dat', 'w')
fout.write('0 ' + str(forceTop[1]) + ' ' + str(lengthTop) + '\n')

time_start = time.time()
while t < 400:
	prob.initialize(f=uNbc, umsk=uDbc, uvalue=vDbc)
	t += 1
	du = prob.solve(globalIter=10, solidIter=50)
	disp += du
	pore = prob.getCurrentPore()  # pore pressure at node (reduced)
	flux = prob.getCurrentFlux()  # Darcy flux at GP
	stress = prob.getCurrentStress()  # effective stress at GP
	strain = prob.getCurrentStrain()  # disp. grad at GP
	volume_strain = trace(strain)  # volumetric strain
	dev_strain = symmetric(strain) - volume_strain * k / dim  # deviatoric strain
	shear = sqrt(2. * inner(dev_strain, dev_strain))  # shear strain
	fab = prob.getLocalFabric()  # fabric tensor at GP
	dev_fab = 4. * (fab - trace(fab) / dim * k)
	anis = sqrt(.5 * inner(dev_fab, dev_fab))
	p = prob.getEquivalentPorosity()  # porosity at GP
	rot = prob.getLocalAvgRotation()  # average rotation at GP
	saveGauss2D(name='./result/gauss/time_' + str(t) + '.dat', strain=strain, fabric=fab, stress=stress)
	dom = prob.getDomain()  # domain updated (Lagrangian)
	proj = Projector(dom)
	flux = proj(flux)  # Darcy flux at node (reduced)
	p = proj(p)  # porosity at node (reduced)
	shear = proj(shear)  # shear strain at node (reduced)
	anis = proj(anis)
	rot = proj(rot)
	saveVTK("./result/vtk/undrain_%d.vtu" % t, disp=disp, pore=pore, flux=flux, shear=shear, p=p, anis=anis, rot=rot)
	sig = proj(stress)  # effective stress at node (reduced)
	sig_bound = interpolate(sig, FunctionOnBoundary(dom))
	traction = matrix_mult(sig_bound, dom.getNormal())
	tractTop = traction * topSurf
	forceTop = integrate(tractTop, where=FunctionOnBoundary(dom))
	lengthTop = integrate(topSurf, where=FunctionOnBoundary(dom))
	fout.write(str(t * vel * dt / ly) + ' ' + str(forceTop[1]) + ' ' + str(lengthTop) + '\n')

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