#! python
# -*- coding: utf-8 -*-
"""
Created on Wed Jan 25 13:30:52 2017

@author: Adam Deller, UNIVERSITY COLLEGE LONDON.

Tools for calculating the Stark effect in Rydberg helium
using the Numerov method.

Based on:

    Stark structure of the Rydberg states of alkali-metal atoms
        M. L. Zimmerman et al. Phys. Rev. A, 20 2251 (1979)
        http://dx.doi.org/10.1103/PhysRevA.20.2251

    Rydberg atom diamagnetism
        M. M. KASH, PhD thesis, Massachusetts Institute of Technology.
        Dept. of Physics. (1988)
        http://hdl.handle.net/1721.1/14367

    Cold atoms and molecules by Zeeman deceleration and Rydberg-
    Stark deceleration
        S. D. Hogan, Habilitationsschrift, ETH Zurich (2012)
        http://dx.doi.org/10.3929/ethz-a-007577485

"""
from __future__ import print_function, division
from math import ceil, log, exp
import numpy as np
from numba import jit
from tqdm import trange
from .drake1999 import quantum_defects
import errno
import os
from datetime import datetime
from scipy.constants import h, hbar, c, alpha, m_e, e, epsilon_0, atomic_mass, pi, Rydberg, physical_constants
a_0 = physical_constants['Bohr radius'][0]
E_h = physical_constants['Hartree energy'][0]

# derived constants
mass_He = 4.002602 * atomic_mass
Z = 2
mu_me = (mass_He - m_e) / mass_He
mu_M = m_e / mass_He
En_h_He = E_h * mu_me
a_0_He = a_0 / mu_me

def version_table():
    import matplotlib
    import scipy
    import numba
    import numpy as np
    import pandas as pd
    
    versions = [np.__version__, scipy.__version__, matplotlib.__version__, pd.__version__, numba.__version__]
    indicies = ['Numpy', 'Scipy', 'Matplotlib', 'Pandas', 'Numba']
    return pd.DataFrame(versions, index=indicies, columns=['Version'])

@jit
def get_nl_vals(nmin, nmax, m):
    """ n and l vals for each matrix column, using range n_min to n_max.
    """
    n_rng = np.arange(nmin, nmax + 1)
    n_vals = np.array([], dtype='int32')
    l_vals = np.array([], dtype='int32')
    for n in n_rng:
        l_rng = np.arange(np.abs(m), n)
        n_vals = np.append(n_vals, np.array(np.zeros_like(l_rng) + n))
        l_vals = np.append(l_vals, l_rng)
    return n_vals, l_vals

@jit
def get_nlm_vals(nmin, nmax):
    """ n, l and m vals for each matrix column, using range n_min to n_max.
    """
    n_rng = np.arange(nmin, nmax + 1)
    n_vals = np.array([], dtype='int32')
    l_vals = np.array([], dtype='int32')
    m_vals = np.array([], dtype='int32')
    for n in n_rng:
        l_rng = np.arange(0, n)
        for l in l_rng:
            m_rng = np.arange(-l, l+1)
            n_vals = np.append(n_vals, np.array(np.zeros_like(m_rng) + n))
            l_vals = np.append(l_vals, np.array(np.zeros_like(m_rng) + l))
            m_vals = np.append(m_vals, m_rng)
    return n_vals, l_vals, m_vals

@jit
def get_J_vals(S, L_vals, diff):
    """ J = L + diff; unless L == 0, in which case J = S.
    """
    J_vals = L_vals + diff
    J_vals[L_vals == 0] = S
    return J_vals

@jit
def get_triplet_nLJ(nmin, nmax, m):
    """ n and L and J vals for each matrix column, using range n_min to n_max.
    """
    n_rng = np.arange(nmin, nmax + 1, dtype='int32')
    n_vals = np.array([], dtype='int32')
    L_vals = np.array([], dtype='int32')
    J_vals = np.array([], dtype='int32')
    for n in n_rng:
        l_rng = np.arange(m, n, dtype='int32')
        for l in l_rng:
            if l == 0:
                n_vals = np.append(n_vals, n)
                L_vals = np.append(L_vals, 0)
                J_vals = np.append(J_vals, 1)
            else:
                n_vals = np.append(n_vals, np.zeros(3, dtype='int32') + n)
                L_vals = np.append(L_vals, np.zeros(3, dtype='int32') + l)
                J_vals = np.append(J_vals, np.arange(l - 1, l + 2, dtype='int32'))
    return n_vals, L_vals, J_vals

@jit
def get_qd(S, n_vals, L_vals, J_vals):
    """ Calculate quantum defects.
    """
    iterations = int(10)
    num_cols = len(n_vals)
    qd = np.zeros(num_cols)
    for i in range(num_cols):
        n = n_vals[i]
        L = L_vals[i]
        J = J_vals[i]
        if L in quantum_defects[S]:
            if J in quantum_defects[S][L]:
                delta = quantum_defects[S][L][J]
                # calculate quantum defects
                qd_i = delta[0]
                for rep in range(iterations):
                    # repeat to get convergence
                    m = n - qd_i
                    defect = delta[0]
                    for j, d in enumerate(delta[1:]):
                        defect = defect + d*m**(-2.0*(j + 1))
                    qd_i = defect
                    qd[i] = defect
            else:
                qd[i] = np.nan
        else:
            qd[i] = 0.0
    return qd

@jit
def En_0(neff):
    """ Field-free energy. Ignores extra correction terms.

        -- atomic units --
    """
    energy = np.array([])
    for n in neff:
        en = -0.5 * n**-2.0
        energy = np.append(energy, en)
    return energy

@jit
def W_n(S, n_vals, L_vals, J_vals):
    """ Field-free energy. Includes extra correction terms.

        -- atomic units --
    """
    neff = n_vals - get_qd(S, n_vals, L_vals, J_vals)
    energy = np.array([])
    for i, n in enumerate(n_vals):
        en = -0.5 * (neff[i]**-2.0 - 3.0 * alpha**2.0 / (4.0 * n**4.0) + \
                     mu_M**2.0 * ((1.0 + (5.0 / 6.0) * (alpha * Z)**2.0)/ n**2.0))
        energy = np.append(energy, en)
    return energy

@jit
def E_zeeman(m_vals, B_z):
    """ Energy shift due to the interaction of the orbital angular momentum of the Rydberg electron with the magnetic field.

        -- atomic units --
    """
    return m_vals * B_z * (1/2)

def chose_step(nmax, lmax, step_params):
    interp_type = step_params[0]
    if interp_type == 'flat':
        return step_params[1]
    elif interp_type == 'poly':
        return chose_step_poly(nmax, lmax, step_params)
    elif interp_type == 'sigmoid':
        return chose_step_sigmoid(nmax, lmax, step_params)
    
@jit
def chose_step_poly(nmax, lmax, step_params):
    interp_type, step_low, step_high, interp_params = [*step_params]
    exponent, round_step = [*interp_params]
    return np.round( (lmax/nmax)**exponent * (step_high-step_low) + step_low, round_step)

@jit
def chose_step_sigmoid(nmax, lmax, step_params):
    interp_type, step_low, step_high, interp_params = [*step_params]
    width, mid, round_step = [*interp_params]
    return np.round( (( step_high-step_low )/( 1+np.exp( -width*( (lmax/nmax)-mid ) ) )) + step_low, round_step)

@jit
def wf_numerov(n, l, nmax, rmin, step):
    """ Use the Numerov method to find the wavefunction for state n*, l, where
        n* = n - delta.

        nmax ensures that wavefunctions from different values of n can be aligned.
    """
    l = float(l)             # horrible hack. Otherwise jit fails if l = int(>81).
    W1 = -0.5 * n**-2.0
    W2 = (l + 0.5)**2.0
    rmax = 2 * nmax * (nmax + 15)
    r_in = n**2.0 - n * (n**2.0 - l*(l + 1.0))**0.5
    step_sq = step**2.0
    # ensure wf arrays will align using nmax
    if n == nmax:
        i = 0
        r_sub2 = rmax
    else:
        i = int(ceil(log(rmax / (2 * n * (n + 15))) / step))
        r_sub2 = rmax * exp(-i*step)
    i += 1

    # initialise
    r_sub1 = rmax * exp(-i*step)
    rvals = [r_sub2, r_sub1]
    g_sub2 = 2.0 * r_sub2**2.0 * (-1.0 / r_sub2 - W1) + W2
    g_sub1 = 2.0 * r_sub1**2.0 * (-1.0 / r_sub1 - W1) + W2