%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 = 6
plot_linewidth = 3
plot_markersize = 15

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 geom, 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

bulk_Nb_fixeddelta_pdos_model_different_Delta1 = orm.load_node(uuid='1ccd750d-5ee4-4d3b-a4b0-0fa61f3ef8e2')
bulk_Nb_fixeddelta_pdos_model_different_Delta1b = orm.load_node(uuid='8da3e68b-ebfb-4a2c-b6ec-fa2e843b5949')
bulk_Nb_oneshot_pdos_model_different_Delta1a = orm.load_node(uuid='5e477894-fb40-461d-8620-e22686ca5faa')
bulk_Nb_oneshot_pdos_model_different_Delta1b = orm.load_node(uuid='e7a03d85-8070-4588-8019-76744701b947')

calcnode = bulk_Nb_fixeddelta_pdos_model_different_Delta1

delta_fixeddelta = collect_from_aiida_iterator(calcnode, lambda node: float(node.inputs.parameters['%block bdgdeltapairingpairing1'].split(' ')[13]))

pdos_fixeddelta = collect_from_aiida_iterator(calcnode, lambda node: get_pdos(node)[2], dtype=object)

Es_fixeddelta = collect_from_aiida_iterator(calcnode, lambda node: get_pdos(node)[1], dtype=object)

idx = np.argsort(delta_fixeddelta)
delta_fixeddelta = delta_fixeddelta[idx]
pdos_fixeddelta = pdos_fixeddelta[idx]
Es_fixeddelta = Es_fixeddelta[idx]

calcnode = bulk_Nb_fixeddelta_pdos_model_different_Delta1b

delta_fixeddelta2 = collect_from_aiida_iterator(calcnode, lambda node: float(node.inputs.parameters['%block bdgdeltapairingpairing1'].split(' ')[13]))

pdos_fixeddelta2 = collect_from_aiida_iterator(calcnode, lambda node: get_pdos(node)[2], dtype=object)

Es_fixeddelta2 = collect_from_aiida_iterator(calcnode, lambda node: get_pdos(node)[1], dtype=object)

delta_fixeddelta2 = delta_fixeddelta2[Es_fixeddelta2[1]]
idx = np.argsort(delta_fixeddelta2)
delta_fixeddelta2 = delta_fixeddelta2[idx]
pdos_fixeddelta2 = np.stack(pdos_fixeddelta2[0][Es_fixeddelta2[1]])
pdos_fixeddelta2 = pdos_fixeddelta2[idx]
Es_fixeddelta2 = np.stack(Es_fixeddelta2[0][Es_fixeddelta2[1]])
Es_fixeddelta2 = Es_fixeddelta2[idx]

delta_fixeddelta = np.concatenate((delta_fixeddelta2, delta_fixeddelta))
pdos_fixeddelta = np.concatenate((pdos_fixeddelta2, np.stack(pdos_fixeddelta)))
Es_fixeddelta = np.concatenate((Es_fixeddelta2, np.stack(Es_fixeddelta)))

calcnode = bulk_Nb_oneshot_pdos_model_different_Delta1a

delta_oneshot = collect_from_aiida_iterator(calcnode, lambda node: float(node.inputs.parameters['%block bdgdeltapairingpairing1'].split(' ')[13]))

pdos_oneshot = collect_from_aiida_iterator(calcnode, lambda node: get_pdos(node)[2], dtype=object)

Es_oneshot = collect_from_aiida_iterator(calcnode, lambda node: get_pdos(node)[1], dtype=object)

idx = np.argsort(delta_oneshot)
delta_oneshot = delta_oneshot[idx]
pdos_oneshot = pdos_oneshot[idx]
Es_oneshot = Es_oneshot[idx]

calcnode = bulk_Nb_oneshot_pdos_model_different_Delta1b

delta_oneshot2 = collect_from_aiida_iterator(calcnode, lambda node: float(node.inputs.parameters['%block bdgdeltapairingpairing1'].split(' ')[13]))

pdos_oneshot2 = collect_from_aiida_iterator(calcnode, lambda node: get_pdos(node)[2], dtype=object)

Es_oneshot2 = collect_from_aiida_iterator(calcnode, lambda node: get_pdos(node)[1], dtype=object)

delta_oneshot2 = delta_oneshot2[Es_oneshot2[1]]
idx = np.argsort(delta_oneshot2)
delta_oneshot2 = delta_oneshot2[idx]
pdos_oneshot2 = np.stack(pdos_oneshot2[0][Es_oneshot2[1]])
pdos_oneshot2 = pdos_oneshot2[idx]
Es_oneshot2 = np.stack(Es_oneshot2[0][Es_oneshot2[1]])
Es_oneshot2 = Es_oneshot2[idx]

delta_oneshot = np.concatenate((delta_oneshot2, delta_oneshot))
pdos_oneshot = np.concatenate((pdos_oneshot2, np.stack(pdos_oneshot)))
Es_oneshot = np.concatenate((Es_oneshot2, np.stack(Es_oneshot)))

dEs = Es_fixeddelta[:,1] - Es_fixeddelta[:,0]

spin = 0
integrated_differences_total = 100*np.sum(np.abs(np.sum(pdos_fixeddelta[:,spin,:,:], axis=1) - np.sum(pdos_oneshot[:,spin,:,:], axis=1))*dEs.reshape(-1,1), axis=-1) / np.sum(np.sum(pdos_fixeddelta[:,spin,:,:], axis=1)*dEs.reshape(-1,1), axis=-1)

plt.plot(delta_fixeddelta, integrated_differences_total, 'rs--', linewidth=plot_linewidth, markersize=plot_markersize)
plt.xlabel(r'$\Delta$ [meV]')
plt.ylabel('DOS Integrated Difference [%]')
plt.xscale('log')
ax = plt.gca()
for axis in ['top','bottom','left','right']:
    ax.spines[axis].set_linewidth(axes_linewidth)
#plt.savefig('model-rel-int-diff.pdf', bbox_inches='tight')
plt.show()