from aiida import load_profile
import aiida.orm as orm
from io import TextIOWrapper
import matplotlib.pyplot as plt
import matplotlib.transforms as mtransforms
import numpy as np
from scipy.integrate import quad
from scipy.interpolate import CubicSpline
import sisl

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

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

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

plt.rcParams['figure.figsize'] = (6*oldwidth, 6*oldheight)
plt.rcParams.update({'font.size': 64, 'ps.fonttype': 42, 'axes.linewidth': 1})

def construct_J_pair_matrices(struct, arraydata, return_sc=False):
    # NB: cell should be an array whose columns are the lattice vectors
    na = len(struct.sites)
    is_present = np.zeros((na, na), dtype=bool)
    vals = np.empty((na, na))
    scs = np.empty((na, na), dtype=object)
    cell = np.array(struct.cell).T

    # Construct matrices
    Js = arraydata.get_array('Jiso')
    pairs = arraydata.get_array('pair')
    sc = arraydata.get_array('sc')
    
    for i, J in enumerate(Js):
        ia, ja = pairs[i]
        if is_present[ia, ja]:
            if np.linalg.norm(cell @ sc[i]) < np.linalg.norm(cell @ scs[ia, ja]):
                # This one is preferred because smaller distance between pairs
                vals[ia, ja] = J
                vals[ja, ia] = J
                scs[ia, ja] = sc[i]
                scs[ja, ia] = sc[i]
            elif np.allclose(sc[i], scs[ia, ja]):
                if not np.isclose(J, vals[ia, ja]):
                    warnings.warn(f'Inconsistency in J for pair {ia}, {ja}, |deltaJ| = {1000*np.abs(J - vals[ia, ja])} meV')
        else:
            vals[ia, ja] = J
            vals[ja, ia] = J
            is_present[ia, ja] = True
            is_present[ja, ia] = True
            scs[ia, ja] = sc[i]
            scs[ja, ia] = sc[i]

    if return_sc:
        return is_present, vals, scs
    else:
        return is_present, vals

def get_orb_idx(geom, atom_idx=None, species_label=None, orb_label=None, return_labels=False):
    idx = []
    labels = []
    if orb_label is None:
        # By default accept all orbitals
        orb_label = ''
        warnings.warn('"orb_label" was not set so defaulting to empty string (include all orbitals)')
    for i, a in enumerate(geom.atoms):
        include_orb = False
        if atom_idx is not None:
            include_orb = i == atom_idx
            if include_orb and species_label is not None:
                # Check if atom_idx and species_label are consistent
                if not a.tag == species_label:
                    raise ValueError(f'Both "atom_idx" and "species_label" were set but atom {atom_idx} is not of species {species_label} (actual species: {a.tag})')
        elif species_label is not None:
            include_orb = a.tag == species_label
        else:
            raise TypeError('Should set at least one of "atom_idx" or "species_label"')
        if include_orb:
            io = geom.a2o(i)
            ios = []
            ls = []
            for j, o in enumerate(a.orbitals):
                if orb_label in o.name():
                    ios.append(io+j)
                    ls.append(o.name())
            idx.append(ios)
            labels.append(ls)
    if return_labels:
        return idx, labels
    else:
        return idx

def wrap_buffer(buffer):
    from zipfile import ZipExtFile
    if isinstance(buffer, ZipExtFile):
        return TextIOWrapper(buffer)
    else:
        buffer.readable = buffer._fhandle.readable
        buffer.writable = buffer._fhandle.writable
        buffer.closed = buffer._fhandle.closed
        buffer.flush = buffer._fhandle.flush
        return TextIOWrapper(buffer)

NanoRibbon_Se2_soc_armchair7_sep3_relaxed_pdos1 = orm.load_node(uuid='856475c1-bc32-43b6-801f-510873fdaab3')

with NanoRibbon_Se2_soc_armchair7_sep3_relaxed_pdos1.outputs.retrieved.open('aiida.PDOS.xml', 'rb') as fh:
    pdos_data = sisl.get_sile_class('PDOS.xml')(wrap_buffer(fh)).read_data()

NanoRibbon_Se2_soc_armchair7_sep3_relaxed_FM_testrun1b = orm.load_node(uuid='a2de1d08-13a9-48a5-9fc0-5624e6761d8f')

# Get indices of atoms for which we want the PDOS

geom = sisl.Geometry.new.ase(NanoRibbon_Se2_soc_armchair7_sep3_relaxed_FM_testrun1b.inputs.structure.get_ase())

fxyz = np.array([(row[0], row[1], row[2]) for row in np.round(geom.fxyz, decimals=2)], dtype=[('a', np.float64), ('b', np.float64), ('c', np.float64)])

idx_a_c = np.argsort(fxyz, order=('a', 'c'))

idx_Cr = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cr')]

idx_Nb = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Nb')]

idx_Cl = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cl')]

idx_Cl_lower = np.arange(geom.na)[idx_Cl][(geom.xyz[idx_Cl,2] > 16.)]

def get_idx_CrCl3(iCr, geom, idx_Cl_lower):
    # Get idx of Cr and 3 nearest neighbour Cl which are closest to NbSe2
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=3, atoms=idx_Cl_lower)
    return geom.sc2uc(idx)

def get_idx_Nb(iCr, geom, idx_Nb):
    # Get idx of 3 nearest neighbour Nb atoms to a Cr
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=7, atoms=idx_Nb)
    return geom.sc2uc(idx)

linewidth = 5
axis_linewidth = 3

Es = pdos_data[1]
Eidx = np.logical_and(Es >= -0.15, Es <= 0.15)
x = Es[Eidx]

idx_Cr1 = np.flip(np.roll(idx_Cr[::2], -1))
spin = 0
for i in range(int(np.ceil(idx_Cr1.size/2))):
    idx = get_orb_idx(pdos_data[0], atom_idx=idx_Cr1[i], species_label='Cr', orb_label='3d')
    ydata = pdos_data[2][:,:,:]
    ydata_plot = np.sum(ydata[spin,idx[0],:], axis=0)
    plt.plot(Es, ydata_plot+i*0.1, linewidth=linewidth)
for y in np.linspace(0.0, 0.6, num=7):
    plt.hlines(y, -0.3, 0.3, color='k', linewidth=axis_linewidth)
plt.grid(visible=True, linestyle='--')
for axis in ['top','bottom','left','right']:
    plt.gca().spines[axis].set_linewidth(axis_linewidth)
plt.ylim((0, 0.7))
plt.xlim((-0.1, 0.1))
plt.xlabel(r'$E - E_F$ (eV)')
plt.ylabel('PDOS (1/eV)')
ticks = ['']*14
ticks[::2] = np.round(np.linspace(0., 0.6, num=7), decimals=1)
plt.yticks(np.linspace(0.0, 0.65, num=14), ticks)
plt.xticks(np.linspace(-0.08, 0.08, num=5))
ax2 = plt.gca().twinx()
ax2.set_ylabel('Position along ribbon')
ax2.set_ylim((0, 0.7))
ax2.set_yticks([0.05, 0.65], ['center', 'edge'])
#plt.savefig('Cr-3d-atom-separated-DOS.pdf', bbox_inches='tight')
plt.show()

TB2J_dict = {}
pdos_dict = {}

NanoRibbon_soc_armchair7_sep3_relaxed_pdos1b_finer_grids = orm.load_node(uuid='53664383-afd2-4495-8000-b3c62beff22c')

with NanoRibbon_soc_armchair7_sep3_relaxed_pdos1b_finer_grids.outputs.retrieved.open('aiida.PDOS.xml', 'rb') as fh:
    pdos_data = sisl.get_sile_class('PDOS.xml')(wrap_buffer(fh)).read_data()

# Get indices of atoms for which we want PDOS

geom = sisl.Geometry.new.ase(NanoRibbon_soc_armchair7_sep3_relaxed_pdos1b_finer_grids.inputs.structure.get_ase())

fxyz = np.array([(row[0], row[1], row[2]) for row in np.round(geom.fxyz, decimals=2)], dtype=[('a', np.float64), ('b', np.float64), ('c', np.float64)])

idx_a_c = np.argsort(fxyz, order=('a', 'c'))

idx_Cr = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cr')]

# Get PDOS at Ef

idx_Cr1 = np.flip(idx_Cr[::2])
spin = 0
dE = 2e-3
DOS_at_Ef = []
for i in range(int(np.ceil(idx_Cr1.size/2))):
    idx = get_orb_idx(pdos_data[0], atom_idx=idx_Cr1[i], species_label='Cr', orb_label='3d')
    ydata = pdos_data[2][:,:,:]
    ydata_plot = np.sum(ydata[spin,idx[0],:], axis=0)
    pdos_f = CubicSpline(pdos_data[1], ydata_plot)
    DOS_at_Ef.append(quad(pdos_f, -dE, dE)[0])

# Retrieve J's

NanoRibbon_soc_armchair7_sep3_TB2J_testrun1 = orm.load_node(uuid='02a32d39-3f17-4e70-9ce2-13cfeba85b54')
input_struct = NanoRibbon_soc_armchair7_sep3_TB2J_testrun1.inputs.nodes.remote_data.creator.inputs.remote_stash.creator.outputs.output_structure

is_present, vals = construct_J_pair_matrices(
    input_struct,
    NanoRibbon_soc_armchair7_sep3_TB2J_testrun1.outputs.exchange
)

# Get indices of atoms for which we want J's

geom = sisl.Geometry.new.ase(input_struct.get_ase())

fxyz = np.array([(row[0], row[1], row[2]) for row in np.round(geom.fxyz, decimals=2)], dtype=[('a', np.float64), ('b', np.float64), ('c', np.float64)])

idx_a_c = np.argsort(fxyz, order=('a', 'c'))

idx_Cr = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cr')]

idx_Nb = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Nb')]

idx_Cl = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cl')]

idx_Cl_lower = np.arange(geom.na)[idx_Cl][(geom.xyz[idx_Cl,2] > 16.)]

def get_idx_CrCl3(iCr, geom, idx_Cl_lower):
    # Get idx of Cr and 3 nearest neighbour Cl which are closest to NbSe2
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=3, atoms=idx_Cl_lower)
    return geom.asc2uc(idx)

def get_idx_Nb(iCr, geom, idx_Nb):
    # Get idx of 3 nearest neighbour Nb atoms to a Cr
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=6.5, atoms=idx_Nb)
    return geom.asc2uc(idx)

# Extract J's for the Cr-Nb pairs we are interested in
idx_Cr1 = np.flip(idx_Cr[::2])
nJs = int(np.ceil(idx_Cr1.size/2))
ydata_TB2J = np.empty(nJs)
for i in range(nJs):
    ij = idx_Cr1[i]
    kjs = get_idx_Nb(ij, geom, idx_Nb)
    max = 0
    for iNbSe2 in kjs:
        if np.abs(vals[ij, iNbSe2]) > max:
            max = vals[ij, iNbSe2]
    ydata_TB2J[i] = 1000*max

TB2J_dict['Nb-above, FM, U1'] = np.flip(ydata_TB2J)

pdos_dict['Nb-above, FM, U1'] = np.flip(DOS_at_Ef)

NanoRibbon_Se2_soc_armchair7_sep3_relaxed_pdos1 = orm.load_node(uuid='856475c1-bc32-43b6-801f-510873fdaab3')

with NanoRibbon_Se2_soc_armchair7_sep3_relaxed_pdos1.outputs.retrieved.open('aiida.PDOS.xml', 'rb') as fh:
    pdos_data = sisl.get_sile_class('PDOS.xml')(wrap_buffer(fh)).read_data()

NanoRibbon_Se2_soc_armchair7_sep3_TB2J_FM_testrun1 = orm.load_node(uuid='780a26c3-53d3-446f-b2e4-1c9ee80f2b93')
NanoRibbon_Se2_soc_armchair7_sep3_relaxed_FM_testrun1b = orm.load_node(uuid='a2de1d08-13a9-48a5-9fc0-5624e6761d8f')
NanoRibbon_Se2_soc_armchair7_sep3_TB2J_FM_testrun1b = orm.load_node(uuid='2e0cf3e7-7323-47c0-bbe4-dbb3c459bee5')

is_present, vals = construct_J_pair_matrices(
    NanoRibbon_Se2_soc_armchair7_sep3_relaxed_FM_testrun1b.inputs.structure,
    NanoRibbon_Se2_soc_armchair7_sep3_TB2J_FM_testrun1b.outputs.exchange
)

# Get indices of atoms for which we want PDOS and J's

geom = sisl.Geometry.new.ase(NanoRibbon_Se2_soc_armchair7_sep3_relaxed_FM_testrun1b.inputs.structure.get_ase())

fxyz = np.array([(row[0], row[1], row[2]) for row in np.round(geom.fxyz, decimals=2)], dtype=[('a', np.float64), ('b', np.float64), ('c', np.float64)])

idx_a_c = np.argsort(fxyz, order=('a', 'c'))

idx_Cr = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cr')]

idx_Nb = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Nb')]

idx_Cl = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cl')]

idx_Cl_lower = np.arange(geom.na)[idx_Cl][(geom.xyz[idx_Cl,2] > 16.)]

def get_idx_CrCl3(iCr, geom, idx_Cl_lower):
    # Get idx of Cr and 3 nearest neighbour Cl which are closest to NbSe2
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=3, atoms=idx_Cl_lower)
    return geom.asc2uc(idx)

def get_idx_Nb(iCr, geom, idx_Nb):
    # Get idx of 3 nearest neighbour Nb atoms to a Cr
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=7, atoms=idx_Nb)
    return geom.asc2uc(idx)

# Extract J's for the Cr-Nb pairs we are interested in
idx_Cr1 = np.flip(np.roll(idx_Cr[::2], -1))
nJs = int(np.ceil(idx_Cr1.size/2))
ydata_TB2J = np.empty(nJs)
for i in range(nJs):
    ij = idx_Cr1[i]
    kjs = get_idx_Nb(ij, geom, idx_Nb)
    max = 0
    for iNbSe2 in kjs:
        if np.abs(vals[ij, iNbSe2]) > max:
            max = vals[ij, iNbSe2]
    ydata_TB2J[i] = max*1000

# Get PDOS at Ef

idx_Cr1 = np.flip(np.roll(idx_Cr[::2], -1))
spin = 0
dE = 2e-3
DOS_at_Ef = []
for i in range(int(np.ceil(idx_Cr1.size/2))):
    idx = get_orb_idx(pdos_data[0], atom_idx=idx_Cr1[i], species_label='Cr', orb_label='3d')
    ydata = pdos_data[2][:,:,:]
    ydata_plot = np.sum(ydata[spin,idx[0],:], axis=0)
    pdos_f = CubicSpline(pdos_data[1], ydata_plot)
    DOS_at_Ef.append(quad(pdos_f, -dE, dE)[0])

TB2J_dict['Se-below, FM, U1'] = ydata_TB2J

pdos_dict['Se-below, FM, U1'] = np.array(DOS_at_Ef)

# Retrieve J's

NanoRibbon_Se2_soc_armchair7_sep3_TB2J_AFM_testrun1 = orm.load_node(uuid='908d15d8-4d77-428c-8cba-fdb56a68f144')
input_struct = NanoRibbon_Se2_soc_armchair7_sep3_TB2J_AFM_testrun1.inputs.nodes.remote_data.creator.inputs.remote_stash.creator.inputs.structure

is_present, vals = construct_J_pair_matrices(
    input_struct,
    NanoRibbon_Se2_soc_armchair7_sep3_TB2J_AFM_testrun1.outputs.exchange
)

# Get indices of atoms for which we want J's

geom = sisl.Geometry.new.ase(input_struct.get_ase())

fxyz = np.array([(row[0], row[1], row[2]) for row in np.round(geom.fxyz, decimals=2)], dtype=[('a', np.float64), ('b', np.float64), ('c', np.float64)])

idx_a_c = np.argsort(fxyz, order=('a', 'c'))

idx_Cr = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cr')]

idx_Nb = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Nb')]

idx_Cl = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cl')]

idx_Cl_lower = np.arange(geom.na)[idx_Cl][(geom.xyz[idx_Cl,2] > 16.)]

def get_idx_CrCl3(iCr, geom, idx_Cl_lower):
    # Get idx of Cr and 3 nearest neighbour Cl which are closest to NbSe2
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=3, atoms=idx_Cl_lower)
    return geom.asc2uc(idx)

def get_idx_Nb(iCr, geom, idx_Nb):
    # Get idx of 3 nearest neighbour Nb atoms to a Cr
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=7, atoms=idx_Nb)
    return geom.asc2uc(idx)

# Extract J's for the Cr-Nb pairs we are interested in
idx_Cr1 = np.flip(np.roll(idx_Cr[::2], -1))
nJs = int(np.ceil(idx_Cr1.size/2))
ydata_TB2J = np.empty(nJs)
for i in range(nJs):
    ij = idx_Cr1[i]
    kjs = get_idx_Nb(ij, geom, idx_Nb)
    max = 0
    for iNbSe2 in kjs:
        if np.abs(vals[ij, iNbSe2]) > max:
            max = vals[ij, iNbSe2]
    ydata_TB2J[i] = max*1000

TB2J_dict['Se-below, AFM, U1'] = ydata_TB2J

# Retrieve J's

NanoRibbon_h2_soc_armchair7_sep3_TB2J_FM_run1 = orm.load_node(uuid='30fbea7c-bb7a-48a2-82f1-942738586c9d')
input_struct = NanoRibbon_h2_soc_armchair7_sep3_TB2J_FM_run1.inputs.nodes.remote_data.creator.inputs.remote_stash.creator.inputs.structure

is_present, vals = construct_J_pair_matrices(
    input_struct,
    NanoRibbon_h2_soc_armchair7_sep3_TB2J_FM_run1.outputs.exchange
)

# Get indices of atoms for which we want J's

geom = sisl.Geometry.new.ase(input_struct.get_ase())

fxyz = np.array([(row[0], row[1], row[2]) for row in np.round(geom.fxyz, decimals=2)], dtype=[('a', np.float64), ('b', np.float64), ('c', np.float64)])

idx_a_c = np.argsort(fxyz, order=('a', 'c'))

idx_Cr = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cr')]

idx_Nb = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Nb')]

idx_Cl = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cl')]

idx_Cl_lower = np.arange(geom.na)[idx_Cl][(geom.xyz[idx_Cl,2] > 16.)]

def get_idx_CrCl3(iCr, geom, idx_Cl_lower):
    # Get idx of Cr and 3 nearest neighbour Cl which are closest to NbSe2
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=3, atoms=idx_Cl_lower)
    return geom.asc2uc(idx)

def get_idx_Nb(iCr, geom, idx_Nb):
    # Get idx of 3 nearest neighbour Nb atoms to a Cr
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=7, atoms=idx_Nb)
    return geom.asc2uc(idx)

# Extract J's for the Cr-Nb pairs we are interested in
idx_Cr1 = np.flip(np.roll(idx_Cr[::2], -1))
nJs = int(np.ceil(idx_Cr1.size/2))
ydata_TB2J = np.empty(nJs)
for i in range(nJs):
    ij = idx_Cr1[i]
    kjs = get_idx_Nb(ij, geom, idx_Nb)
    max = 0
    for iNbSe2 in kjs:
        if np.abs(vals[ij, iNbSe2]) > max:
            max = vals[ij, iNbSe2]
    ydata_TB2J[i] = max*1000

TB2J_dict['h-below, FM, U1'] = ydata_TB2J

NanoRibbon_Se2_soc_armchair7_sep3_relaxed_abinitio_Us_AFM_pdos1 = orm.load_node(uuid='3dc9a92f-131d-4165-b659-f0edfc789018')

with NanoRibbon_Se2_soc_armchair7_sep3_relaxed_abinitio_Us_AFM_pdos1.outputs.retrieved.open('aiida.PDOS.xml', 'rb') as fh:
    pdos_data = sisl.get_sile_class('PDOS.xml')(wrap_buffer(fh)).read_data()

# Retrieve J's

NanoRibbon_Se2_soc_armchair7_sep3_TB2J_AFM_abinitio_Us_run1 = orm.load_node(uuid='251fb04b-36fb-45ce-9797-60549e171872')
input_struct = NanoRibbon_Se2_soc_armchair7_sep3_TB2J_AFM_abinitio_Us_run1.inputs.nodes.remote_data.creator.inputs.remote_stash.creator.inputs.structure

is_present, vals = construct_J_pair_matrices(
    input_struct,
    NanoRibbon_Se2_soc_armchair7_sep3_TB2J_AFM_abinitio_Us_run1.outputs.exchange
)

# Get indices of atoms for which we want PDOS and J's

geom = sisl.Geometry.new.ase(input_struct.get_ase())

fxyz = np.array([(row[0], row[1], row[2]) for row in np.round(geom.fxyz, decimals=2)], dtype=[('a', np.float64), ('b', np.float64), ('c', np.float64)])

idx_a_c = np.argsort(fxyz, order=('a', 'c'))

idx_Cr = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cr')]

idx_Nb = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Nb')]

idx_Cl = idx_a_c[geom.atoms.Z[idx_a_c] == sisl.PeriodicTable().Z('Cl')]

idx_Cl_lower = np.arange(geom.na)[idx_Cl][(geom.xyz[idx_Cl,2] > 16.)]

def get_idx_CrCl3(iCr, geom, idx_Cl_lower):
    # Get idx of Cr and 3 nearest neighbour Cl which are closest to NbSe2
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=3, atoms=idx_Cl_lower)
    return geom.asc2uc(idx)

def get_idx_Nb(iCr, geom, idx_Nb):
    # Get idx of 3 nearest neighbour Nb atoms to a Cr
    geom.set_nsc(a=3, b=3, c=3)
    idx = geom.close(iCr, R=7, atoms=idx_Nb)
    return geom.asc2uc(idx)

# Extract J's for the Cr-Nb pairs we are interested in
idx_Cr1 = np.flip(np.roll(idx_Cr[::2], -1))
nJs = int(np.ceil(idx_Cr1.size/2))
ydata_TB2J = np.empty(nJs)
for i in range(nJs):
    ij = idx_Cr1[i]
    kjs = get_idx_Nb(ij, geom, idx_Nb)
    max = 0
    for iNbSe2 in kjs:
        if np.abs(vals[ij, iNbSe2]) > max:
            max = vals[ij, iNbSe2]
    ydata_TB2J[i] = max*1000

# Get PDOS at Ef

idx_Cr1 = np.flip(np.roll(idx_Cr[::2], -1))
spin = 0
dE = 2e-3
DOS_at_Ef = []
for i in range(int(np.ceil(idx_Cr1.size/2))):
    idx = get_orb_idx(pdos_data[0], atom_idx=idx_Cr1[i], species_label='Cr', orb_label='3d')
    ydata = pdos_data[2][:,:,:]
    ydata_plot = np.sum(ydata[spin,idx[0],:], axis=0)
    pdos_f = CubicSpline(pdos_data[1], ydata_plot)
    DOS_at_Ef.append(quad(pdos_f, -dE, dE)[0])

TB2J_dict['Se-below, AFM, U2'] = ydata_TB2J

pdos_dict['Se-below, AFM, U2'] = np.array(DOS_at_Ef)

plt.rcParams.update({'font.size': 32, 'ps.fonttype': 42, 'axes.linewidth': 1})

markersize=16
linewidth = 5
axis_linewidth = 3

fig = plt.figure(figsize=(6*oldwidth, 2.25*oldheight))
axs = fig.subplot_mosaic([['a)', 'b)']])
plt.subplots_adjust(wspace=0.2, hspace=0.0)

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

for l, ydata in pdos_dict.items():
    axs['a)'].plot(ydata, 's--', markersize=markersize, label=l)

axs['a)'].set_xlabel(r'Cr index')
axs['a)'].set_ylabel(f'Cr 3d PDOS at $E_F$')
axs['a)'].set_xticks(ticks=np.arange(7), labels=np.arange(1,8))
axs['a)'].ticklabel_format(axis='y', scilimits=(0,0))
axs['a)'].grid(visible=True, linestyle='--')
axs['a)'].legend()
for axis in ['top','bottom','left','right']:
    axs['a)'].spines[axis].set_linewidth(axis_linewidth)

for l, ydata in TB2J_dict.items():
    axs['b)'].plot(ydata**2, 's--', label=l, markersize=markersize)
axs['b)'].set_ylabel(r'$J^2_\mathrm{Cr-Nb}$ (meV$^2$)')
axs['b)'].set_xlabel(r'Cr index')
axs['b)'].set_xticks(ticks=np.arange(7), labels=np.arange(1,8))#, labels=['center', 'edge'])
axs['b)'].grid(visible=True, linestyle='--')
axs['b)'].legend()
for axis in ['top','bottom','left','right']:
    axs['b)'].spines[axis].set_linewidth(axis_linewidth)
    
#fig.savefig('pdos-J2-trends.pdf', bbox_inches='tight')
plt.show()