""" Author: Ning Guo <ceguo@connect.ust.hk>
    run `mv biaxialSmooth.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 ./biaxialSmooth.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 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.0001; confining=-1.e5;
lx = 0.05; ly = 0.1; # sample size, 50mm by 100mm
nx = 8; ny = 16; # sample discretization, 8 by 16 quadrilateral elements
mydomain = Rectangle(l0=lx,l1=ly,n0=nx,n1=ny,order=2,integrationOrder=2)
dim = mydomain.getDim()
k = kronecker(mydomain)
numg = 4*nx*ny; # number of Gauss points, 4 GP each element (reduced integration)
nump = 16; # number of processes for multiprocessing

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

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

t=0

stress = prob.getCurrentStress() # initial stress
proj = Projector(mydomain)
sig = proj(stress) # project Gauss point value to nodal value
sig_bounda = interpolate(sig,FunctionOnBoundary(mydomain)) # interpolate
traction = matrix_mult(sig_bounda,mydomain.getNormal()) # boundary traction
x = mydomain.getX() # nodal coordinate
bx = FunctionOnBoundary(mydomain).getX()
topSurf = whereZero(bx[1]-sup(bx[1]))
tractTop = traction*topSurf # traction at top surface
forceTop = integrate(tractTop,where=FunctionOnBoundary(mydomain)) # resultant force at top
lengthTop = integrate(topSurf,where=FunctionOnBoundary(mydomain)) # length of top surface
fout=file('./result/biaxial_surf.dat','w')
fout.write('0 '+str(forceTop[1])+' '+str(lengthTop)+'\n')
# Dirichlet BC positions, smooth at bottom and top, fixed at the center of bottom
Dbc = whereZero(x[1])*[0,1]+whereZero(x[1]-ly)*[0,1]+whereZero(x[1])*whereZero(x[0]-.5*lx)*[1,1]
# Dirichlet BC values
Vbc = whereZero(x[1])*[0,0]+whereZero(x[1]-ly)*[0,vel]+whereZero(x[1])*whereZero(x[0]-.5*lx)*[0,0]
# Neumann BC, constant confining pressure
Nbc = whereZero(bx[0])*[-confining,0]+whereZero(bx[0]-lx)*[confining,0]

time_start = time.time()
while t < 100: # apply 100 load steps

   prob.initialize(f=Nbc, specified_u_mask=Dbc, specified_u_val=Vbc) # initialize BC
   t += 1
   du=prob.solve(iter_max=100) # get solution: nodal displacement

   disp += du
   stress=prob.getCurrentStress()
   
   dom = prob.getDomain() # domain is updated Lagrangian formulation
   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))
   lengthTop = integrate(topSurf,where=FunctionOnBoundary(dom))
   fout.write(str(t*vel/ly)+' '+str(forceTop[1])+' '+str(lengthTop)+'\n')
      
   vR=prob.getLocalVoidRatio()
   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/biaxialSmooth_%d.vtu"%t,disp=disp,shear=shear,e=vR)

prob.getCurrentPacking(pos=(),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()
