"""Utility functions for the wannier_quality method of the spinWannier.WannierTBmodel.WannierTBmodel class.
"""

import pickle
import numpy as np
import matplotlib.pyplot as plt
import shutil
from spinWannier.wannier_utils import (
    real_to_W_gauge,
    W_gauge_to_H_gauge,
    split_spn_dict,
    get_kpoint_path,
    load_lattice_vectors,
    reciprocal_lattice_vectors,
    check_file_exists,
)
from mpl_toolkits.axes_grid1 import make_axes_locatable

# import sys
# sys.path.append('/home/lv268562/bin/python_scripts/wannier_quality/wann_quality_calculation')
# from wannier_quality import parse_eigenval_file, plot_err_vs_energy


def get_fermi(path="."):
    """Extract Fermi energy from the DOSCAR file.

    Args:
        path (str, optional): Path to the directory with the DOSCAR file. Defaults to ".".

    Returns:
        float: Fermi energy in eV.
    """
    with open(path + "/DOSCAR", "r") as fr:
        for i, line in enumerate(fr):
            if i == 5:
                EF = line.split()[-2]
                with open(path + "/FERMI_ENERGY.in", "w") as fw:
                    fw.write(
                        f"# FERMI ENERGY extracted with the get_EF.py script\n{EF}"
                    )
                return float(EF)


def get_band_at_kpoint_from_EIGENVAL(
    EIGENVAL_path="./EIGENVAL",
    target_band=1,
    target_kpoint_string="0.0000000E+00  0.0000000E+00  0.0000000E+00",
):
    """
    Get the energy of the target_band at the target_kpoint from the EIGENVAL file.

    Args:
        EIGENVAL_path (str, optional): Path to the EIGENVAL file. Defaults to './EIGENVAL'.
        target_band (int, optional): Target band. Defaults to 1.
        target_kpoint_string (str, optional): Target k-point string. Defaults to '0.0000000E+00  0.0000000E+00  0.0000000E+00'.

    Returns:
        float: Energy of the target_band at the target_kpoint.
    """
    with open(EIGENVAL_path, "r") as fr:
        lines = fr.readlines()
        for i, line in enumerate(lines):
            if target_kpoint_string in line:
                energy = float(lines[i + target_band].split()[1])
                return energy


def get_fermi_corrected_by_matching_bands(
    path=".",
    nsc_calculation_path="../0_nsc_for_wann_25x25_frozmaxmargin_0.2eV",
    corrected_at_kpoint="0.0000000E+00  0.0000000E+00  0.0000000E+00",
    corrected_at_band=11,
    sc_calculation_path="../sc",
    fout_name="FERMI_ENERGY_corrected.in",
):
    """Get the Fermi energy from the self-consistent calculation and correct it so that the band at the target_kpoint and target_band is at the same energy in the non-self-consistent calculation.

    Args:
        path (str, optional): Path to the directory with the DOSCAR file. Defaults to ".".
        nsc_calculation_path (str, optional): Path to the non-self-consistent calculation directory. Defaults to '../0_nsc_for_wann_25x25_frozmaxmargin_0.2eV'.
        corrected_at_kpoint (str, optional): Target k-point string. Defaults to '0.0000000E+00  0.0000000E+00  0.0000000E+00'.
        corrected_at_band (int, optional): Target band. Defaults to 11.
        sc_calculation_path (str, optional): Path to the self-consistent calculation directory. Defaults to "../sc".
        fout_name (str, optional): Output file name. Defaults to "FERMI_ENERGY_corrected.in".

    Returns:
        float: Corrected Fermi energy in eV.
    """

    Fermi_sc = get_fermi(sc_calculation_path)

    sc_band = get_band_at_kpoint_from_EIGENVAL(
        EIGENVAL_path=sc_calculation_path + "/EIGENVAL",
        target_band=corrected_at_band,
        target_kpoint_string=corrected_at_kpoint,
    )
    nsc_band = get_band_at_kpoint_from_EIGENVAL(
        EIGENVAL_path=nsc_calculation_path + "/EIGENVAL",
        target_band=corrected_at_band,
        target_kpoint_string=corrected_at_kpoint,
    )

    # WE REQUIRE THAT:
    #   band_nsc - Fermi_nsc = band_sc - Fermi_sc
    Fermi_nsc = Fermi_sc + nsc_band - sc_band

    with open(nsc_calculation_path + "/" + fout_name, "w") as fw:
        fw.write(
            f"# {Fermi_nsc:.8f} eV = {Fermi_sc:.8f} + ({nsc_band:.8f} - {sc_band:.8f}) = Fermi_sc + (E_band{corrected_at_band}_nsc - E_band{corrected_at_band}_sc) ... self-consistent Fermi from {sc_calculation_path} corrected so that band {corrected_at_band} at k-point {corrected_at_kpoint} is at the same energy (relative to the recpective Fermi energies) for the self-consistent (path {sc_calculation_path}) and non-self-consistent (path {nsc_calculation_path}) calculation\n{Fermi_nsc:.8f}"
        )

    return float(Fermi_nsc)


def compare_eigs_bandstructure_at_exact_kpts(
    dft_bands,
    wann_bands,
    num_kpoints,
    num_wann,
    f_name_out="WannierBerri_quality_error_Fermi_corrected.dat",
):
    """Compare the DFT and Wannierized band structures at the exact k-points.

    Args:
        dft_bands (np.array): DFT bands.
        wann_bands (np.array): Wannierized bands.
        num_kpoints (int): Number of k-points.
        num_wann (int): Number of Wannierized bands.
        f_name_out (str, optional): Output file name. Defaults to 'WannierBerri_quality_error_Fermi_corrected.dat'.

    Returns:
        np.array: Array with the DFT energy in the first column and the Wannierization error in the second column.
    """

    error_by_energy = np.zeros((num_kpoints * num_wann, 2))
    for i in range(num_kpoints):
        # save the DFT energy as x variable
        error_by_energy[i * num_wann : (i + 1) * num_wann, 0] = dft_bands[i, :]

        # save the Wannierization error as y variable
        error_by_energy[i * num_wann : (i + 1) * num_wann, 1] = (
            dft_bands[i, :] - wann_bands[i, :]
        ) ** 2

    # sort 'error_by_energy' by energy - by the first column
    error_by_energy = error_by_energy[error_by_energy[:, 0].argsort()]

    # save to file
    header = "E (eV)\t(E_wann - E_DFT)**2"
    fmt = "%.12f\t%.12e"
    np.savetxt(f_name_out, error_by_energy, delimiter="\t", header=header, fmt=fmt)
    return error_by_energy


def duplicate_kpoints_for_home_made(data, NK):
    """Duplicate also the last k-point (in  dictionary the keys are unique, so actually the data in the dictionaries where keys are k-points contain only one of each k-point, so if k-path starts and ends with the same k-point, only the first one is recorded.

    Args:
        data (np.array): Data to duplicate.
        NK (int): Number of k-points.

    Returns:
        np.array: Duplicated data.
    """
    dimensions = list(np.array(data).shape)
    # print(dimensions[0])
    # print(NK)
    n_blocs = (dimensions[0]) // (NK - 1)
    dimensions[0] += n_blocs

    data_duplicates = np.zeros(dimensions)

    data_duplicates[:NK, :] = np.real(data[:NK, :])
    for i in range(1, n_blocs - 1):
        data_duplicates[i * NK, :] = np.real(data[i * (NK - 1), :])
        data_duplicates[i * NK + 1 : i * NK + 1 + NK, :] = np.real(data[
            1 + i * (NK - 1) : 1 + i * (NK - 1) + NK, :
        ])
    # last block different - missing last value (supposed to be the same as the first value)
    i = n_blocs - 1
    data_duplicates[i * NK, :] = np.real(data[i * (NK - 1), :])
    data_duplicates[i * NK + 1 : i * NK + 1 + NK - 2, :] = np.real(data[
        1 + i * (NK - 1) : 1 + i * (NK - 1) + NK, :
    ])
    data_duplicates[-1, :] = np.real(data[0, :])
    return data_duplicates


def get_NKpoints(OUTCAR="OUTCAR"):
    """Return number of kpoints from bands calculation from OUTCAR.

    Args:
        OUTCAR (str, optional): OUTCAR path. Defaults to 'OUTCAR'.

    Returns:
        int: number of k-points stated in the OUTCAR
    """
    with open(OUTCAR, "r") as fr:
        for line in fr:
            if "NKPTS =" in line:
                return int(line.split()[3])


def parse_eigenval_file(fin, spin=0):
    """Parse the EIGENVAL file and return the kpoints, bands, number of kpoints, and number of bands.

    Args:
        fin (str): Path to the EIGENVAL file.
        spin (int, optional): Spin index. Defaults to 0.

    Returns:
        np.array: kpoints.
        np.array: bands.
        int: number of kpoints.
        int: number of bands.
    """
    flag_kpoint = False
    flag_bands = False
    i_k = -1
    i_b = -1
    num_bands = 99999999
    with open(fin, "r") as fr:
        for i, line in enumerate(fr):

            l = line.split()

            if i == 5:
                num_bands = int(l[2])
                num_kpoints = int(l[1])
                kpoints = np.zeros((num_kpoints, 3))
                bands = np.zeros((num_kpoints, num_bands))

            if flag_bands == True:
                i_b += 1
                bands[i_k, i_b] = l[spin + 1]

            if flag_kpoint == True:
                i_k += 1
                kpoints[i_k, :] = l[0:3]
                flag_kpoint = False
                flag_bands = True

            # empty line
            if l == []:
                flag_kpoint = True
            else:
                if l[0] == str(num_bands):
                    flag_bands = False
                    i_b = -1

    kpoints = kpoints.astype(np.float32)
    bands = bands.astype(np.float32)

    kpoints_path = np.zeros((num_kpoints,))
    for i in range(num_kpoints - 1):
        kpoints_path[i + 1] = (
            kpoints_path[i] + np.sum((kpoints[i + 1, :] - kpoints[i, :]) ** 2) ** 0.5
        )

    return kpoints_path, bands, num_kpoints, num_bands


def plot_err_vs_energy(
    error_by_energy,
    Ef,
    title="Wannierization RMS error vs. energy",
    fig_name_out="wannier_quality_error_by_energy.png",
    savefig=True,
    showfig=True,
):
    """Plot the error vs. energy.

    Args:
        error_by_energy (np.array): Error vs. energy.
        Ef (float): Fermi energy.
        title (str, optional): Title of the plot. Defaults to "Wannierization RMS error vs. energy".
        fig_name_out (str, optional): Output file name. Defaults to "wannier_quality_error_by_energy.png".
        savefig (bool, optional): Save the figure. Defaults to True.
        showfig (bool, optional): Show the figure. Defaults to True.
    """
    plt.semilogy(
        error_by_energy[:, 0] - Ef,
        error_by_energy[:, 1] ** 0.5,
        "ok",
        linewidth=1,
        markersize=1,
    )
    plt.title(title, fontsize=8)
    plt.xlabel(r"$E - E_\mathrm{F}$ (eV)")
    plt.ylabel(r"$\|E_\mathrm{wann} - E_\mathrm{DFT}\|$ (eV)")
    if savefig:
        plt.savefig(fig_name_out, dpi=400)
    if showfig:
        plt.show()
    plt.close()


def plot_err_vs_bands(
    kpoints,
    kpath,
    kpath_ticks,
    Eigs_k,
    E_diff,
    S_diff,
    fout="ERRORS_ALL_band_structure.jpg",
    yaxis_lim=None,
    savefig=True,
    showfig=True,
):
    """Output a figure with RMSE_E, RMSE_Sx, RMSE_Sy, and RMSE_Sz-projected band structure.

    Args:
        kpoints (np.array): kpoints.
        kpath (np.array): kpath.
        kpath_ticks (list): kpath ticks.
        Eigs_k (dict): Eigs_k.
        E_diff (np.array): E_diff.
        S_diff (np.array): S_diff.
        fout (str, optional): Output file name. Defaults to 'ERRORS_ALL_band_structure.jpg'.
        yaxis_lim (list, optional): y-axis limits. Defaults to None.
        savefig (bool, optional): Save the figure. Defaults to True.
        showfig (bool, optional): Show the figure. Defaults to True.
    """
    NW = len(Eigs_k[list(Eigs_k.keys())[0]])
    Nk = len(kpoints) // (len(kpath_ticks) - 1)

    fig, axes = plt.subplots(2, 2, figsize=[9, 6])
    fig.suptitle("Error of Wannier interpolation", fontsize=14)
    spin_name = [r"Energy (eV)", r"$S_z$", r"$S_x$", r"$S_y$"]
    for i, S in enumerate([E_diff, S_diff[:, 2], S_diff[:, 0], S_diff[:, 1]]):
        ax = axes[i // 2, i % 2]
        ax.axhline(linestyle="--", color="k")
        if yaxis_lim:
            ax.set_ylim(yaxis_lim)
        ax.set_xlim([min(kpath), max(kpath)])
        ax.set_title(spin_name[i], fontsize=13, pad=8)

        if i == 0:
            # colorbar limits for energy
            vmin = 0
            vmax = 0.002
        else:
            # colorbar limits for spin
            vmin = 0
            vmax = 0.05
        sc = ax.scatter(
            [[k_dist for i in range(NW)] for k_dist in kpath],
            [Eigs_k[kpoint] for kpoint in kpoints],
            c=S.reshape(-1, NW),
            cmap="YlOrRd",
            s=0.2,
            vmin=vmin,
            vmax=vmax,
        )
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        cbar = fig.colorbar(sc, cax=cax, orientation="vertical")

        for j in range(1, len(kpath_ticks) - 1):
            ax.axvline(x=kpath[j * Nk], color="#000000", linestyle="-", linewidth=0.75)

        # if i == 0 or i == 2:
        ax.set_ylabel(r"$E - E_\mathrm{F}$ (eV)", fontsize=13)
        # if i == 0 or i == 1:
        secax = ax.secondary_xaxis("top")
        secax.tick_params(labelsize=9)  # axis='both', which='major',
        secax.set_xlabel(r"$k$-distance (1/$\mathrm{\AA}$)", fontsize=10)
        # no primary x-axis
        ax.set_xticks([])

        # if i == 2 or i == 3:
        idx = np.array(range(0, Nk * len(kpath_ticks), Nk))
        idx[-1] += -1
        ax.set_xticks(kpath[idx])
        ax.set_xticklabels(kpath_ticks, fontsize=12)
        ax.yaxis.set_tick_params(labelsize=11)

        # cbar.set_label(r'$S_\mathrm{z}$')
        # sc.set_clim(vmin=colorbar_Sz_lim[0], vmax=colorbar_Sz_lim[1])

    # plt.suptitle(fig_caption)
    plt.tight_layout()
    # plt.show()
    fout = check_file_exists(fout)
    if savefig:
        plt.savefig(fout, dpi=400)
    if showfig:
        plt.show()
    plt.close()
    # print(
    #     "ERROR band structure printed\n\n--------------------\n\n======================="
    # )


def integrate_error(error_by_energy, E_min=-1e3, E_max=1e3):
    """Integrate the error in 'f_name_in' in the energy range [E_min, E_max] included.

    Args:
        error_by_energy (np.array): Error vs. energy.
        E_min (float, optional): Minimum energy. Defaults to -1e3.
        E_max (float, optional): Maximum energy. Defaults to 1e3.

    Returns:
        np.array: Array with the integrated error.
    """
    E = error_by_energy[:, 0]
    err_array = error_by_energy[(E_min <= E) & (E <= E_max)][:, 1:]
    err_array_squared = np.power(err_array, 2)
    return np.power(err_array_squared.mean(axis=0), 0.5)


def get_frozen_window_min_max(wannier90winfile="wannier90.win"):
    """Get the frozen window min and max from the wannier90.win file.

    Args:
        wannier90winfile (str, optional): Path to the wannier90.win file. Defaults to 'wannier90.win'.

    Returns:
        float: Frozen window min.
        float: Frozen window max.
    """
    dis_froz_min = None
    dis_froz_max = None

    with open(wannier90winfile, "r") as fr:
        for line in fr:
            if "dis_froz_min" in line:
                dis_froz_min = float(
                    line.split("dis_froz_min")[1].split("=")[1].split(" ")[0]
                )
            if "dis_froz_max" in line:
                dis_froz_max = float(
                    line.split("dis_froz_max")[1].split("=")[1].split(" ")[0]
                )
    return dis_froz_min, dis_froz_max


def wannier_quality_calculation(
    kpoint_matrix,
    NK,
    kpath_ticks,
    Fermi_nsc_wann,
    num_wann,
    discard_first_bands=0,
    sc_dir="0_self-consistent",
    nsc_dir="1_non-self-consistent",
    wann_dir="2_wannier",
    bands_dir="1_band_structure",
    tb_model_dir="2_wannier/tb_model_wann90",
    band_for_Fermi_correction=None,
    kpoint_for_Fermi_correction="0.0000000E+00  0.0000000E+00  0.0000000E+00",
    yaxis_lim=None,
    savefig=True,
    showfig=True,
):
    """Calculate the quality of the Wannierization.

    Needed files from VASP:

        = nsc_calculation_path:
            - EIGENVAL
            - DOSCAR

        = dft_bands_folder:
            - EIGENVAL
            - "Sxyz_exp_values_from_spn_file.dat" (get automatically from wannier90.spn_formatted!!)
            - OUTCAR

    Args:
        kpoint_matrix (np.array): K-point matrix.
        NK (int): Number of k-points.
        kpath_ticks (list): K-path ticks.
        Fermi_nsc_wann (float): Fermi energy in the non-self-consistent calculation.
        num_wann (int): Number of Wannierized bands.
        discard_first_bands (int, optional): Number of bands to discard. Defaults to 0.
        sc_dir (str, optional): Path to the self-consistent calculation directory. Defaults to '0_self-consistent'.
        nsc_dir (str, optional): Path to the non-self-consistent calculation directory. Defaults to '1_non-self-consistent'.
        wann_dir (str, optional): Path to the Wannier calculation directory. Defaults to '2_wannier'.
        bands_dir (str, optional): Path to the bands calculation directory. Defaults to '1_band_structure'.
        tb_model_dir (str, optional): Path to the wannier90.tb_model directory. Defaults to '2_wannier/tb_model_wann90'.
        band_for_Fermi_correction (int, optional): Band for Fermi correction. Defaults to None.
        kpoint_for_Fermi_correction (str, optional): K-point for Fermi correction. Defaults to '0.0000000E+00  0.0000000E+00  0.0000000E+00'.
        yaxis_lim (list, optional): Y-axis limits. Defaults to None.
        savefig (bool, optional): Save the figure. Defaults to True.
        showfig (bool, optional): Show the figure. Defaults to True.

    Returns:
        np.array: Error by energy.
    """

    # make sure directories don't have / at the end
    for dir_name in [sc_dir, nsc_dir, wann_dir, bands_dir, tb_model_dir]:
        if dir_name[-1] == "/":
            dir_name = dir_name[:-1]

    if (
        band_for_Fermi_correction is None
    ):  # if not specified, take the first non-discarded band
        band_for_Fermi_correction = discard_first_bands + 1

    # ================================================== CONSTANTS ==================================================================

    labels = ["G", "K", "M", "G"]

    S_DFT_fname = "Sxyz_exp_values_from_spn_file.dat"
    hr_R_name = tb_model_dir + "/hr_R_dict.pickle"  # "hr_R_dict_sym.pickle"
    spn_R_name = tb_model_dir + "/spn_R_dict.pickle"

    deltaE_around_EF = (
        0.5  # eV; integrate error in this +- window around E_F for plotting
    )
    deltaE2_around_EF = (
        0.1  # eV; integrate error in this +- window around E_F for plotting
    )

    # =============================================================================================================================

    # ======================== GET HOME-MADE SPIN TEXTURE ================================
    A = load_lattice_vectors(win_file=f"{wann_dir}/wannier90.win")
    G = reciprocal_lattice_vectors(A)
    kpoints, kpoints_cart, kpath = get_kpoint_path(kpoint_matrix, G, Nk=NK)

    # interpolate Hamiltonian
    with open(hr_R_name, "rb") as fr:
        hr_R_dict = pickle.load(fr)

    H_k_W = real_to_W_gauge(kpoints, hr_R_dict)
    Eigs_k, U_mn_k = W_gauge_to_H_gauge(H_k_W, U_mn_k={}, hamiltonian=True)

    shutil.copyfile(
        nsc_dir + "/FERMI_ENERGY_corrected.in", "./FERMI_ENERGY_corrected.in"
    )
    E_F = Fermi_nsc_wann
    # take all elements from Eigs_k and subtract Fermi_nsc_wann
    for key in Eigs_k.keys():
        Eigs_k[key] -= Fermi_nsc_wann

    # interpolate spin operator
    with open(spn_R_name, "rb") as fr:
        spn_R_dict = pickle.load(fr)

    S_mn_R_x, S_mn_R_y, S_mn_R_z = split_spn_dict(
        spn_R_dict, spin_names=["x", "y", "z"]
    )

    Sx_k_W = real_to_W_gauge(kpoints, S_mn_R_x)
    Sy_k_W = real_to_W_gauge(kpoints, S_mn_R_y)
    Sz_k_W = real_to_W_gauge(kpoints, S_mn_R_z)

    # print('len kpoints', len(kpoints))
    # print('len H_k_W', len(list(H_k_W.keys())))

    S_mn_k_H_x = W_gauge_to_H_gauge(Sx_k_W, U_mn_k=U_mn_k, hamiltonian=False)
    S_mn_k_H_y = W_gauge_to_H_gauge(Sy_k_W, U_mn_k=U_mn_k, hamiltonian=False)
    S_mn_k_H_z = W_gauge_to_H_gauge(Sz_k_W, U_mn_k=U_mn_k, hamiltonian=False)

    # =====================================================================================

    # COMPARE with DFT
    dft_kpoints, dft_bands, num_kpoints_dft, num_bands = parse_eigenval_file(
        bands_dir + "/EIGENVAL"
    )
    dft_bands = dft_bands[:, discard_first_bands : discard_first_bands + num_wann]

    # get Fermi for the bands non-self-consistent calculation
    Fermi_nsc_bands = get_fermi_corrected_by_matching_bands(
        path=".",
        nsc_calculation_path=bands_dir,
        corrected_at_kpoint=kpoint_for_Fermi_correction,
        corrected_at_band=band_for_Fermi_correction,
        sc_calculation_path=sc_dir,
        fout_name="FERMI_ENERGY_corrected.in",
    )
    dft_bands -= Fermi_nsc_bands

    # make 2D again: (NKpoints, num_wann)
    E_to_compare = np.array([Eigs_k[key] for key in Eigs_k.keys()])
    # print("E_to_compare_shape", E_to_compare.shape)
    E_to_compare_with_duplicates = duplicate_kpoints_for_home_made(E_to_compare, NK)
    # print("E_to_compare_with_duplicates_shape", E_to_compare_with_duplicates.shape)

    # print("50", E_to_compare_with_duplicates[50,:])
    # print("51", E_to_compare_with_duplicates[51,:])

    S_mn_k_H_x_to_compare = duplicate_kpoints_for_home_made(
        np.array([np.diag(S_mn_k_H_x[key]) for key in S_mn_k_H_x.keys()]), NK
    )
    S_mn_k_H_y_to_compare = duplicate_kpoints_for_home_made(
        np.array([np.diag(S_mn_k_H_y[key]) for key in S_mn_k_H_y.keys()]), NK
    )
    S_mn_k_H_z_to_compare = duplicate_kpoints_for_home_made(
        np.array([np.diag(S_mn_k_H_z[key]) for key in S_mn_k_H_z.keys()]), NK
    )

    # print('k-points from dict keys', list(S_mn_k_H_x.keys()))

    S_to_compare_with_duplicates = np.array(
        [S_mn_k_H_x_to_compare, S_mn_k_H_y_to_compare, S_mn_k_H_z_to_compare]
    )
    S_shape = S_to_compare_with_duplicates.shape

    # make the spin axis the last instead of the first one
    S_to_compare_with_duplicates = S_to_compare_with_duplicates.swapaxes(0, 1)
    S_to_compare_with_duplicates = S_to_compare_with_duplicates.swapaxes(1, 2)

    # S_to_compare_with_duplicates = duplicate_kpoints(S_to_compare, NK)

    # error_by_energy = compare_eigs_bandstructure_at_exact_kpts(
    #     dft_bands,
    #     E_to_compare_with_duplicates,
    #     num_kpoints_dft,
    #     num_wann,
    #     f_name_out="home-made_quality_error_Fermi_corrected.dat",
    # )

    # plot_err_vs_energy(
    #     error_by_energy,
    #     Ef=0,
    #     title="Wannierization RMS error vs. energy",
    #     fig_name_out="wannier_quality_error_by_energy_home-made_Fermi_corrected.png",
    #     savefig=savefig,
    #     showfig=showfig,
    # )

    # ------------- COMPARE spin texture --------------------

    # load spin expectation values and select relevant bands
    S_DFT = np.loadtxt(f"{bands_dir}/{S_DFT_fname}")
    # select relevant bands
    NK = get_NKpoints(OUTCAR=f"{bands_dir}/OUTCAR")
    S_DFT = S_DFT.reshape(NK, -1, 3)
    S_DFT_to_compare = S_DFT[:, discard_first_bands : discard_first_bands + num_wann, :]

    # print("S_to_compare dimensions", S_to_compare_with_duplicates.shape)
    # print("S_to_compare", S_to_compare)

    # print("50", S_to_compare_with_duplicates[50,:,:])
    # print("51", S_to_compare_with_duplicates[51,:,:])

    S_diff = np.abs(
        S_DFT_to_compare.reshape(-1, 3) - S_to_compare_with_duplicates.reshape(-1, 3)
    )
    E_diff = np.abs(dft_bands.reshape(-1) - E_to_compare_with_duplicates.reshape(-1))

    # print('!!!!!!!!!!!Fermi_nsc_wann', Fermi_nsc_wann)

    # plot error-colored band structure
    plot_err_vs_bands(
        kpoints,
        kpath,
        kpath_ticks,
        Eigs_k,
        E_diff,
        S_diff,
        fout="ERRORS_ALL_band_structure_home-made_Fermi_corrected.jpg",
        yaxis_lim=yaxis_lim,
        savefig=savefig,
        showfig=showfig,
    )

    # E_F was already subtracted
    E = dft_bands.reshape(-1)

    # make final matrix
    error_E_S_by_energy = np.vstack(
        [E, E_diff, S_diff[:, 0], S_diff[:, 1], S_diff[:, 2]]
    ).T

    # print("Error matrix shape:", error_E_S_by_energy.shape)

    # order the '' matrix by energy (0th column)
    error_E_S_by_energy = error_E_S_by_energy[error_E_S_by_energy[:, 0].argsort()]

    np.savetxt(
        "home-made_quality_error_S_E_Fermi_corrected.dat",
        error_E_S_by_energy,
        header="E (eV)\t|Delta E| (eV)\t|Delta S_x|\t|Delta S_y|\t|Delta S_z|",
    )

    # plot all the errors
    fig, axes = plt.subplots(2, 2, figsize=[6, 4])

    axes[0, 0].semilogy(
        error_E_S_by_energy[:, 0], error_E_S_by_energy[:, 1], "ko", markersize=2
    )
    axes[0, 0].set_title(r"$E$", fontsize=14)
    axes[0, 0].set_ylabel(r"|$E_\mathrm{DFT} - E_\mathrm{wann}$| (eV)")

    axes[0, 1].semilogy(
        error_E_S_by_energy[:, 0], error_E_S_by_energy[:, 4], "bo", markersize=2
    )
    axes[0, 1].set_title(r"$S_z$", fontsize=14)
    axes[0, 1].set_ylabel(r"|$S_{z, \mathrm{DFT}} - S_{z, \mathrm{wann}}$|")

    axes[1, 0].semilogy(
        error_E_S_by_energy[:, 0], error_E_S_by_energy[:, 2], "ro", markersize=2
    )
    axes[1, 0].set_title(r"$S_x$", fontsize=14)
    axes[1, 0].set_ylabel(r"|$S_{x, \mathrm{DFT}} - S_{x, \mathrm{wann}}$|")
    axes[1, 0].set_xlabel(r"$E - E_\mathrm{F}$ (eV)")

    axes[1, 1].semilogy(
        error_E_S_by_energy[:, 0], error_E_S_by_energy[:, 3], "go", markersize=2
    )
    axes[1, 1].set_title(r"$S_y$", fontsize=14)
    axes[1, 1].set_ylabel(r"|$S_{y, \mathrm{DFT}} - S_{y, \mathrm{wann}}$|")
    axes[1, 1].set_xlabel(r"$E - E_\mathrm{F}$ (eV)")

    fig.suptitle("Error of Wannier interpolation", fontsize=14)
    plt.tight_layout()

    # apply ylim
    if yaxis_lim:
        for ax in axes.flatten():
            ax.set_xlim(yaxis_lim)
    if savefig is True:
        plt.savefig("ERRORS_all_home-made_Fermi_corrected.png", dpi=400)
    if showfig is True:
        plt.show()
    plt.close()

    dis_froz_min, dis_froz_max = get_frozen_window_min_max(
        wannier90winfile=f"{wann_dir}/wannier90.win"
    )

    Fermi_sc = float(np.loadtxt(f"{sc_dir}/FERMI_ENERGY.in"))
    mean_error_whole_range = integrate_error(error_E_S_by_energy, E_min=-1e6, E_max=1e6)
    mean_error_up_to_Fermi = integrate_error(error_E_S_by_energy, E_min=-1e6, E_max=0)
    mean_error_frozen_window = integrate_error(
        error_E_S_by_energy,
        E_min=dis_froz_min - Fermi_nsc_wann,
        E_max=dis_froz_max - Fermi_nsc_wann,
    )
    mean_error_around_Fermi = integrate_error(
        error_E_S_by_energy, E_min=-deltaE_around_EF, E_max=deltaE_around_EF
    )
    mean_error_around_Fermi2 = integrate_error(
        error_E_S_by_energy, E_min=-deltaE2_around_EF, E_max=deltaE_around_EF
    )

    with open("error_home-made_integrated_Fermi_corrected.dat", "w") as fw:
        fw.write(
            "#                                           \tRMSE_E (eV)\tRMSE_Sx  \tRMSE_Sy  \tRMSE_Sz\n"
        )
        fw.write(
            f"in the whole energy range                \t"
            + "\t".join([f"{val:.6f}" for val in mean_error_whole_range])
            + "\n"
        )
        fw.write(
            f"up to Fermi                                 \t"
            + "\t".join([f"{val:.6f}" for val in mean_error_up_to_Fermi])
            + "\n"
        )
        fw.write(
            f"in the frozen window ({dis_froz_min-Fermi_nsc_wann:.3f} to {dis_froz_max-Fermi_sc:.3f} eV)\t"
            + "\t".join([f"{val:.6f}" for val in mean_error_frozen_window])
            + "\n"
        )
        fw.write(
            f"window +- {deltaE_around_EF:.3f} eV around Fermi level    \t"
            + "\t".join([f"{val:.6f}" for val in mean_error_around_Fermi])
            + "\n"
        )
        fw.write(
            f"window +- {deltaE2_around_EF:.3f} eV around Fermi level    \t"
            + "\t".join([f"{val:.6f}" for val in mean_error_around_Fermi2])
        )

    # ------------ ABSOLUTE VALUE OF SPIN ------------------

    S_abs = np.linalg.norm(S_to_compare_with_duplicates.reshape(-1, 3), axis=1)
    # print("S_abs shape", S_abs.shape)

    # combined matrix
    E_S_abs = np.vstack([E, S_abs]).T
    # print("S_abs_E shape", E_S_abs.shape)

    # order the '' matrix by energy (0th column)
    E_S_abs = E_S_abs[E_S_abs[:, 0].argsort()]

    # save txt file, with statistics
    Sabs_mean = np.mean(S_abs)
    Sabs_median = np.median(S_abs)
    Sabs_over_one = len(S_abs[S_abs > 1.0]) / len(S_abs)
    header = f"S over 1 ratio = {Sabs_over_one:.8f}\nS mean = {Sabs_mean:.8f}\nS median = {Sabs_median:.8f}\nE (eV)\t|S|"
    np.savetxt("home-made_Sabs_vs_E_Fermi_corrected.dat", E_S_abs, header=header)

    # plot and save S_abs vs. E
    plot_title = "S_abs"

    fig, ax = plt.subplots(1, 1, figsize=[3.7, 2.5])
    ax.plot(E_S_abs[:, 0], E_S_abs[:, 1], "ko", markersize=2)
    ax.set_title("Interpolated spin magnitudes", fontsize=13)
    ax.set_ylabel(r"|$S$|", fontsize=13)
    ax.set_xlabel(r"$E - E_\mathrm{F}$ (eV)", fontsize=13)
    plt.tight_layout()
    if savefig is True:
        plt.savefig(plot_title + "_vs_E_home-made_Fermi_corrected.png", dpi=400)
    if showfig is True:
        plt.show()
    plt.close()

    # plot and save S_abs histogram
    fig, ax = plt.subplots(1, 1, figsize=[3.5, 2.5])
    plt.hist(S_abs.flatten(), bins=100)
    plt.xlabel("|S|", fontsize=13)
    plt.ylabel("counts", fontsize=13)
    plt.title(
        "Spin magnitude histogram", #f"histogram of abs values of diagonal elements of spin operator\n- in k-space (eigenvalues) home-made interpolation",
        fontsize=13,
    )
    plt.tight_layout()
    if savefig is True:
        plt.savefig(plot_title + "_S_histogram_home-made_Fermi_corrected.png", dpi=400)
    if showfig is True:
        plt.show()
    plt.close()
