#
# Example of field creation/discretisation and output
# to test post-treatment (vtk, hdf5 ...)
#
# Usage :
# python testVisu.py
# or
# mpirun -np NB python testVisu.py
# NB = number of mpi processes


import parmepy as pp
import math
from parmepy.fields.continuous import Field
from parmepy.domain.obstacle.sphere import Sphere
from parmepy.operator.analytic import Analytic
from parmepy.domain.obstacle.planes import SubPlane
pi = math.pi
from parmepy.operator.penalization import Penalization
from parmepy.operator.monitors.printer import Printer
from parmepy.problem.simulation import Simulation
import numpy as np
cos = np.cos
sin = np.sin
nb = 5
Lx = Ly = Lz = 2


## ======== 3D ======== ##

# Domain (box-shaped)
dom = pp.Box(dimension=3, length=[Lx, Ly, Lz], origin=[-1., -1., -1.])

# global resolution for the grid
resol3D = [nb, nb, nb]

# Continuous fields (a scalar and a vector)


# Example of function to init the scalar field
def func3D(res, x, y, z, t):
    res[0][...] = 1.#cos(t * x) + sin(y) + z
    return res


scal3D = Field(domain=dom, formula=func3D, name='Scal')
scalRef = Field(domain=dom, formula=func3D, name='Ref')
resolution3D = np.ones(3) * nb
ghosts = np.zeros(3) + 2
from parmepy.mpi.topology import Cartesian
topo = Cartesian(dom, dom.dimension, resolution3D, ghosts=ghosts)

scalRef.discretize(topo)
scalRef.initialize(topo=topo)
hh = topo.mesh.space_step[2]
planes = SubPlane(dom, normal=[0, 0, 1], point=[0, -2*hh, 2*hh],
                  lengths=[0.8, 0.8, 0.5])

op = Analytic(scal3D, resolutions={scal3D: resolution3D},
              topo=topo)
op.discretize()
op.setUp()
simu = Simulation(start=0.0, end=2., nbIter=100, max_iter=1000000)
op.apply(simu)

sc3D = scal3D.discretization(topo)
scRef = scalRef.discretization(topo)

assert np.allclose(sc3D.data[0], scRef.data[0])
## # Monitor : print scal3D to vtk file.
from parmepy.constants import HDF5
printer = Printer(variables=[scal3D], topo=topo,
                  formattype=HDF5, prefix='testSlice', subset=planes)
printerRef = Printer(variables=[scalRef], topo=topo,
                     formattype=HDF5, prefix='testRef')
printer.setUp()
printerRef.setUp()
printerRef.apply(simu)

simu.initialize()

while not simu.isOver:
    op.apply(simu)
#    printer.apply(simu)
    printerRef.apply(simu)
    simu.advance()
   # simu.isOver=True
printerRef.finalize()
sc3D.data[0][...] = 0.0

## Example to postprocess some hdf5 files and save only a required section : 

import glob
import parmepy.tools.io_utils as io
filepath = io.io.defaultPath()
prefix = '/testRef'
# The list of all file with testRef_...
filelist = glob.glob(filepath + prefix + '*.h5')
print filepath + prefix
filelist.sort()
print filelist
# for each file in dir ...
# 1 - Dowload the full-grid hdf file and use it to fill in a field
# First iter:
from parmepy.operator.monitors.reader import Reader
import os
name = os.path.basename(filelist[0])
name = name.rsplit('.')[0]

reader0 = Reader(variables=[scal3D], topo=topo, prefix=name)
reader0.setUp()
reader0.apply()
printer.apply(simu)
simu.initialize()
for pref in filelist[1:]:
    name = os.path.basename(pref)
    name = name.rsplit('.')[0]
    reader = Reader(variables=[scal3D], topo=topo, prefix=name)
    reader.setUp()
    simu.printState()
        #print 'pre', scal3D.norm(topo)
    reader.apply()
    # 2 - Use the predifined printer on planes to save data
    printer.apply(simu)
    simu.advance()
##     #print scalRef.norm(topo)
##     #print scal3D.norm(topo)
##     #assert np.allclose(sc3D.data[0][topo.mesh.compute_index], scRef.data[0][topo.mesh.compute_index])


velo = Field(domain=dom, name='Velocity', is_vector=True)
vorti = Field(domain=dom, name='vty', is_vector=True)
reader = Reader(variables=[velo, vorti], topo=topo, prefix='curl_io_00000')