# -*- coding: utf-8 -*-
"""
Created on Wed Aug 06 13:43:55 2014

@author: Jens von der Linden
"""

from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
from future.builtins import (ascii, bytes, chr, dict, filter, hex, input,
                             int, map, next, oct, open, pow, range, round,
                             str, super, zip)
"""Python 3.x compatibility"""

import eigenvalue_goedbloed as eg
from numba import jit, float64


def newcomb_f_16(r, k, m, b_z, b_theta, **kwargs):
    r"""
    Return f from Newcomb's paper.

    Parameters
    ----------
    r : ndarray of floats
        radius
    k : float
        axial periodicity number
    m : float
        azimuthal periodicity number
    b_z : ndarray of floats
        axial magnetic field
    b_theta : ndarray of floats
        azimuthal mangetic field

    Returns
    -------
    f : ndarray of floats
        f from Newcomb's paper

    Notes
    -----
    Implements equation:
    :math:`f = \frac{r (k r B_{z}+ m B_{\theta})^2}{k^{2} r^{2} + m^{2}}`


    Reference
    ---------
    Newcomb (1960) Hydromagnetic Stability of a Diffuse Linear Pinch
    Equation (16)
    """
    params_num = {'r': r, 'k': k, 'm': m, 'b_z': b_z, 'b_theta': b_theta}
    params_denom = {'r': r, 'k': k, 'm': m}
    #if type(r) == float:
    #    all_f_g.all_f.append([r, r*f_num_wo_r(**params_num)/f_denom(**params_denom)])
    return r*f_num_wo_r(**params_num)/f_denom(**params_denom)


def goedbloed_f_9_106(r, k, m, b_z, b_theta):
    r"""
    Return f from Newcomb's paper.

    Parameters
    ----------
    r : ndarray of floats
        radius
    k : float
        axial periodicity number
    m : float
        azimuthal periodicity number
    b_z : ndarray of floats
        axial magnetic field
    b_theta : ndarray of floats
        azimuthal mangetic field
    Returns
    -------
    f : ndarray of floats
        f from Newcomb's paper

    Reference
    ---------
    Goedbloed (2004) Principles of Magnetohydrodynamics
    Equation (9.106)
    """
    params = {'r': r, 'k': k, 'm': m, 'b_z': b_z, 'b_theta': b_theta}
    f = eg.f(**params)
    return r**3*f**2/(m**2 + k**2*r**2)


def jardin_f_8_78(r, k, m, b_z, b_theta, q):
    r"""
    Return f from Newcomb's paper.

    Parameters
    ----------
    r : ndarray of floats
        radius
    k : float
        axial periodicity number
    m : float
        azimuthal periodicity number
    b_z : ndarray of floats
        axial magnetic field
    b_theta : ndarray of floats
        azimuthal mangetic field
    q: ndarray of floats
        safety factor

    Returns
    -------
    f : ndarray of floats
        f from Newcomb's paper

    Reference
    ---------
    Jardin (2010) Computational Methods in Plasma Physics
    Equation (8.78)
    """
    return r*b_theta**2*(m - k*q)**2/(k**2*r**2 + m**2)


#@jit(float64(float64, float64, float64))
@jit
def f_denom(r, k, m):
    r"""
    Return denominator of f from Newcomb's paper.

    Parameters
    ----------
    r: ndarray of floats
        radius
    k: float
        axial periodicity number
    m: float
        azimuthal periodicity number

    Returns
    -------
    f_denom: ndarray of floats
         denominator of f from Newcomb's paper

    Notes
    -----
    f denominator :math:`k^{2}r^{2}+m^{2}`

    Reference
    ---------
    Newcomb (1960) Hydromagnetic Stability of a Diffuse Linear Pinch
    Equation (16)
    """
    return k**2*r**2 + m**2


#@jit(float64(float64, float64, float64, float64, float64))
@jit
def f_num_wo_r(r, k, m, b_z, b_theta):
    r"""
    Return numerator of f without r from Newcomb's paper.

    Parameters
    ----------
    r : ndarray of floats
        radius
    k : float
        axial periodicity number
    m : float
        azimuthal periodicity number
    b_z : ndarray of floats
        axial magnetic field
    b_theta : ndarray of floats
        azimuthal mangetic field

    Returns
    -------
    f_num_wo_r : ndarray of floats
        numerator of f without r from Newcomb's paper

    Notes
    -----
    f numerator without r: :math:`(k r B_{z}+m B_{\theta})^{2}`

    Reference
    ---------
    Newcomb (1960) Hydromagnetic Stability of a Diffuse Linear Pinch
    Equation (16)
    """
    return (k*r*b_z + m*b_theta)**2


def f_prime(r, k, m, b_z, b_theta, b_theta_prime, **kwargs):
    r"""
    return f_prime
    """
    term1 = -2.*k**2*r**2/(k**2*r**2 + m**2)**2 * (b_z*k*r + m*b_theta)**2
    term2 = r/(k**2*r**2 + m**2)*(2.*b_z*k + 2.*m*b_theta_prime)*(b_z*k*r + m*b_theta)
    term3 = 1./(k**2*r**2 + m**2)*(b_z*k*r + m*b_theta)**2
    return term1 + term2 + term3
@@ -432,42 +432,8 @@
 %22%22%0A%0A
-import eigenvalue_goedbloed as eg%0A
 from
@@ -1474,720 +1474,8 @@
 )%0A%0A%0A
-def goedbloed_f_9_106(r, k, m, b_z, b_theta):%0A    r%22%22%22%0A    Return f from Newcomb's paper.%0A%0A    Parameters%0A    ----------%0A    r : ndarray of floats%0A        radius%0A    k : float%0A        axial periodicity number%0A    m : float%0A        azimuthal periodicity number%0A    b_z : ndarray of floats%0A        axial magnetic field%0A    b_theta : ndarray of floats%0A        azimuthal mangetic field%0A    Returns%0A    -------%0A    f : ndarray of floats%0A        f from Newcomb's paper%0A%0A    Reference%0A    ---------%0A    Goedbloed (2004) Principles of Magnetohydrodynamics%0A    Equation (9.106)%0A    %22%22%22%0A    params = %7B'r': r, 'k': k, 'm': m, 'b_z': b_z, 'b_theta': b_theta%7D%0A    f = eg.f(**params)%0A    return r**3*f**2/(m**2 + k**2*r**2)%0A%0A%0A
 def 
@@ -3939,12 +3939,13 @@
 erm2 + term3
+%0A

Remove dependence on eigenvalue module.