%matplotlib notebook

from sympy import init_printing
from sympy import S
from sympy import sin, cos, tanh, exp, pi, sqrt

from boutdata.mms import x, y, z, t
from boutdata.mms import DDX, DDZ

import os, sys
# If we add to sys.path, then it must be an absolute path
common_dir = os.path.abspath('./../../../../common/')
# Sys path is a list of system paths
sys.path.append(common_dir)
from CELMAPy.MES import get_metric, make_plot, BOUT_print

init_printing()

def poisson(f, g, metric):
    return\
    DDZ(f, metric=metric)*DDX(g, metric=metric)\
    -\
    DDX(f, metric=metric)*DDZ(g, metric=metric)\

"""
Explanation: Exact solution used in MES runs
We would like to MES the operation
$$
\left{\left(\partial_\rho \phi\right)^2, n\right}
$$
Using cylindrical geometry.
End of explanation
"""


folder = '../mixModeAndGaussian/'
metric = get_metric()

"""
Explanation: Initialize
End of explanation
"""


# Initialization
the_vars = {}

"""
Explanation: Define the variables
End of explanation
"""


# We need Lx
from boututils.options import BOUTOptions
myOpts = BOUTOptions(folder)
Lx = eval(myOpts.geom['Lx'])

# Two gaussians

# The skew sinus
# In cartesian coordinates we would like a sinus with with a wave-vector in the direction
# 45 degrees with respect to the first quadrant. This can be achieved with a wave vector
# k = [1/sqrt(2), 1/sqrt(2)]
# sin((1/sqrt(2))*(x + y))
# We would like 2 nodes, so we may write
# sin((1/sqrt(2))*(x + y)*(2*pi/(2*Lx)))
# Rewriting this to cylindrical coordinates, gives
# sin((1/sqrt(2))*(x*(cos(z)+sin(z)))*(2*pi/(2*Lx)))

# The gaussian
# In cartesian coordinates we would like
# f = exp(-(1/(2*w^2))*((x-x0)^2 + (y-y0)^2))
# In cylindrical coordinates, this translates to
# f = exp(-(1/(2*w^2))*(x^2 + y^2 + x0^2 + y0^2 - 2*(x*x0+y*y0) ))
#   = exp(-(1/(2*w^2))*(rho^2 + rho0^2 - 2*rho*rho0*(cos(theta)*cos(theta0)+sin(theta)*sin(theta0)) ))
#   = exp(-(1/(2*w^2))*(rho^2 + rho0^2 - 2*rho*rho0*(cos(theta - theta0)) ))

# A parabola
# In cartesian coordinates, we have
# ((x-x0)/Lx)^2
# Chosing this function to have a zero value at the edge yields in cylindrical coordinates
# ((x*cos(z)+Lx)/(2*Lx))^2

# Scaling with 40 to get S in order of unity

w = 0.8*Lx
rho0 = 0.3*Lx
theta0 = 5*pi/4
the_vars['n'] = 40*sin((1/sqrt(2))*(x*(cos(z)+sin(z)))*(2*pi/(2*Lx)))*\
                  exp(-(1/(2*w**2))*(x**2 + rho0**2 - 2*x*rho0*(cos(z - theta0)) ))*\
                  ((x*cos(z)+Lx)/(2*Lx))**2

# Mixmode

# Need the x^3 in order to let the second derivative of the field go towards one value when rho -> 0
#    (needed in Arakawa brackets)
# Mutliply with a mix of modes
# Multiply with a tanh in order to make the variation in x more homogeneous

# Scaling with 10 to make variations in phi comparable to those of n

the_vars['phi'] = 10*(6+((x/(Lx))**3)*\
                cos(2*z)*\
                (
                   cos(2*pi*(x/Lx)) + sin(2*pi*(x/Lx))
                 + cos(3*2*pi*(x/Lx)) + cos(2*2*pi*(x/Lx))                
                )\
                *(1/2)*(1-tanh((1/8)*(x))))

"""
Explanation: Define the function to take the derivative of
NOTE:

z must be periodic
The field $f(\rho, \theta)$ must be of class infinity in $z=0$ and $z=2\pi$
The field $f(\rho, \theta)$ must be single valued when $\rho\to0$
The field $f(\rho, \theta)$ must be continuous in the $\rho$ direction with $f(\rho, \theta + \pi)$
Eventual BC in $\rho$ must be satisfied
End of explanation
"""


the_vars['S'] = poisson(
                        DDX(the_vars['phi'], metric=metric)**2.0,
                        the_vars['n'],
                        metric=metric
                       )

"""
Explanation: Calculating the solution
End of explanation
"""


make_plot(folder=folder, the_vars=the_vars, plot2d=True, include_aux=False, save=False)

"""
Explanation: Plot
End of explanation
"""


BOUT_print(the_vars, rational=False)

"""
Explanation: Print the variables in BOUT++ format
End of explanation
"""


