%load_ext aiida

from aiida import load_profile
import aiida.orm as orm
import sisl
import numpy as np
import matplotlib.pyplot as plt

profile = 'BdG'   # Modify to match name of profile archive was imported into
load_profile(profile)

(oldwidth, oldheight) = plt.rcParams['figure.figsize']

plt.rcParams['figure.figsize'] = (2*oldwidth, 2*oldheight)
plt.rcParams.update({'font.size': 32})

plt.style.use('seaborn-v0_8-deep')

axes_linewidth = 3
plot_linewidth = 4

bulk_Nb_fixeddelta_pdos_kgrid_smearing1 = orm.load_node(uuid='8c083b71-f2ea-4f8f-8c83-4e141a03074e')
bulk_Nb_fixeddelta_pdos_adaptive1 = orm.load_node(uuid='9d6b948a-fe21-4d40-92b3-75b03a1eb54e')

def get_pdos_nk(node):
    return int(node.inputs.parameters['%block pdoskgridmonkhorstpack'].split()[0])

nk_dict = {}
for node in bulk_Nb_fixeddelta_pdos_kgrid_smearing1.called:
    nk_dict[get_pdos_nk(node)] = node

with nk_dict[800].outputs.retrieved.open('aiida.PDOS.xml') as fh:
    geom, E800, PDOS800 = sisl.get_sile_class('PDOS.xml')(fh).read_data()

with bulk_Nb_fixeddelta_pdos_adaptive1.outputs.retrieved.open('aiida.PDOS.xml') as fh:
    geom, E_adaptive, PDOS_adaptive = sisl.get_sile_class('PDOS.xml')(fh).read_data()

plt.plot(1000*E800, np.sum(PDOS800[0,:,:], axis=0), 'b-', label=r'uniform', linewidth=plot_linewidth-1)
plt.plot(1000*E_adaptive, np.sum(PDOS_adaptive[0,:,:], axis=0), 'r--', label=r'adaptive', linewidth=plot_linewidth)
#plt.legend(bbox_to_anchor=(1,1))
plt.legend()
plt.xlabel(r'$E - E_F$ [meV]')
plt.ylabel(r'DOS [1/eV]')
ax = plt.gca()
for axis in ['top','bottom','left','right']:
    ax.spines[axis].set_linewidth(axes_linewidth)
#plt.savefig('DOS-adaptive-kgrid.pdf', bbox_inches='tight')
plt.show()