import Converter.PyTree as C
import Post.PyTree as P
import Connector.PyTree as X
import Connector.ToolboxIBM as IBM
import Transform.PyTree as T
import Converter.Internal as Internal
import Fast.PyTree as Fast
import Fast.Utils as Utils
import FastS.PyTree as FastS

#==========================================================
# Body surface
tb = C.convertFile2PyTree("tryIBMWorkflow.cgns")

# fichier tc_restart.cgns a ete sauvegarde et contient les champs pression/masse volumique aux points IBM
t,tc,ts,graph=Fast.load('restart', 'tc_restart', split='single')

#==========================================================
# Extraction Cp, Cf, ... sur les surfaces par interpolation
#==========================================================
#--------------------------------------------------------
# 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)

if model == 'Euler': bcType = 0
else: bcType = 3 # Musker

dimPb = Internal.getNodeFromName(tb, 'EquationDimension')
dimPb = Internal.getValue(dimPb)

[RoInf, RouInf, RovInf, RowInf, RoeInf, PInf, TInf, cvInf, MInf,
 ReInf, Cs, Gamma, RokInf, RoomegaInf, RonutildeInf, Mus, Cs, Ts, Pr] = C.getState(tb)

vars = ['Density', 'VelocityX', 'VelocityY', 'VelocityZ', 'Temperature']
varsIBC = ['Density', 'VelocityX', 'VelocityY', 'VelocityZ', 'Temperature']
if model != 'Euler': vars+=['ViscosityEddy']
zw = IBM.extractIBMWallFields(tc, tb=tb)

RoUInf2I = 1./(RouInf*RouInf+RovInf*RovInf+RowInf*RowInf)

C._initVars(zw,'{Cp}=2*%f*({Pressure}-%f)*%f'%(RoInf,PInf,RoUInf2I))
if model != 'Euler': C._initVars(zw,'{Cf}=2.*%f*{Density}*{utau}**2*%f'%(RoInf,RoUInf2I))
C.convertPyTree2File(zw, 'wall.cgns')

#=================================
# Calcul de mut/mu dans le volume
#=================================
if model != 'Euler':
    betas = Mus*(Ts+Cs)/(Ts**(3./2.))
    C._initVars(t,'{centers:ViscosityMolecular} = %20.16g*sqrt({centers:Temperature})/(1.+%20.16g/{centers:Temperature})'%(betas,Cs))
    C._initVars(t,'{centers:mutsmu}=({centers:ViscosityEddy})/({centers:ViscosityMolecular})-1.')
    C._cpVars(t, 'centers:mutsmu', tc, 'mutsmu')
    X._setInterpTransfers(t, tc, variables=['mutsmu'], variablesIBC=None)
tc = None

#===============================
# En 2D, extrait un seul plan k
#================================
if dimPb == 2:
    t = T.subzone(t, (1,1,1), (-1,-1,1))
    C._initVars(tb, 'CoordinateZ', 0.) # forced

#==============================
# Sortie champs aux noeuds
#==============================
print('Extracting field in nodes')
vars = ['centers:Density','centers:VelocityX', 'centers:VelocityY', 'centers:Temperature','centers:ViscosityEddy', 
    'centers:TurbulentSANuTilde','centers:ViscosityMolecular', 'centers:mutsmu', 'centers:cellN']
for v in vars: t = C.center2Node(t, v)
Internal._rmNodesByName(t, 'FlowSolution#Centers')
C._initVars(t, '{Velocity}=sqrt({VelocityX}*{VelocityX}+{VelocityY}*{VelocityY})')
C.convertPyTree2File(t, 'out.cgns')

#=============================
# Extraction convergence
#=============================
print('Extracting residuals in residual.dat')
cvg_node = Internal.getNodeFromName(t, 'ZoneConvergenceHistory')
if cvg_node is not None: FastS.extractConvergenceHistory(t, 'residus.dat')

#=============================
# slices
#=============================
print('Extracting slice')
p = P.isoSurfMC(t, 'CoordinateY', 0.5)
C.convertPyTree2File(p, "slice.cgns")