# Entree: 
# - un maillage 2D (un seul plan) multibloc
# - ou un maillage 3D multibloc 
# avec BCs + raccords + (solution initiale et/ou Reference State)
# Sortie: t.cgns et tc.cgns pour compute.py

import Converter.PyTree as C
import Generator.PyTree as G
import Transform.PyTree as T
import Converter.Internal as Internal
import Connector.PyTree as X
import Fast.PyTree as Fast
import Fast.Utils as Utils

NP = Utils.getArgs1()
if NP > 0: print('Preparing for %d processes.'%NP)
if NP > 0: import Distributor2.PyTree as D2

FILE = 'tryMBWorkflow.cgns'

# INPUT: maillage + BC + GC
t = C.convertFile2PyTree(FILE)

# Get dim
dim = 3
node = Internal.getNodeFromName(t, 'EquationDimension')
if node is not None: dim = Internal.getValue(node)

# Split size
if NP > 0:
    t = T.splitSize(t, R=NP, type=2, minPtsPerDir=9)
    t = X.connectMatch(t, dim=dim)
    stats = D2._distribute(t, NP, useCom='match')

# Solution initiale
eqs = Internal.getNodeFromType(t, 'GoverningEquations_t')
Model = 'NSTurbulent'
if eqs is not None: Model = Internal.getValue(eqs)

ret = C.isNamePresent(t, 'centers:Density')
if ret != 1: # Density not present
    state = Internal.getNodeFromType(t, 'ReferenceState_t')
    if state is None:
        raise ValueError('Reference state is missing in input cgns.')
    vars = ['Density', 'MomentumX', 'MomentumY', 'MomentumZ',
            'EnergyStagnationDensity']
    for v in vars:
        node = Internal.getNodeFromName(state, v)
        if node is not None:
            val = float(node[1][0])
            C._initVars(t, 'centers:'+v, val)
        else:
            raise ValueError(v + ' is missing in ReferenceState.')
    if Model == 'NSTurbulent':
        vars = ['TurbulentSANuTildeDensity']
        for v in vars:
            node = Internal.getNodeFromName(state, v)
            if node is not None:
                val = float(node[1][0])
                C._initVars(t, 'centers:'+v, val)
ret = C.isNamePresent(t, 'centers:TurbulentDistance')
if Model == 'NSTurbulent' and ret != 1: # Wall distance not present
    import Dist2Walls.PyTree as DTW
    walls = C.extractBCOfType(t, 'BCWall')
    DTW._distance2Walls(t, walls, loc='centers', type='ortho')

# Ajout des ghost-cells
C.addState2Node__(t, 'EquationDimension', dim)
Internal._addGhostCells(t, t, 2, adaptBCs=1, fillCorner=0)

if dim == 2: 
    T._addkplane(t)
    T._contract(t, (0,0,0), (1,0,0), (0,1,0), 0.01)
    T._makeDirect(t)

# Raccords
tc = C.node2Center(t)
tc = X.setInterpData(t, tc, nature=1, loc='centers', storage='inverse', 
                     sameName=1, dim=dim)
C._rmVars(tc, 'FlowSolution')
Fast.save(t, 't.cgns', split='single', NP=-NP)

# Copie la distribution de t dans tc
if NP > 0: D2._copyDistribution(tc,t)
C.convertPyTree2File(tc, 'tc.cgns')