#!/usr/bin/env python
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from math import erf

# Q, Eb, Easymp, Zb
nn_backup_oldprobablywrongsurface = np.array(
[[ -0.3 ,  -20.64450740269419 ,  -22.585227807462356 ,  2.0772631480363986 ],
[ -0.25 ,  -20.8490398809165 ,  -22.747354605938817 ,  2.1019185995794176 ],
[ -0.2 ,  -21.016994269823005 ,  -22.874819540903 ,  2.12797285831064 ],
[ -0.15000000000000002 ,  -21.15329038800612 ,  -22.97303386857885 ,  2.156103397500163 ],
[ -0.10000000000000003 ,  -21.26242083529775 ,  -23.04671223370818 ,  2.1859713774880625 ],
[ -0.050000000000000044 ,  -21.348469768625634 ,  -23.099890629065484 ,  2.217524141842634 ],
[ -5.551115123125783e-17 ,  -21.415028546716695 ,  -23.135996872947704 ,  2.2506893154347023 ],
[ 0.04999999999999993 ,  -21.465285799784755 ,  -23.157947209594596 ,  2.285150234695392 ],
[ 0.09999999999999992 ,  -21.50199151050099 ,  -23.168213249198363 ,  2.320026524079246 ],
[ 0.1499999999999999 ,  -21.52742871415249 ,  -23.168872308967767 ,  2.35500514897085 ],
[ 0.1999999999999999 ,  -21.54359608704241 ,  -23.161675169479334 ,  2.3908512417735923 ],
[ 0.24999999999999983 ,  -21.55220713014844 ,  -23.148090701350434 ,  2.4264677827225367 ],
[ 0.2999999999999999 ,  -21.554665134647625 ,  -23.12933933524405 ,  2.461976609975075 ]]
)

nn = np.array(
[[ -0.3 ,  -26.623967299253614 ,  -28.57137163676987 ,  2.075312873477223 ],
[ -0.25 ,  -26.721685563362314 ,  -28.634300143852002 ,  2.1005966420316886 ],
[ -0.2 ,  -26.801035325565937 ,  -28.68030534574057 ,  2.12800566803265 ],
[ -0.15000000000000002 ,  -26.86425118488444 ,  -28.711722523583823 ,  2.157516436778961 ],
[ -0.10000000000000003 ,  -26.913423247878022 ,  -28.730630355159825 ,  2.1882940896434646 ],
[ -0.050000000000000044 ,  -26.95049903344413 ,  -28.73886833016352 ,  2.2206459710020336 ],
[ -5.551115123125783e-17 ,  -26.97723285963948 ,  -28.738046002077997 ,  2.2539228958515714 ],
[ 0.04999999999999993 ,  -26.995201354170067 ,  -28.72962190119392 ,  2.287770244202866 ],
[ 0.09999999999999992 ,  -27.005839101306865 ,  -28.714888296233397 ,  2.3221186691402997 ],
[ 0.1499999999999999 ,  -27.01041496797999 ,  -28.69496738475587 ,  2.3563554457750695 ],
[ 0.1999999999999999 ,  -27.0098870670912 ,  -28.670824898851205 ,  2.390354758557775 ],
[ 0.24999999999999983 ,  -27.005264941061842 ,  -28.643280989694045 ,  2.425256718024927 ],
[ 0.2999999999999999 ,  -26.997469967422553 ,  -28.61302165632594 ,  2.4601511084493244 ],]
)

nn[:,1] *= 96.4853075
nn[:,2] *= 96.4853075

vasp_backup_oldprobablywrongsurface = np.array(
[[ -0.30, -.20641083E+02, -.22584612E+02, 2.076469858],
[ -0.25, -.20845835E+02, -.22750576E+02, 2.101397463],
[ -0.20, -.21013513E+02, -.22880925E+02, 2.128479673],
[ -0.15, -.21149507E+02, -.22981168E+02, 2.158491493],
[ -0.10, -.21258447E+02, -.23056380E+02, 2.18884129],
[ -0.05, -.21344766E+02, -.23110684E+02, 2.220005943],
[ 0.00, -21.411814, -23.147636, 2.251306189],
[ 0.05, -.21462966E+02, -.23170213E+02, 2.285490309],
[ 0.10, -.21500498E+02, -.23181089E+02, 2.318608252],
[ 0.15, -.21526657E+02, -.23182200E+02, 2.352746481],
[ 0.20, -.21543072E+02, -.23175433E+02, 2.38744197],
[ 0.25, -.21551533E+02, -.23162163E+02, 2.422236742],
[ 0.30, -.21553204E+02, -.23143500E+02, 2.45748053]]
)

vasp = np.array(
[[ -0.30, -.26616153E+02, -.28513095E+02, 2.07677243548476],
[ -0.25, -.26715026E+02, -.28574734E+02, 2.10267164940857],
[ -0.20, -.26795497E+02, -.28619779E+02, 2.12994666773284],
[ -0.15, -.26859981E+02, -.28650270E+02, 2.1585605258384],
[ -0.10, -.26910239E+02, -.28668646E+02, 2.18842768872838],
[ -0.05, -.26948370E+02, -.28676732E+02, 2.21865932589169],
[ 0.00, -.26976081E+02, -.28676030E+02, 2.24949173348615],
[ 0.05, -.26994583E+02, -.28668196E+02, 2.28261843221517],
[ 0.10, -.27005386E+02, -.28654340E+02, 2.31658949277447],
[ 0.15, -.27009654E+02, -.28635534E+02, 2.34991605948969],
[ 0.20, -.27008018E+02, -.28612775E+02, 2.38416614367356],
[ 0.25, -.27001699E+02, -.28586662E+02, 2.41985251022372],
[ 0.30, -.26991333E+02, -.28557913E+02, 2.45234556888136]]
)

vasp[:,1] *= 96.4853075
vasp[:,2] *= 96.4853075

mpl.rc("font", size=12)
fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(3.69,3))

ax2 = plt.subplot(1,2,1)

#plt.plot( nn[:,0], nn[:,1] - nn[:,2] +0.03322326412951426*96.4853075, label='NN', color='blue', marker='o' )
plt.plot( nn[:,0], nn[:,1] - nn[:,2], label='NN', color='blue', marker='o' )
plt.plot( vasp[:,0], vasp[:,1] - vasp[:,2], label='DFT', color='red', marker='+', linestyle='--' )
#plt.errorbar( nn[:,0], nn[:,1] - nn[:,2], yerr=len(nn)*[4.2], label='NN', color='blue', marker='o' )

plt.legend(loc='upper right', numpoints=1, frameon=False, borderaxespad=0)
plt.xlim(-0.35,0.35)
plt.ylim(150., 200.)
plt.tick_params(length=6, width=1)
plt.tick_params(axis='x', labelsize=9)
ax2.set_xticks([-0.3, -0.1, 0.1, 0.3])
plt.ylabel('$E_b$ (kJ/mol)')
plt.xlabel(r'Q ($\mathrm{\AA}$)')
#plt.annotate('(a)', xy=(0.15, 175.), size=12)

ax3 = plt.subplot(1,2,2)

plt.plot( nn[:,0], nn[:,3], label='NN', color='blue', marker='o' )
plt.plot( vasp[:,0], vasp[:,3], label='DFT', color='red', marker='+', linestyle='--' )

plt.xlim(-0.35,0.35)
plt.ylim(2.0,2.5)
plt.tick_params(length=6, width=1)
plt.tick_params(axis='y', labelleft=False, labelright=True)
plt.tick_params(axis='x', labelsize=9)
plt.ylabel(r'$Z_b$ ($\mathrm{\AA}$)')
ax3.yaxis.set_label_position("right")
ax3.set_xticks([-0.3, -0.1, 0.1, 0.3])
plt.xlabel(r'Q ($\mathrm{\AA}$)')
#plt.annotate('(b)', xy=(-0.25, 2.28), size=12)

plt.tight_layout()
plt.subplots_adjust(hspace = 0.0, wspace = 0.08)
plt.subplots_adjust(left=0.18, right=0.84)
plt.savefig('coupling.pdf')
#plt.show()