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

# ENCUT=600 eV, 2x2 supercell, kpoints: 4	6	8	10	11	13	15
Eb2x2_600 = np.array(
[[4,	108.109,	100.950,	109.527,	103.931,	104.410,	106.475,	105.0],
[5,	103.771,	108.278,	104.209,	101.846,	103.879,	106.108,	106.6],
[6,	112.598,	100.335,	105.989,	105.941,	103.982,	103.480,	105.0],
[7,	91.256,	111.595,	102.621,	103.364,	104.406,	104.561,	106.2],
[8,	111.170,	97.074,	104.291,	105.777,	104.870,	101.985,	103.0],
[9,	115.367,	106.027,	105.748,	101.290,	102.757,	105.709,	105.1],
[10,	105.063,	104.580,	99.650,	107.108,	104.850,	103.789,	103.1]]
)

# ENCUT=600 eV, 3x3 supercell, kpoints: 4	6	8	10
Eb3x3_600 = np.array(
[[4,	93.590,	100.769,	100.586,	98.810],
[5,	99.408,	96.273,	100.132,	100.402],
[6,	97.344,	103.577,	96.938,	99.862],
[7,	102.998,	93.957,	97.382,	98.965],
[8,	91.352,	98.926,	96.967,	97.083],
[9,	98.260,	97.932,	98.810,	98.309],
[10,	98.550,	98.376,	97.566,	96.890]]
)

# ENCUT=550 eV, 4x4 supercell, kpoints: 6
Eb4x4_550 = np.array(
[[4,	102.245],
[5,	99.331],
[6,	96.099],
[7,	98.280],
[8,	97.276],
[9,	99.457]]
)

mpl.rc("text", usetex=True)
fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(4.,4.))

plt.subplot(2,1,1)
plt.plot( Eb2x2_600[:,0], Eb2x2_600[:,1:], marker='o' )
y_avg = (Eb3x3_600[-1,-1] + Eb3x3_600[-2,-1]) / 2.
plt.plot( [0,13], [y_avg, y_avg], linestyle='--', c='black' )
plt.fill_between([0,13], [y_avg+4.2, y_avg+4.2], [y_avg-4.2, y_avg-4.2], color='grey', alpha=0.5)
plt.legend(['4', '6', '8', '10', '11', '13', '15'])
plt.xlim(3.5,12.5)
plt.ylim(90.,115.)
plt.tick_params(labelbottom=False, length=6, width=1, direction='in', top=True, right=True)
plt.ylabel('Barrier (kJ/mol)')
plt.annotate('2x2', xy=(8.5, 92.), size=18)

plt.subplot(2,1,2)
plt.plot( Eb3x3_600[:,0], Eb3x3_600[:,1:], marker='o' )
plt.plot( Eb4x4_550[:,0], Eb4x4_550[:,1:], marker='o' )
#plt.plot( Eb3_400_x, Eb3_400, marker='o' )
#y_avg = (Eb4_400[5][2] + Eb4_400[6][2]) / 2
plt.plot( [0,13], [y_avg, y_avg], linestyle='--', c='black' )
plt.fill_between([0,13], [y_avg+4.2, y_avg+4.2], [y_avg-4.2, y_avg-4.2], color='grey', alpha=0.5)
plt.legend(['4', '6', '8', '10', '6 (4x4)'])
plt.xlim(3.5,12.5)
plt.ylim(90.,104.)
plt.annotate('3x3', xy=(8.5, 101.), size=18)
plt.xlabel('Number of layers')
plt.ylabel('Barrier (kJ/mol)')
plt.tick_params(length=12, width=1, direction='in', top=True, right=True)

plt.tight_layout()
plt.subplots_adjust(wspace=0, hspace=0)
plt.savefig('convergence.pdf')
#plt.show()