#!/usr/bin/python
#import sys
#sys.path.insert(0,'/scratch/mimeau/install-parmes3/')
import parmepy as pp
from parmepy.f2py import fftw2py
import numpy as np
import math as m
from parmepy.mpi.topology import Cartesian
from parmepy.operator.advection import Advection
from parmepy.operator.stretching import Stretching
from parmepy.operator.poisson import Poisson
from parmepy.operator.diffusion import Diffusion
from parmepy.operator.redistribute import Redistribute
from parmepy.problem.navier_stokes import NSProblem
from parmepy.operator.monitors.printer import Printer
from parmepy.operator.monitors.energy_enstrophy import Energy_enstrophy
from dataTG import dim, nb, NBGHOSTS, ADVECTION_METHOD, VISCOSITY, \
    WITH_PROJ, OUTPUT_FREQ, FILENAME, simu


## ----------- A 3d problem -----------
print " ========= Start Navier-Stokes 3D (Vortex Ring benchmark) ========="

## pi constant
pi = m.pi

## Domain
box = pp.Box(dim, length=[2.0 * pi, 2.0 * pi, 2.0 * pi])

## Global resolution
nbElem = [nb] * dim


## Function to compute TG velocity
def computeVel(x, y, z):
    vx = 0.
    vy = 0.
    vz = 1.
    return vx, vy, vz


## Function to compute reference vorticity
def computeVort(x, y, z):
    xc = 6. / 2.
    yc = 6. / 2.
    zc = 6. / 6.
    R = 1.5
    sigma = R / 3.
    Gamma = 0.75
    dist = m.sqrt((x - xc) ** 2 + (y - yc) ** 2)
    s2 = (z - zc) ** 2 + (dist - R) ** 2
    wx = 0.
    wy = 0.
    wz = 0.
    if (dist != 0.):
        cosTheta = (x - xc) / dist
        sinTheta = (y - yc) / dist
        wTheta = Gamma / (pi * sigma ** 2) * m.exp(-(s2 / sigma ** 2))
        wx = - wTheta * sinTheta
        wy = wTheta * cosTheta
        wz = 0.
    return wx, wy, wz

## Fields
velo = pp.Field(domain=box, formula=computeVel,
                name='Velocity', isVector=True)
vorti = pp.Field(domain=box, formula=computeVort,
                 name='Vorticity', isVector=True)

## Usual Cartesian topology definition
ghosts = np.ones((box.dimension)) * NBGHOSTS
topo = Cartesian(box, box.dimension, nbElem,
                 ghosts=ghosts)

## Operators
advec = Advection(velo, vorti,
                  resolutions={velo: nbElem,
                               vorti: nbElem},
                  method=ADVECTION_METHOD
                  )

stretch = Stretching(velo, vorti,
                     resolutions={velo: nbElem,
                                  vorti: nbElem},
                     topo=topo
                     )

diffusion = Diffusion(vorti,
                      resolution=nbElem,
                      viscosity=1.0e-6
                     )

poisson = Poisson(velo, vorti,
                  resolutions={velo: nbElem,
                               vorti: nbElem},
                  projection=WITH_PROJ)

## Diagnostics related to the problem

energy = Energy_enstrophy(velo, vorti,
                          topo=topo,
                          viscosity=VISCOSITY,
                          frequency=OUTPUT_FREQ,
                          prefix=FILENAME)

printer = Printer(fields=[vorti, velo],
                  frequency=100,
                  prefix='./res/vort_ring_',
                  ext='.vtk')

distrAdvStr = Redistribute([vorti, velo], advec, stretch)
#distrStrPoiss = Redistribute([vorti, velo], stretch, poisson)

distrStrAdv = Redistribute([vorti, velo], stretch, advec)

## Define the problem to solve

pb = NSProblem(operators=[advec, distrAdvStr, stretch, distrStrPoiss,
                          diffusion, poisson],
               simulation=simu, monitors=[printer], dumpFreq=-1)

## Setting solver to Problem
pb.setUp()

print 'all topologies', box.topologies


## Solve problem
#poisson.apply(simu)
pb.solve()

## end of time loop ##

# Clean memory buffers
fftw2py.clean_fftw_solver(box.dimension)