#!/usr/bin/python
import parmepy as pp
import parmepy.f2py
scales = parmepy.f2py.scales2py
import numpy as np
import mpi4py.MPI as MPI
import math

pi = math.pi

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
print "start simulation for process number", rank

pbDim = 3

# Number of MPI processus
nbProcs = comm.Get_size()

# Physical Domain description
Lx = Ly = Lz = 2 * pi
myDomain3d = pp.Box(dimension=3, length=[Lx, Ly, Lz], origin=[0., 0., 0.])
nbCells = np.asarray((16, 16, 32))

# ===== Initialize the particular solver =====
# First choose the number of mpi process in each dir
#topodims = MPI.Compute_dims(nbProcs, 3)
if nbProcs > 1:
    topodims = [1, 2, nbProcs / 2]
else:
    topodims = [1, 1, 1]

# It must be such that nbCells is a multiple of topodims
# and such that nbCells/topodims is a multiple of 2, 4 or 5.
localres, localoffset, stab_coeff = \
    scales.init_advection_solver(nbCells, myDomain3d.length,
                                 topodims, order='p_O2')

if(rank == 0):
    print "Start initialisation ..."

# ===== Create and initialize Fields =====
#scal3D = np.zeros((localres), dtype='float64', order='Fortran')
scal3D = np.zeros((localres), dtype='float64')

r02 = ((myDomain3d.length / 10.).min()) ** 2  # ??

# Coordinates of the lowest point of the current subdomain
step = myDomain3d.length / nbCells
coord0 = localoffset * localres * step
vz = np.zeros((localres), dtype='float64', order='Fortran')
vx = np.zeros((localres), dtype='float64', order='Fortran')  # vz.copy()
vy = np.zeros((localres), dtype='float64', order='Fortran')  # vz.copy()
period = 1.0
coef = 2. * pi / period

for k in range(localres[2]):
    rz = (step[2] * k + coord0[2]) ** 2
    for j in range(localres[1]):
        ry = (step[1] * j + coord0[1]) ** 2
        for i in range(localres[0]):
            rx = (step[0] * i + coord0[0]) ** 2
            rr = rx + ry + rz
            if(rr < r02):
                scal3D[i, j, k] = (1 - rr / r02) ** 4
            vx[i, j, k] = coef * (0.5 - coord0[1] - step[1] * j)
            vy[i, j, k] = coef * (step[0] * i + coord0[0] - 0.5)

goodScal = scal3D.copy()
goodVelo = vx.copy()

time_step = 0.01
finalTime = 2. * period
currentTime = 0.0

comm.Barrier()

if(rank == 0):
    print "End of initialisation ..."

#while currentTime<finalTime:

# Simulation
t0 = MPI.Wtime()
scal3D = scales.solve_advection(time_step, vx, vy, vz, scal3D)

print "elapsed time on processus ", rank, ": ", MPI.Wtime() - t0

currentTime += time_step