# - compute (pyTree) -
import Fast.PyTree as Fast
import Converter.Internal as Internal
import Fast.Utils as Utils
import CPlot.PyTree as CPlot
NIT = 1000
# NP: 0: Seq, NP> 0: running with mpi, distributed on NP procs
NP = Utils.getArgs1()
if NP > 0:
    import Converter.Mpi as Cmpi
    import FastS.Mpi as FastS
    rank = Cmpi.rank
else:
    import FastS.PyTree as FastS
    rank=0

# Reading input files: t.cgns/restart.cgns and tc.cgns
t,tc,ts,graph=Fast.load('t.cgns', 'tc.cgns', split='single', 
                        restart=False, NP=NP)

it0 = 0; time0 = 0.
first = Internal.getNodeFromName1(t, 'Iteration')
if first is not None: it0 = Internal.getValue(first)
first = Internal.getNodeFromName1(t, 'Time')
if first is not None: time0 = Internal.getValue(first)

# Numerics
numb = {}
numb["temporal_scheme"]    = "implicit"
numb["ss_iteration"]       = 1
numz = {}
numz["time_step"]          = 0.002
numz["time_step_nature"]   = "local"
numz["cfl"]                = 5.
numz["scheme"]             = "roe_min"

Fast._setNum2Zones(t, numz); Fast._setNum2Base(t, numb)

(t, tc, metrics) = FastS.warmup(t, tc, graph)

time_step = Internal.getNodeFromName(t, 'time_step')
time_step = Internal.getValue(time_step)
for it in xrange(NIT):
    FastS._compute(t, metrics, it, tc, graph)
    if it%100 == 0:
        if rank == 0: print '- %d / %d - %f'%(it+it0, NIT+it0, time0)
        FastS.display_temporal_criteria(t, metrics, it, format='double')
        CPlot.display(t, dim=2, mode=3, scalarField=3)

    time0 += time_step

# time stamp
Internal.createUniqueChild(t, 'Iteration', 'DataArray_t', value=it0+NIT)
Internal.createUniqueChild(t, 'Time', 'DataArray_t', value=time0)
Fast.save(tc,'tc_restart.cgns', split='single', NP=NP)
Fast.save(t,'restart.cgns', split='single', NP=NP)