#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import numpy as np
import scipy.special
import os

import plot
import plot_params
import constants as cst
import fleches



script_path = os.path.dirname(os.path.realpath(__file__))
# File generated by libpotentials.git/testing
saved_file = os.path.join(script_path, "potential_HermanSkillman.csv")

hs_data = np.loadtxt(saved_file, delimiter=',', skiprows=0, dtype=float)

Xe_Z0_Ip = 12.265625 * cst.eV_to_Eh

Umin = -5.2 * Xe_Z0_Ip
Umax = 0.25

GammaE = 1.25 * Xe_Z0_Ip

r = np.linspace(-5.0, 5.0, 1000)
cs = 1.0
U = -cs / abs(r)
U[U<Umin] = np.NaN

lr = r[-1] - r[0]
lU = Umax - Umin

r0 = abs(cs/Xe_Z0_Ip)

fig = plot.figure(figsize = (10.0, 7.0))

ax1, ax1b = plot.get_axis_two_scales(fig,
                                     scale_y = cst.Eh_to_eV,
                                     ax2_ylabel = 'Energy [eV]')

# Ip
#ax1.plot([-r0, r0], [-Xe_Z0_Ip, -Xe_Z0_Ip], '-m')

# U(r) (Coulomb)
ax1.plot(r, U, '--k', label = 'Coulomb')

# U(r) (Herman-Skillman)
ax1.plot(+hs_data[:,0], hs_data[:,1], '-b', label = 'Atomic')
ax1.plot(-hs_data[:,0], hs_data[:,1], '-b')

# Continuum
ax1.fill([r[0],r[-1],r[-1],r[0]],[0.0,0.0,Umax,Umax], fill=False, hatch='/', alpha = 0.5)



leg = ax1.legend(loc = "best")
leg.get_frame().set_alpha(0.75)

ax1.set_xlim((r[0], r[-1]))
ax1.set_ylim((0.98*Umin, Umax))

ax1.set_xlabel('r [bohr]')
ax1.set_ylabel('Energy [Hartree]')

#for ext in ['pdf', 'svg']:
#    plot.savefig('heating_atomic_potential.' + ext)
plot.show()

