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

from builtins import range
from esys.escript import *
from esys.finley import Brick
from esys.weipa import saveVTK
from esys.escript.pdetools import Projector
from msFEM3D import MultiScale
from saveGauss import saveGauss3D
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.0001; confining=-1.e5;
lx = 0.05; ly = 0.05; lz = 0.1; # sample dimension
nx = 8; ny = 8; nz = 16; # discretization
mydomain = Brick(l0=lx,l1=ly,l2=lz,n0=nx,n1=ny,n2=nz,order=2,integrationOrder=2) # 20-noded,8-Gauss hexahedral element
dim = 3; k = kronecker(mydomain)
numg = 8*nx*ny*nz; # number of Gauss points
packNo=list(range(0,numg,8))

prob = MultiScale(domain=mydomain,ng=numg,useMPI=True,rtol=1e-2) # mpi is activated

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()) # traction on boundary
x = mydomain.getX()
bx = FunctionOnBoundary(mydomain).getX()
topSurf = whereZero(bx[2]-lz)
tractTop = traction*topSurf # traction on top surface
forceTop = integrate(tractTop,where=FunctionOnBoundary(mydomain)) # force on top
areaTop = integrate(topSurf,where=FunctionOnBoundary(mydomain))
fout=file('./result/resultant.dat','w')
fout.write('0 '+str(forceTop[2])+' '+str(areaTop)+'\n')

# Dirichlet BC, rough loading ends
Dbc = whereZero(x[2])*[1,1,1]+whereZero(x[2]-lz)*[1,1,1]
Vbc = whereZero(x[2])*[0,0,0]+whereZero(x[2]-lz)*[0,0,vel]
# Neumann BC, constant lateral confining
Nbc = whereZero(bx[0])*[-confining,0,0]+whereZero(bx[0]-sup(bx[0]))*[confining,0,0]+whereZero(bx[1])*[0,-confining,0]+whereZero(bx[1]-sup(bx[1]))*[0,confining,0]

time_start = time.time()
while t < 100:

   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()
   proj = Projector(dom)
   sig = proj(stress)

   sig_bounda = interpolate(sig,FunctionOnBoundary(dom))
   traction = matrix_mult(sig_bounda,dom.getNormal())
   tractTop = traction*topSurf
   forceTop = integrate(tractTop,where=FunctionOnBoundary(dom))
   areaTop = integrate(topSurf,where=FunctionOnBoundary(dom))
   fout.write(str(t*vel/lz)+' '+str(forceTop[2])+' '+str(areaTop)+'\n')
      
   vR=prob.getLocalVoidRatio()
   rotation=prob.getLocalAvgRotation()
   fabric=prob.getLocalFabric()
   strain = prob.getCurrentStrain()
   saveGauss3D(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./3.*inner(dev_strain,dev_strain))
   saveVTK("./result/vtk/triaxialRough_%d.vtu"%t,disp=disp,shear=shear,e=vR,rot=rotation)
   
prob.getCurrentPacking(pos=packNo,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()
