# IBM preprocessing on Cartesian grids for FastS
import Converter.PyTree as C
import Generator.PyTree as G
import Converter.Internal as Internal
import Post.PyTree as P
import Transform.PyTree as T
import Initiator.PyTree as I
import Fast.PyTree as Fast
import Fast.Utils as Utils
import Connector.ToolboxIBM as IBM
import Dist2Walls.PyTree as DTW

# Ce script est sequentiel. Si NP > 0, on prepare pour un
# calcul parallele sur NP procs
NP = Utils.getArgs1() # Mpi NP > 0
if NP > 0: print 'Preparing for %d processes.'%NP
if NP > 0: import Distributor2.PyTree as D2
split = 'single' # 'multiple' : one file per proc
#=======================================================
# Input data
#=======================================================
# vmin: nb of pts per Cartesian grid
# DEPTH: nb of interpolated and ghost cells
vmin = 15; 
DEPTH = 2
# list of body surfaces as separated CGNS bases
bodySurfaceFile = 'case.cgns'
refineFinestLevel=False
refineNearBodies=False
factor=1.2
check = True # True: some intermediate files are written for checks
#=======================================================
# End of input data
#=======================================================
#-------------------------------------------------------
# Read body mesh
tb = C.convertFile2PyTree(bodySurfaceFile)
#--------------------------------------------------------
# Get Reference State and model from body pyTree
model = Internal.getNodeFromName(tb, 'GoverningEquations')
if model is None: raise ValueError, 'GoverningEquations is missing in input cgns.'
# model: Euler, NSLaminar, NSTurbulent
model = Internal.getValue(model)

# reference state
refstate = C.getState(tb)
# dimension du pb
dimPb = Internal.getNodeFromName(tb, 'EquationDimension')
dimPb = Internal.getValue(dimPb)
if dimPb == 2: C._initVars(tb, 'CoordinateZ=0.') # forced

#--------------------------------------------------------
# Generates the full Cartesian mesh
t = C.convertFile2PyTree("restart.cgns")
RefState = Internal.getNodeFromName(t,'ReferenceState')
Gamma = Internal.getNodeFromName(RefState,"Gamma")
Gamma = Internal.getValue(Gamma)
CvInf = Internal.getNodeFromName(RefState,"Cv")
CvInf = Internal.getValue(CvInf)
RGP = (Gamma-1.)*CvInf
C._initVars(t,'centers:Pressure={centers:Density}*{centers:Temperature}*%g'%RGP)
C._initVars(t,'centers:U2={centers:VelocityX}**2+{centers:VelocityY}**2+{centers:VelocityZ}**2')
C._initVars(t,'centers:a={centers:Density}/({centers:Pressure}*%g)'%Gamma)
C._initVars(t,'centers:Mach=sqrt({centers:U2}/{centers:a})')
C._rmVars(t,["centers:a","centers:U2","centers:Pressure"])
t = P.computeDiff(t,"centers:Mach")
if check: C.convertPyTree2File(t,"in.cgns")
C._rmVars(t,["centers:Mach"])
t = IBM.adaptIBMMesh(t, tb, vmin, sensor='centers:diffMach', factor=factor, \
                         refineFinestLevel=refineFinestLevel, \
                         refineNearBodies=refineNearBodies,\
                         variables=['centers:Density','centers:VelocityX','centers:VelocityY','centers:VelocityZ','centers:Temperature'])

#------------------------------------------------------
# distribute the mesh over NP processors
if NP > 0:
    print 'distribution over %d processors'%NP
    t, stats = D2.distribute(t, NP)
    if check: print stats

#------------------------------------------------
# Add reference state to the pyTree and init fields
# Add viscosity if model is not Euler
if model != "Euler": C._initVars(t, 'centers:ViscosityEddy', 0.)
C._addState(t, state=refstate)
C._addState(t, 'GoverningEquations', model)
C._addState(t, 'EquationDimension', dimPb)

#----------------------------------------
# Computes distance field
#----------------------------------------
if dimPb == 2:
    z0 = Internal.getNodeFromType2(t,"Zone_t")
    bb = G.bbox(z0); dz = bb[5]-bb[2]
    tb2 = C.initVars(tb,'CoordinateZ',dz*0.5)
    t = DTW.distance2Walls(t,tb2,type='ortho',signed=0, dim=dimPb,loc='centers')
else:
    t = DTW.distance2Walls(t,tb,type='ortho',signed=0, dim=dimPb,loc='centers')
    
#----------------------------------------
# Create IBM info
#----------------------------------------
t,tc = IBM.prepareIBMData(t, tb, frontType=1)

#----------------------------------------
# arbre donneur
#----------------------------------------
if NP > 0:
    D2._copyDistribution(tc,t)
Fast.save(tc, 'tc.cgns', split='single', NP=-NP)

#----------------------------------------
# Extraction des coordonnees des pts IBM
#----------------------------------------
if check:
    tibm = IBM.extractIBMInfo(tc)
    C.convertPyTree2File(tibm, 'IBMInfo.cgns')

#----------------------------------------
# arbre de calcul 
#----------------------------------------
del tc # to free memory before initialization
Fast.save(t, 't.cgns', split=split, NP=-NP)