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 os.path import splitext
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': 32, '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 write_vectors_to_xsf(xsf_file, vectors, precision=None):
    with open(xsf_file, 'r') as fh:
        # First read size of system
        for _ in range(6):
            fh.readline()
        na = int(fh.readline().split(' ')[0])
        if vectors.shape != (na, 3):
            raise ValueError(f'Vectors has wrong shape. Expected shape: ({na}, 3), Actual shape: {vectors.shape}')

        # Write new file with vectors
        fh.seek(0)
        base, ext = splitext(xsf_file)
        outfile = base + '.vector' + ext
        with open(outfile, 'w') as fh_out:
            # First write header
            for _ in range(7):
                fh_out.write(fh.readline())

            # For each atom also write the corresponding vector
            for i in range(na):
                fh_out.write(fh.readline().strip() + ' ' + np.array2string(vectors[i,:], precision=precision)[1:-1] + '\n')

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)

def get_mulliken_pops(node):
    with node.outputs.retrieved.open('aiida.out', 'rb') as fh:
        charges = sisl.io.siesta.stdoutSileSiesta(wrap_buffer(fh)).read_charge('mulliken:<5.2')
    return charges

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_FM2, vals_FM2 = 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 the indices of the atoms for which we want the PDOS and the 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)

# Retrieve the J's for the Cr-Nb pairs on the path from center to edge
idx_Cr1 = np.flip(np.roll(idx_Cr[::2], -1))
nJs = int(np.ceil(idx_Cr1.size/2))
ydata_TB2J = np.empty((nJs, 3))
for i in range(nJs):
    ij = idx_Cr1[i]
    kjs = get_idx_Nb(ij, geom, idx_Nb)
    Js = []
    for iNbSe2 in kjs:
        Js.append(vals_FM2[ij, iNbSe2]**2)
        #print(ij, kjs, is_present_FM2[ij, iNbSe2])
    ydata_TB2J[i, :] = Js

markersize=16
linewidth = 5
axis_linewidth = 3

fig = plt.figure(figsize=(6*oldwidth, 4.5*oldheight))
axs = fig.subplot_mosaic([['a)', 'b)'], ['c)', 'd)']])
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')
    if label in ['a)', 'b)']:
        ax.set_axis_off()

geom_top = plt.imread('nanoribbon-armchair7-sep3-Se2.top.png')
axs['a)'].imshow(geom_top)
box = axs['a)'].get_position()
box.y0 = box.y0 + 0.07
box.y1 = box.y1 + 0.07
box.x0 = box.x0 - 0.01
box.x1 = box.x1 - 0.01

geom_side = plt.imread('nanoribbon-armchair7-sep3-Se2.side.png')
axs['b)'].imshow(geom_side)
box = axs['b)'].get_position()
box.y0 = box.y0 + 0.07
box.y1 = box.y1 + 0.07

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])

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

labels = [r'Nb$_1$', r'Nb$_2$', r'Nb$_3$']
for j in range(3):
    axs['d)'].plot(1000*1000*ydata_TB2J[:,j], 's--', label=labels[j], markersize=16)
axs['d)'].set_ylabel(r'$J^2_\mathrm{Cr-Nb}$ (meV$^2$)')
axs['d)'].set_xlabel(r'Cr index')
axs['d)'].set_xticks(ticks=np.arange(7), labels=np.arange(1,8))#, labels=['center', 'edge'])
axs['d)'].grid(visible=True, linestyle='--')
axs['d)'].legend()
for axis in ['top','bottom','left','right']:
    axs['d)'].spines[axis].set_linewidth(axis_linewidth)

#fig.savefig('fig3.pdf', bbox_inches='tight')
plt.show()

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

fname = 'structure.xsf'
NanoRibbon_Se2_soc_armchair7_sep3_relaxed_FM_testrun1b.inputs.structure.export(fname)
write_vectors_to_xsf(fname, vectors=get_mulliken_pops(NanoRibbon_Se2_soc_armchair7_sep3_relaxed_FM_testrun1b)[0,:,2:])
# Now 'structure.vector.xsf' will contain the geometry together with Mulliken magnetic moments for each atom