from dgpy import *
from math import *
import time, os, sys
import gmshPartition
import numpy as np

nbSteps = int(sys.argv[1])
compError = 1

#### Physical constants ###
gamma=1.4
Rd=287.0
p0=1.0e5
g=9.80616
Cv=Rd/(gamma-1.0)
Cp=Rd*(1.0+1.0/(gamma-1.0))
###########################
#### Begin and end time ###
Ti=0.0
Tf=150.0
###########################

#Period of the oscillations
period = 100

order=4
dimension=3

model = GModel()
name='periodic_convection'
if Msg.GetCommSize()>1:
  partStr='_part_%i' % Msg.GetCommSize()
else:
  partStr=''
model.load(name + partStr + '_3d.msh')

groups = dgGroupCollection(model, dimension, order)
groups.splitGroupsByPhysicalTag();

claw = dgEulerAtmLaw(dimension)
if compError:
  ####Load reference solution from file###
  reloaded =dgDofContainer(groups,claw.getNbFields())
  reloaded.readMsh("refSolERK"+str(period)+"_16"+"/subwindow2/subwindow2.idx")
solution = dgDofContainer(groups, claw.getNbFields())

XYZ = groups.getFunctionCoordinates();

def initialCondition(cmap,FCT, XYZ) :
  FCT[:,:] = 0

def boundaryForcing(cmap, FCT,  XYZ, time, rhoHs, thetaHs) :
  thetac=5
  xc=500.0
  yc=500.0
  zc=0.0-600
  rc=250.0
  currentTime=time[0]    
  r=np.sqrt((XYZ[:,0]-xc)**2+(XYZ[:,1]-yc)**2+(XYZ[:,2]-zc)**2)
  thetaPert = np.empty(thetaHs.shape)
  outR = (r>rc)
  thetaPert[outR]=thetaHs[outR]
  inR=np.invert(outR)
  thetaPert[inR]=thetaHs[inR]+thetac/2.0*(1.0+np.cos(pi*r[inR]/rc))*np.sin(currentTime/period*2.0*pi)**2
  FCT[:,0]=0
  FCT[:,1]=0
  FCT[:,2]=0
  FCT[:,3]=0
  FCT[:,4]=rhoHs*thetaPert-rhoHs*thetaHs

def thetaHydrostatic(cmap, FCT) :
  FCT[:]=300

def exnerHydrostatic(cmap, FCT,  XYZ, thetaHs) :
  FCT[:]=1.0-g/(Cp*thetaHs)*XYZ[:,2]

def rhoHydrostatic(cmap, FCT, thetaHs, exnerHs) :
  FCT[:]=p0/(Rd*thetaHs)*exnerHs**(Cv/Rd)

def rhoThetaHydrostatic(cmap, FCT, rhoHs, thetaHs) :
  FCT[:]=rhoHs*thetaHs
        
def getVelocity(cmap, FCT, sol, rhoHs) :
  invrho=1/(rhoHs+sol[:,0])
  FCT[:,0]=sol[:,1]*invrho
  FCT[:,1]=sol[:,2]*invrho
  FCT[:,2]=sol[:,3]*invrho

def getRhop(cmap, FCT, sol) :
  FCT[:]=sol[:,0]

def getpp(cmap, FCT, sol, rhoHs, thetaHs) :
  rho=rhoHs+sol[:,0]
  rhoTheta = rhoHs*thetaHs+sol[:,4]
  pHs=p0*(rhoHs*thetaHs*Rd/p0)**gamma
  p=p0*(rhoTheta*Rd/p0)**gamma
  FCT[:]=p-pHs

def getpHs(cmap, FCT, sol, rhoHs, thetaHs) :
  pHs=p0*(rhoHs*thetaHs*Rd/p0)**gamma
  FCT[:]=pHs

def getp(cmap, FCT, sol, rhoHs, thetaHs) :
  rhoTheta = rhoHs*thetaHs+sol[:,4]
  p=p0*(rhoTheta*Rd/p0)**gamma
  FCT[:]=p

def getThetap(cmap, FCT, sol, rhoHs, thetaHs) :
  rho=rhoHs+sol[:,0]
  FCT[:]=(sol[:,4]+(rhoHs*thetaHs))/rho-thetaHs

thetaHsC = functionNumpy(1, thetaHydrostatic,[])
exnerHsC = functionNumpy(1, exnerHydrostatic, [XYZ, thetaHsC])
rhoHsC = functionNumpy(1, rhoHydrostatic, [thetaHsC, exnerHsC])
rhoThetaHsC = functionNumpy(1, rhoThetaHydrostatic, [rhoHsC, thetaHsC])

uvw=functionNumpy(3, getVelocity, [solution.getFunction(), rhoHsC])
rhop=functionNumpy(1, getRhop, [solution.getFunction()])
pp=functionNumpy(1, getpp, [solution.getFunction(), rhoHsC, thetaHsC])
thetap=functionNumpy(1, getThetap, [solution.getFunction(), rhoHsC, thetaHsC])

initF=functionNumpy(claw.getNbFields(), initialCondition, [XYZ])
solution.interpolate(initF)

rhoHs = dgDofContainer(groups, 1)
rhoHs.interpolate(rhoHsC)
rhoThetaHs = dgDofContainer(groups, 1)
rhoThetaHs.interpolate(rhoThetaHsC)

claw.setHydrostaticState(rhoHs,rhoThetaHs)
claw.setPhysicalConstants(gamma,Rd,p0,g)

nuh = functionConstant(.4);
nuv = functionConstant(.4);
zero = functionConstant(0);
claw.setNu(nuh, nuv, zero)

boundaryWall = claw.newBoundaryWall()
timeFunction = function.getTime()
boundaryForcingF=functionNumpy(5, boundaryForcing, [XYZ, timeFunction, rhoHsC, thetaHsC])
toReplace = VectorFunctorConst([function.getSolution()])
replaceBy = VectorFunctorConst([boundaryForcingF])
boundaryHeated = claw.newOutsideValueBoundaryGeneric("",toReplace,replaceBy)
claw.addBoundaryCondition('none', boundaryWall)
claw.addBoundaryCondition('bottom', boundaryWall)
claw.addBoundaryCondition('left', boundaryWall)
claw.addBoundaryCondition('right', boundaryWall)
claw.addBoundaryCondition('top', boundaryWall)
claw.addBoundaryCondition('bottom_surf', boundaryWall)
claw.addBoundaryCondition('top_surf', boundaryWall)
claw.addBoundaryCondition('bottom_bottom_surf', boundaryHeated)
claw.addBoundaryCondition('top_bottom_surf', boundaryWall)


#petsc = linearSystemPETScBlockDouble()
#petsc.setParameter("petscOptions",  "-pc_type lu -ksp_type preonly -pc_factor_mat_solver_package superlu_dist")

##petsc.setParameter("petscOptions",  "-ksp_atol 1e-12 -ksp_rtol 1e-12")
#if (Msg.GetCommSize()==1):
#    petsc.setParameter("petscOptions",  "-pc_type lu")#-ksp_atol 1e-12 -ksp_rtol 1e-12")
#dof = dgDofManager.newDGBlock(groups, claw.getNbFields(), petsc)
#timeIter = dgIMEXRK(claw, dof, 2)     #timeorder
#timeIter.getNewton().setVerb(1)     #Verbosity


claw.setFilterMode(FILTER_LINEAR);
#petscIm = linearSystemPETScBlockDouble()

petscIm = linearSystemPETScDouble()
#petscIm.setParameter("petscOptions",  "-ksp_type preonly -pc_type lu")
#if (Msg.GetCommSize()>1):
#  Msg.Fatal("Direct LU is for serial jobs only. Please change the solver or run serial")
#petscIm.setParameter("petscOptions",  "-ksp_atol 1.e-6 -ksp_rtol 1.e-9 -ksp_type gmres -ksp_gmres_restart 100 -pc_type jacobi -pc_factor_levels 6")
#petscIm.setParameter("petscOptions",  "-ksp_atol 1e-6 -ksp_rtol 1e-9 -pc_type asm")
#petscIm.setParameter("petscOptions",  "-ksp_atol 1e-6 -ksp_rtol 1e-9 -pc_type hypre -pc_hypre_type boomeramg -pc_hypre_boomeramg_max_iter 2 -pc_hypre_boomeramg_rtol 0 -ksp_monitor -log_summary")
#petscIm.setParameter("petscOptions",  "-ksp_atol 1e-6 -ksp_rtol 1e-9 -pc_type hypre -pc_hypre_type euclid -ksp_monitor -log_summary")
#petscIm.setParameter("petscOptions",  "-ksp_atol 1e-6 -ksp_rtol 1e-9")
#petscIm.setParameter("petscOptions",  "-ksp_atol 1e-6 -ksp_rtol 1e-9  -pc_type asm -pc_factor_in_place")

#mat_mkl_pardiso_65 : number of threads
#mat_mkl_pardiso_68 : verbosity level
#mat_mkl_pardiso_60 : out-of-core
#petscIm.setParameter("petscOptions",  "-ksp_type preonly -pc_type lu -pc_factor_mat_solver_package mkl_pardiso -mat_mkl_pardiso_68 0")

petscIm.setParameter("petscOptions",  "-ksp_type preonly -pc_type lu -pc_factor_mat_solver_package mumps")

# Out o core -mat_mumps_icntl_22 1

#dofIm = dgDofManager.newDGBlock(groups, claw.getNbFields(), petscIm)
dofIm = dgDofManager.newDG(groups, claw, petscIm)
#Last argument: compute sparsity from mesh? If not, computed from equation
timeIter = dgIMEXRK(claw, dofIm, IMEX_ARK3, None, True)     #timeorder 
timeIter.getNewton().setVerb(1)     #Verbosity

#if ( os.path.isfile("/scratch/sblaise/matrix.bin") ):
#  if (Msg.GetCommRank()==0):
#    print("Reloading matrix...")
#  petscIm.loadBinaryMat("/scratch/sblaise/matrix.bin")
#  if (Msg.GetCommRank()==0):
#    print("Done")


#timeIter = dgERKLinear(claw, groups, DG_ERK_22)

if compError: 
  # Compute relative error
  def getRefSolSq(cmap, FCT,  refSol) :
    FCT[:,:] = refSol[:,:]**2;
  solRefSq=functionNumpy(5, getRefSolSq, [reloaded.getFunction()])
  #
  def getError(cmap, FCT, refSol, compSol) :
    FCT[:,:] = (refSol[:,:]-compSol[:,:])**2
  error = functionNumpy(5, getError, [reloaded.getFunction(), solution.getFunction()])
  integratorError = dgFunctionIntegrator(groups, error)
  intErr = fullMatrixDouble(5,1)
  integratorRefSq = dgFunctionIntegrator(groups, solRefSq)
  intRefSq = fullMatrixDouble(5,1)

#Export data
def getExp(cmap, FCT, uvw, rhop, thetap, pp, rhoHs, thetaHs) :
  FCT[:,0]=uvw[:,0]
  FCT[:,1]=uvw[:,1]
  FCT[:,2]=uvw[:,2]
  FCT[:,3]=rhop[:]
  FCT[:,4]=thetap[:]
  FCT[:,5]=pp[:]
  FCT[:,6]=rhoHs[:]
  FCT[:,7]=thetaHs[:]
nCompExp=[3,1,1,1,1,1]
namesExp=["uvw","rhop","thetap","pp","rhoHsC", "thetaHsC"]
Exp=functionNumpy(sum(nCompExp), getExp, [uvw, rhop, thetap, pp, rhoHsC, thetaHsC])

#solution.exportFunctionVtk(Exp,'output/export', 0, 0,"solution",nCompExp,namesExp)

dt=(Tf-Ti)/nbSteps
if (Msg.GetCommRank()==0):
    print ("Time step:",dt)

t=Ti
n_export=0
start = time.clock()

for i in range(0,nbSteps):
  elapsed = (time.clock() - start)
  if (Msg.GetCommRank()==0):
    print ('[',elapsed,' s]')
  #if (i%50==0 or i==nbSteps-1):
  #  solution.exportFunctionVtk(Exp,'output/export', i, i,"solution",nCompExp,namesExp)
  #  if (Msg.GetCommRank()==0):
  #    elapsed = (time.clock() - start)
  #    print ('Exported. Elapsed time: ',elapsed,' s. Step ',i,' over ',nbSteps)
  norm = timeIter.iterate (solution, dt, t)
#  if (i==0):
#    petscIm.writeBinaryMat("/scratch/sblaise/matrix.bin")
  t=t+dt
  if (Msg.GetCommRank()==0):
    sys.stdout.write('.')
    sys.stdout.flush()
print ('')    
elapsed = (time.clock() - start)
if (Msg.GetCommRank()==0):
    print ('Time elapsed: ',elapsed,' s')

if compError:
  integratorError.compute(intErr) 
  integratorRefSq.compute(intRefSq)
  errp = sqrt(intErr(0,0)/intRefSq(0,0))
  erru = sqrt(intErr(1,0)/intRefSq(1,0))
  errv = sqrt(intErr(2,0)/intRefSq(2,0))
  errw = sqrt(intErr(3,0)/intRefSq(3,0))
  errthetap = sqrt(intErr(4,0)/intRefSq(4,0))
  errall = sqrt((intErr(0,0)+intErr(1,0)+intErr(2,0)+intErr(3,0)+intErr(4,0))/(intRefSq(0,0)+intRefSq(1,0)+intRefSq(2,0)+intRefSq(3,0)+intRefSq(4,0)))
#
  if (Msg.GetCommRank()==0):
    dataFile = file("periodic_convection_16_3d_imexark3_poly4_150s.dat","a")
    dataFile.write("%d %15f %e %e %e %e %e %e \n" % (nbSteps,elapsed,errp,erru,errv,errw,errthetap,errall))
    dataFile.close()
del petscIm
del dofIm
Msg.Exit(0)