%load_ext aiida

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

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, 1.8*oldwidth)
plt.rcParams.update({'font.size': 50})

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

axes_linewidth = 4
plot_linewidth = 5

bulk_Nb_fixeddelta_pdos_eT1 = orm.load_node(uuid='6f8d6c3c-4a0c-4b2f-9435-6b9e47621b32')
bulk_Nb_fixeddelta_pdos_eT1_01K_adaptive = orm.load_node(uuid='81d72ff3-bc47-4a2d-a2fd-f20c57c8fa74')

calcnode = bulk_Nb_fixeddelta_pdos_eT1

def get_pdos(node):
    with node.outputs.retrieved.open('aiida.PDOS.xml') as fh:
        geom, E, PDOS = sisl.get_sile_class('PDOS.xml')(fh).read_data()
    return E, PDOS

def collect_from_aiida_iterator(root_node, func, container='array', dtype=float, ignore_failed=False):
    """Collect results from AiiDA iterator workflow by applying func to each workflow node"""

    if isinstance(root_node, orm.WorkChainNode):
        node_list = root_node.called
    elif isinstance(root_node, list):
        node_list = root_node
    else:
        raise ValueError('First argument should be WorkChainNode or list!')
    
    # Setup container to store results
    n = len(node_list)
    idx_okay = np.full(n, True, dtype=bool)
    if container == 'array':
        res = np.empty(n, dtype=dtype)
    else:
        res = [None]*n

    # Collect results
    for i in range(n):
        try:
            res[i] = func(node_list[i])
        except Exception as e:
            print(e)
            res[i] = np.nan
            idx_okay[i] = False

    if not np.all(idx_okay):
        return res, idx_okay
    else:
        return res

PDOSs = collect_from_aiida_iterator(calcnode, get_pdos, dtype='object')

eTs = collect_from_aiida_iterator(calcnode, lambda node: float(node.inputs.parameters['electronictemperature'].split()[0]))

idx = np.argsort(eTs)

PDOSs = PDOSs[idx]
eTs = eTs[idx]

PDOS_T_01K = get_pdos(bulk_Nb_fixeddelta_pdos_eT1_01K_adaptive)

bulk_Nb_fixeddelta_pdos_kgrid_smearing1 = orm.load_node(uuid='d9d82a6e-8d34-43e5-8408-b9a8a070c199')

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

E800, PDOS800 = PDOSs[1]

with nk_dict[400].outputs.retrieved.open('aiida.PDOS.xml') as fh:
    geom, E400, PDOS400 = sisl.get_sile_class('PDOS.xml')(fh).read_data()

with nk_dict[100].outputs.retrieved.open('aiida.PDOS.xml') as fh:
    geom, E100, PDOS100 = sisl.get_sile_class('PDOS.xml')(fh).read_data()

(width, height) = plt.rcParams['figure.figsize']
fig, axs = plt.subplot_mosaic([['a)', 'b)']],
                              layout='constrained', figsize=(2*width, height))

colors = ['b', 'g', 'r', 'orange']
# Plot with adaptive sampling for 0.1 K
# axs['a)'].plot(1000*PDOS_T_01K[0], np.sum(PDOS_T_01K[1][0,:,:], axis=0), label='0.1 K', linewidth=plot_linewidth, color=colors[i])
for i, eT in enumerate(eTs):
    axs['a)'].plot(1000*PDOSs[i][0], np.sum(PDOSs[i][1][0,:,:], axis=0), label=f'{eT} K', linewidth=plot_linewidth, color=colors[i])
axs['a)'].legend(fontsize=32)
axs['a)'].set_xlabel(r'$E-E_F$ [meV]')
axs['a)'].set_ylabel(r'DOS [1/eV]')
axs['a)'].set_xlim((-4,4))
for axis in ['top','bottom','left','right']:
    axs['a)'].spines[axis].set_linewidth(axes_linewidth)

axs['b)'].plot(1000*E100, np.sum(PDOS100[0,:,:], axis=0), 'b-', label=r'$n_k = 100$', linewidth=plot_linewidth)
axs['b)'].plot(1000*E400, np.sum(PDOS400[0,:,:], axis=0), 'r-', label=r'$n_k = 400$', linewidth=plot_linewidth)
axs['b)'].plot(1000*E800, np.sum(PDOS800[0,:,:], axis=0), 'g-', label=r'$n_k = 800$', linewidth=plot_linewidth)
axs['b)'].legend(fontsize=32)
axs['b)'].set_xlabel(r'$E - E_F$ [meV]')
axs['b)'].set_ylabel(r'DOS [1/eV]')
axs['b)'].set_xlim((-4,4))
for axis in ['top','bottom','left','right']:
    axs['b)'].spines[axis].set_linewidth(axes_linewidth)    

for label, ax in axs.items():
    # label physical distance to the left and up:
    trans = mtransforms.ScaledTranslation(-50/72, 20/72, fig.dpi_scale_trans)
    ax.text(0.0, 1.0, label, transform=ax.transAxes + trans,
            fontsize='large', va='bottom')

#fig.savefig('convergence-testing.pdf', bbox_inches='tight')
plt.show()