#!/usr/bin/python
import parmepy as pp
from parmepy.operator.diffusion import Diffusion
from parmepy.operator.poisson import Poisson
from parmepy.problem.navier_stokes import NSProblem
from parmepy.operator.monitors.energy_enstrophy import Energy_enstrophy
import numpy as np
import math
pi = math.pi
sin = np.sin
cos = np.cos

## Physical Domain description
dim = 3
LL = 2 * pi * np.ones((dim))
dom = pp.Box(dimension=dim, length=LL)
resol = [65, 65, 65]


## Fields
def computeVel(x, y, z):
#    # --------Taylor Green------------
    vx = np.sin(x) * np.cos(y) * np.cos(z)
    vy = - np.cos(x) * np.sin(y) * np.cos(z)
    vz = 0.
    return vx, vy, vz

velocity = pp.Field(domain=dom, formula=computeVel,
                    isVector=True, name='Velocity')

# formula to compute initial vorticity field
coeff = 4 * pi ** 2 * (LL[1] ** 2 * LL[2] ** 2 + LL[0] ** 2 * LL[2] ** 2 +
                       LL[0] ** 2 * LL[1] ** 2) / (LL[0] ** 2 * LL[1] ** 2
                                                   * LL[2] ** 2)

cc = 2 * pi / LL


def computeVort(x, y, z):
    # --------Taylor Green------------
    wx = - np.cos(x) * np.sin(y) * np.sin(z)
    wy = - np.sin(x) * np.cos(y) * np.sin(z)
    wz = 2. * np.sin(x) * np.sin(y) * np.cos(z)
    return wx, wy, wz

vorticity = pp.Field(domain=dom, formula=computeVort,
                     name='Vorticity', isVector=True)


## Definition of the Diffusion and Poisson operators
nu = 1e-3
diffusion = Diffusion(vorticity, resolution=resol, viscosity=nu)

poisson = Poisson(velocity, vorticity, resolutions={velocity: resol,
                                                    vorticity: resol})

output = 'energy.dat'

timeStep = 1e-2

monitor = Energy_enstrophy(velocity, vorticity,
                           resolutions={velocity: resol,
                                        vorticity: resol},
                           viscosity=nu,
                           frequency=1,
                           outputPrefix=output)

diffusion.setUp()
poisson.setUp()
monitor.setUp()

t = 0.0
niter = 1
vorticity.initialize()
print "pre ", vorticity.norm()

while t <= 1.0:
    diffusion.apply(timeStep)
    poisson.apply()
    monitor.apply(t, timeStep, niter)
    t += timeStep

print "post ", vorticity.norm()