# -*- coding: utf-8 -*-
"""
Created on Fri May 23 10:19:28 2014

@author: Jens von der Linden
"""

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

import numpy as np
import sympy as sp
import scipy.interpolate as interp


class Equil_solver(object):

    def ___init__(self, a, points):
        self.r = np.linspace(0, a, points)

    def set_splines(self, param_points):
        splines = {}
        r = self.r
        d_b_theta_r_request = False
        for key, value in param_points.items():
            if key == 'd_b_theta_over_r':
                d_b_theta_r_request = True
            else:
                splines[key] = interp.InterpolatedUnivariateSpline(r, value,
                                                                   k=3)
        if d_b_theta_r_request:
            splines['d_b_theta_over_r'] = self.d_b_theta_over_r(r,
                                                           splines['b_theta'])
        self.splines = splines

    def get_splines(self):
        return self.splines


class Parabolic_nu2(Equil_solver):

    def __init__(self, a=1, points=500, q0=1.0, k=1, b_z0=1, temp=1.0):
        self.r = np.linspace(0, a, points)
        self.q0 = q0
        self.k = k
        self.b_z0 = b_z0
        self.temp = temp
        r = self.r
        self.j0 = self.get_j0()
        param_points = {'j_z': self.j_z(r), 'b_theta': self.b_theta(r),
                        'b_z': self.b_z(r), 'p_prime': self.pprime(r),
                        'pressure': self.pressure(r), 'q': self.q(r),
                        'rho': self.rho(r),
                        'd_b_theta_over_r': None}
        self.set_splines(param_points)

    def get_j0(self):
        self.j0 = 1
        return self.k*2*self.b_z(0)/self.q0

    def j_z(self, r):
        return self.j0*(1 - r**2)**2

    def b_theta(self, r):
        j0 = self.j0
        return j0*r/2 - j0*r**3/2 + j0*r**5/6

    def b_z(self, r):
        b_z0 = self.b_z0
        if isinstance(r, float) or isinstance(r, int):
            return np.ones(1) * b_z0
        else:
            return np.ones(r.size) * b_z0

    def pprime(self, r):
        j0 = self.j0
        return (-j0**2*r/2 + 3*j0**2*r**3/2 - 5*j0**2*r**5/3 + 5*j0**2*r**7/6
                - j0**2*r**9/6)

    def pressure(self, r):
        j0 = self.j0
        return (47*j0**2/720 - j0**2*r**2/4 + 3*j0**2*r**4/8 - 5*j0**2*r**6/18
                + 5*j0**2*r**8/48 - j0**2*r**10/60)

    def rho(self, r):
        return np.ones(r.size)
        #return self.pressure(r)/self.temp

    def q(self, r):
        q_to_return = r*self.k*self.b_z(r)/self.b_theta(r)
        q_to_return[0] = self.q0
        return q_to_return

    def d_b_theta_over_r(self, r, b_theta):
        b_over_r = interp.InterpolatedUnivariateSpline(r, b_theta.derivative()(r)
                                                       /r, k=3)
        return b_over_r.derivative()


class Newcomb_constant_pressure(Equil_solver):
    """
    Creates splines describing the constant pressure profile at the end of
    Newcomb's 1960 paper.
    """

    def __init__(self, a=0.1, r_0i=0.5, k=1, b_z0=0.1, b_thetai=0.1,
                 points=500):
        self.r = np.linspace(a, r_0i, points)
        self.k = k
        self.b_z0 = b_z0
        self.b_thetai = b_thetai
        r = self.r
        param_points = {'j_z': self.j_z(r), 'b_theta': self.b_theta(r),
                        'b_z': self.b_z(r), 'p_prime': self.pprime(r),
                        'pressure': self.pressure(r), 'q': self.q(r),
                        'rho': self.rho(r)}
        self.set_splines(param_points)

    def get_j_z(self, r):
        return np.zeros(r.size)

    def b_theta(self, r):
        return self.b_thetai*self.r_0i/r

    def b_z(self, r):
        np.ones(r.size)*self.b_z0

    def pprime(self, r):
        np.ones(r.size)

    def rho(self, r):
        np.ones(r.size)

    def q(self, r):
        q_to_return = r*self.k*self.b_z(r)/self.b_theta(r)
        return q_to_return


class Smoothed_core_skin(Equil_solver):
    pass


class Sharp_core_skin(Equil_solver):
    pass
@@ -544,26 +544,25 @@
 %0Aclass Equil
-_s
+S
 olver(object
@@ -1348,18 +1348,17 @@
 olic
-_n
+N
 u2(Equil
 _sol
@@ -1349,26 +1349,25 @@
 licNu2(Equil
-_s
+S
 olver):%0A%0A   
@@ -3223,19 +3223,17 @@
 comb
-_c
+C
 onstant
-_p
+P
 ress
@@ -3237,26 +3237,25 @@
 essure(Equil
-_s
+S
 olver):%0A    

Change class names to CapWord style.