from yade import pack
from yade import export
from yade import timing
from yade import plot

if ( 'PFVFLOW' in features ):

  num_spheres=1000# number of spheres
  young=1e6
  compFricDegree = 3 # initial contact friction during the confining phase
  finalFricDegree = 30 # contact friction during the deviatoric loading
  mn,mx=Vector3(0,0,0),Vector3(1,1,0.4) # corners of the initial packing
  graindensity=2600
  toleranceWarning =1.e-11
  toleranceCritical=1.e-6

  O.materials.append(FrictMat(young=young,poisson=0.5,frictionAngle=radians(compFricDegree),density=graindensity,label='spheres'))
  O.materials.append(FrictMat(young=young,poisson=0.5,frictionAngle=0,density=0,label='walls'))
  walls=aabbWalls([mn,mx],thickness=0,material='walls')
  wallIds=O.bodies.append(walls)

  sp=pack.SpherePack()
  sp.makeCloud(mn,mx,-1,0.3333,num_spheres,False, 0.95,seed=1) #"seed" make the "random" generation always the same
  sp.toSimulation(material='spheres')

  triax=TriaxialStressController(
          maxMultiplier=1.+2e4/young, # spheres growing factor (fast growth)
          finalMaxMultiplier=1.+2e3/young, # spheres growing factor (slow growth)
          thickness = 0,
          stressMask = 7,
          max_vel = 0.005,
          internalCompaction=True, # If true the confining pressure is generated by growing particles
  )

  newton=NewtonIntegrator(damping=0.2)

  O.engines=[
          ForceResetter(),
          InsertionSortCollider([Bo1_Sphere_Aabb(),Bo1_Box_Aabb()]),
          InteractionLoop(
                  [Ig2_Sphere_Sphere_ScGeom(),Ig2_Box_Sphere_ScGeom()],
                  [Ip2_FrictMat_FrictMat_FrictPhys()],
                  [Law2_ScGeom_FrictPhys_CundallStrack()],label="iloop"
          ),
          FlowEngine(dead=1,label="flow"),#introduced as a dead engine for the moment, see 2nd section
          GlobalStiffnessTimeStepper(active=1,timeStepUpdateInterval=100,timestepSafetyCoefficient=0.8),
          triax,
          newton
  ]

  triax.goal1=triax.goal2=triax.goal3=-10000

  while 1:
    O.run(1000, True)
    unb=unbalancedForce()
    if unb<0.001 and abs(-10000-triax.meanStress)/10000<0.001:
      break

  setContactFriction(radians(finalFricDegree))

  radius=0
  for b in O.bodies:
    if b.state.mass==0:
        b.state.blockedDOFs='xyzXYZ'
        b.state.vel=(0,0,0)
        b.state.angVel=(0,0,0)
    if b.state.mass>0:
        radius+=b.shape.radius
    #    b.state.blockedDOFs='xyz'
     #   b.state.vel=(0,0,0)
  radius=radius/num_spheres

  triax.dead=True
  while 1:
    O.run(1000, True)
    unb=unbalancedForce()
    if unb<0.001:
      break

  press=1000.
  O.run(10,1)
  flow.dead=0
  flow.meshUpdateInterval=-1
  flow.useSolver=3
  flow.permeabilityFactor=1
  flow.viscosity=0.1
  flow.bndCondIsPressure=[0,0,1,0,0,0]
  flow.bndCondValue=[0,0,press,0,0,0]
  flow.boundaryUseMaxMin=[0,0,0,0,0,0]
  flow.iniVoidVolumes=True
  newton.damping=0.1
  GlobalStiffnessTimeStepper.dead=True
  O.dt=min(0.8*PWaveTimeStep(),0.8*1./1200.*pi/flow.viscosity*graindensity*radius**2)
  O.dynDt=False

  O.run(1,1)

  voidvol=0.0
  voidvoltot=0.0
  nvoids=flow.nCells()
  totalflux=[0] * (nvoids)
  initialvol=[0] * (nvoids)
  bar=[0] * (nvoids)
  initiallevel=O.bodies[flow.wallIds[flow.ymin]].state.pos[1]+(O.bodies[flow.wallIds[flow.ymax]].state.pos[1]-O.bodies[flow.wallIds[flow.ymin]].state.pos[1])/3
  initialymax=O.bodies[flow.wallIds[flow.ymax]].state.pos[1]

  for ii in range(nvoids):
     initialvol[ii]=1./flow.getCellInvVoidVolume(ii)
     bar[ii]=flow.getCellBarycenter(ii)[1]
     voidvoltot+=initialvol[ii]
     if bar[ii]>=initiallevel:
        voidvol+=initialvol[ii]
        flow.setCellPImposed(ii,True)
        flow.setCellPressure(ii,0)

  O.run(1,1)
  iterini=O.time
  flow.saveVtk("./vtk",True)
  for ii in range(nvoids):
    if bar[ii]>=initiallevel:
      if flow.getCellPImposed(ii)==True:
        #totalflux[ii]=flow.getCellFluxFromId(ii)*O.dt
        totalflux[ii]=0.0*O.dt
  #Qx=(flow.getBoundaryFlux(flow.wallIds[flow.xmin]))*O.dt+(flow.getBoundaryFlux(flow.wallIds[flow.xmax]))*O.dt
  #Qz=(flow.getBoundaryFlux(flow.wallIds[flow.zmin]))*O.dt+(flow.getBoundaryFlux(flow.wallIds[flow.zmax]))*O.dt
  #Qin=(flow.getBoundaryFlux(flow.wallIds[flow.ymin]))*O.dt
  #Qout=(flow.getBoundaryFlux(flow.wallIds[flow.ymax]))*O.dt
  Qin=0.0*O.dt
  Qout=0.0*O.dt
  Qx=0.0*O.dt
  Qz=0.0*O.dt

  bubble=0
  deltabubble=0
  deltaovercross=0
  deltaagain=0
  def pressureImbibition():
     global bubble,deltabubble,deltaovercross,deltaagain
     global Qin,Qout,Qz,Qx,totalflux,total
     Qin+=(flow.getBoundaryFlux(flow.wallIds[flow.ymin]))*O.dt
     Qout+=(flow.getBoundaryFlux(flow.wallIds[flow.ymax]))*O.dt
     Qx+=(flow.getBoundaryFlux(flow.wallIds[flow.xmin]))*O.dt+(flow.getBoundaryFlux(flow.wallIds[flow.xmax]))*O.dt
     Qz+=(flow.getBoundaryFlux(flow.wallIds[flow.zmin]))*O.dt+(flow.getBoundaryFlux(flow.wallIds[flow.zmax]))*O.dt

     for ii in range(nvoids):
        if bar[ii]>=initiallevel:
           if flow.getCellPImposed(ii)==True:
              totalflux[ii]+=flow.getCellFluxFromId(ii)*O.dt

     for ii in range(nvoids):
        if bar[ii]>=initiallevel:
           if flow.getCellPImposed(ii)==True:
              if (-totalflux[ii])>initialvol[ii]:
                 deltaflux=-totalflux[ii]-initialvol[ii]
                 kk=0
                 neigh=[nvoids+2,nvoids+2,nvoids+2,nvoids+2]
                 neighok=[nvoids+2,nvoids+2,nvoids+2,nvoids+2]
                 for jj in range(4):
                   if jj<=len(flow.getNeighbors(ii))-1:
                     if flow.getCellPImposed(flow.getNeighbors(ii)[jj])==True:
                       if (-totalflux[flow.getNeighbors(ii)[jj]])<initialvol[flow.getNeighbors(ii)[jj]]:
                         neigh[kk]=flow.getNeighbors(ii)[jj]
                         kk=kk+1
                 if kk==0:
                   totalflux[ii]+=deltaflux
                   bubble=bubble+1
                   deltabubble+=deltaflux
                 if kk!=0:
                   totalflux[ii]+=deltaflux
                   deltafluxEach=deltaflux/kk
                   deltadelta=0
                   for xx in range(4):
                     if kk!=0:
                       if neigh[xx]<=nvoids:
                         jj=neigh[xx]
                         if (-totalflux[jj]+deltafluxEach)>initialvol[jj]:
                           deltadelta+=-totalflux[jj]-initialvol[jj]
                           totalflux[jj]+=-(-totalflux[jj]-initialvol[jj])
                           kk=kk-1
                           neighok[kk]=jj
                   if kk==0:
                     deltaflux2=deltaflux-deltadelta
                     kk2=0
                     neigh2=[nvoids+2]*(4*4)
                     neighok2=[nvoids+2]*(4*4)
                     for xx in range(4):
                        if neigh[xx]<=nvoids:
                          jj=neigh[xx]
                          for gg in range(4):
                            if gg<=len(flow.getNeighbors(jj))-1:
                             if flow.getCellPImposed(flow.getNeighbors(jj)[gg])==True:
                                 if (-totalflux[flow.getNeighbors(jj)[gg]])<initialvol[flow.getNeighbors(jj)[gg]]:
                                    neigh2[kk2]=flow.getNeighbors(jj)[gg]
                                    kk2=kk2+1
                     if kk2==0:
                        bubble=bubble+1
                        deltabubble+=deltaflux2
                     if kk2!=0:
                        deltafluxEach2=deltaflux2/kk2
                        for xx in range(16):
                           if kk2!=0:
                             if neigh2[xx]<=nvoids:
                                gg=neigh2[xx]
                                if (-totalflux[gg]+deltafluxEach2)>initialvol[gg]:
                                   print('check')
                                if (-totalflux[gg]+deltafluxEach2)<=initialvol[gg]:
                                   deltadelta+=deltafluxEach2
                                   totalflux[gg]+=-deltafluxEach2
                                   kk2=kk2-1
                                   neighok2[kk2]=gg
                        if deltaflux!=deltadelta:
                            print('overcross')
                            deltaovercross+=+(deltaflux-deltadelta)
                   if kk!=0:
                     deltafluxEach2=(deltaflux-deltadelta)/kk
                     for xx in range(4):
                      if kk!=0:
                       if neigh[xx]<=nvoids:
                         jj=neigh[xx]
                         if jj!=neighok[0] and jj!=neighok[1] and jj!=neighok[2] and jj!=neighok[3]:
                            if (-totalflux[jj]+deltafluxEach2)<=initialvol[jj]:
                                totalflux[jj]+=-deltafluxEach2
                                deltadelta+=deltafluxEach2
                                kk=kk-1
                                neighok[kk]=jj
                     if deltaflux!=deltadelta:
                      print('again')
                      deltaagain+=+(deltaflux-deltadelta)

     total=0
     for ii in range(nvoids):
        total+=totalflux[ii]
        if bar[ii]>=initiallevel:
           if flow.getCellPImposed(ii)==True:
              if -(totalflux[ii])==initialvol[ii]:
                 flow.setCellPImposed(ii,False)
                 #print 'pp_',ii
              if -(totalflux[ii])>initialvol[ii]:
                 flow.setCellPImposed(ii,False)
                 print('error_',ii)
     total=abs(total)

   #  for ii in range(nvoids):
    #    if bar[ii]>=initiallevel:
  #        if flow.getCellPImposed(ii)==False:
    #          if -(totalflux[ii])!=initialvol[ii]:
     #            print 'error_',ii


  file=open('Test.txt',"w")
  checkdifference=0
  def equilibriumtest():
     global F33,F22,checkdifference,errors
     #unbalanced=unbalancedForce()
     F33=abs(O.forces.f(flow.wallIds[flow.ymax])[1])
     F22=abs(O.forces.f(flow.wallIds[flow.ymin])[1])
     #F11 =abs(O.forces.f(flow.wallIds[flow.xmax])[0]),
     #F00=abs(O.forces.f(flow.wallIds[flow.xmin])[0]),
     #F44=abs(O.forces.f(flow.wallIds[flow.zmin])[2]),
     #F55=abs(O.forces.f(flow.wallIds[flow.zmax])[2]),
     deltaF=abs(F33-F22)
     file.write(str(O.iter)+" "+str(deltaF)+"\n")
     if O.time>=iterini+1.5:
       if checkdifference==0:
         print('check F done')
         if deltaF>0.01*press:
           raise YadeCheckError('Error: too high difference between forces acting at the bottom and upper walls')
           #O.pause()
         checkdifference=1


  once=0
  def fluxtest():
     global once,QinOk
     no=0
     QinOk=-Qin-deltabubble
     error=QinOk-total
     if error>toleranceWarning:
        print("Warning: difference between total water volume flowing through bottom wall and water loss due to air bubble generations",QinOk," vs. total water volume flowing inside dry or partially saturated cells",total)
     if error>toleranceCritical:
        raise YadeCheckError("The difference is more, than the critical tolerance!")
     file.write(str(O.time)+" "+str(QinOk)+" "+str(error)+" \n")
     for ii in range(nvoids):
       if bar[ii]>=initiallevel:
        if flow.getCellPImposed(ii)==True:
           no=1
     if once==0:
        if no==0:
           imbtime=O.time-iterini
           print(imbtime,voidvol,total,QinOk)
           if voidvol-total>toleranceWarning:
             print("Warning: initial volume of dry voids",voidvol," vs. total water volume flowing inside dry or partially saturated cells",total)
           if voidvol-total>toleranceCritical:
             raise YadeCheckError("The difference is more, than the critical tolerance!")
           file.write(str(imbtime)+" "+str(voidvol)+" "+str(total)+" "+str(QinOk)+"\n")
           once=1
           timing.stats()


  def addPlotData():
     global F33,F22,QinOk,total
     plot.addData(i1=O.iter,
        t=O.time,
        Fupper=F33,
        Fbottom=F22,
        Q=QinOk,
        T=total
     )
  plot.live=True
  plot.plots={'t':('Fbottom','Fupper'),' t ':('Q','T')}
  plot.plot()


  def pl():
     flow.saveVtk("./vtk",True)

  O.engines=O.engines+[PyRunner(iterPeriod=100,command='pl()')]
  #O.engines=O.engines+[VTKRecorder(iterPeriod=100,recorders=['spheres'],fileName='./exp')]
  O.engines=O.engines+[PyRunner(iterPeriod=1,command='equilibriumtest()')]
  O.engines=O.engines+[PyRunner(iterPeriod=1,command='pressureImbibition()')]
  O.engines=O.engines+[PyRunner(iterPeriod=1,command='fluxtest()')]
  O.engines=O.engines+[PyRunner(iterPeriod=1,command='addPlotData()')]

  O.timingEnabled=True
  #file.close()
  #plot.saveDataTxt('plots.txt',vars=('i1','t','Fupper','Fbottom','Q','T'))

  import tempfile, shutil
  dirpath = tempfile.mkdtemp()
  for fileName in ['./vtk', './Test.txt' ]:
    if (os.path.exists(fileName)): shutil.move(fileName,dirpath)
    print("File %s moved into %s/ directory"%(fileName,dirpath))

else:
    print("This checkDEM-PFVPressureInjection.py cannot be executed because PFVFLOW is disabled")

