""" Author: Ning Guo <ceguo@connect.ust.hk>
    run `mv retainingSmooth.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 ./retainingSmooth.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 Rectangle
from esys.weipa import saveVTK
from esys.escript.pdetools import Projector
from msFEM2D 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.00017; surcharge=-2.e4; # surcharge equals to the initial vertical stress of the RVE packing; vel<0 passive failure; vel>0 active failure
B = 0.4; H = 0.2; wallH = 0.17; baseH = H-wallH; # setup dimensions
nx = 40; ny = 20 # discretization with 40x20 quads
mydomain = Rectangle(l0=B,l1=H,n0=nx,n1=ny,order=2,integrationOrder=2)
dim = mydomain.getDim()
k = kronecker(mydomain)
numg = 4*nx*ny; nump = 16;
packNo=list(range(0,numg,16))

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

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

t=0
time_start = time.time()

x = mydomain.getX()
bx = FunctionOnBoundary(mydomain).getX()
left = whereZero(x[0])
right = whereZero(x[0]-B)
bottom = whereZero(x[1])
top = whereZero(bx[1]-H)
base = whereZero(x[0]-B)*whereNegative(x[1]-baseH)
wall = whereZero(x[0]-B)*whereNonNegative(x[1]-baseH)
wallBF = whereZero(bx[0]-B)*whereNonNegative(bx[1]-baseH)
# Dirichlet BC, all fixed in space except wall (only fixed in x direction, smooth)
Dbc = left*[1,1]+base*[1,1]+bottom*[1,1]+wall*[1,0]
Vbc = left*[0,0]+base*[0,0]+bottom*[0,0]+wall*[vel,0]
# Neumann BC, apply surcharge at the top surface
Nbc = top*[0,surcharge]

stress = prob.getCurrentStress()
proj = Projector(mydomain)
sig = proj(stress)
sig_bounda = interpolate(sig,FunctionOnBoundary(mydomain))
traction = matrix_mult(sig_bounda,mydomain.getNormal())
tract = traction*wallBF # traction on wall
forceWall = integrate(tract,where=FunctionOnBoundary(mydomain)) # force on wall
lengthWall = integrate(wallBF,where=FunctionOnBoundary(mydomain))
fout=file('./result/pressure.dat','w')
fout.write('0 '+str(forceWall[0])+' '+str(lengthWall)+'\n')

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())
   tract = traction*wallBF
   forceWall = integrate(tract,where=FunctionOnBoundary(dom))
   lengthWall = integrate(wallBF,where=FunctionOnBoundary(dom))
   fout.write(str(t*vel)+' '+str(forceWall[0])+' '+str(lengthWall)+'\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/retainingSmooth_%d.vtu"%t,disp=disp,stress=stress,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()
