# Copyright (c) [2024] []
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import warnings
from argparse import ArgumentParser
from sys import getsizeof
from timeit import default_timer as timer

# use numpy number of threads one
try:
    from threadpoolctl import threadpool_info, threadpool_limits

    user_api = threadpool_info()["user_api"]
    threadpool_limits(limits=1, user_api=user_api)
except:
    print("Warning: threadpoolctl could not make numpy use single thread!")

from pickle import dump, load

import numpy as np
import sisl
from mpi4py import MPI
from numpy.linalg import inv
from scipy.special import roots_legendre

try:
    from tqdm import tqdm

    tqdm_imported = True
except:
    print("Please install tqdm for nice progress bar.")
    tqdm_imported = False


# Pauli matrices
TAU_X = np.array([[0, 1], [1, 0]])
TAU_Y = np.array([[0, -1j], [1j, 0]])
TAU_Z = np.array([[1, 0], [0, -1]])
TAU_0 = np.array([[1, 0], [0, 1]])


# list of accepted input parameters
ACCEPTED_INPUTS = [
    "infile",
    "outfile",
    "scf_xcf_orientation",
    "ref_xcf_orientations",
    "kset",
    "kdirs",
    "ebot",
    "eset",
    "esetp",
    "parallel_solver_for_Gk",
    "padawan_mode",
]

DEFAULT_ARGUMENTS = dict(
    infile=None,
    outfile=None,
    scf_xcf_orientation=np.array([0, 0, 1]),
    ref_xcf_orientations=[
        dict(o=np.array([1, 0, 0]), vw=[np.array([0, 1, 0]), np.array([0, 0, 1])]),
        dict(o=np.array([0, 1, 0]), vw=[np.array([1, 0, 0]), np.array([0, 0, 1])]),
        dict(o=np.array([0, 0, 1]), vw=[np.array([1, 0, 0]), np.array([0, 1, 0])]),
    ],
    kset=2,
    kdirs="xyz",
    ebot=None,
    eset=42,
    esetp=1000,
    parallel_solver_for_Gk=False,
    padawan_mode=True,
)


def commutator(a, b):
    """Shorthand for commutator.

    Commutator of two matrices in the mathematical sense.

    Args:
        a : np.array_like
            The first matrix
        b : np.array_like
            The second matrix

    Returns:
        np.array_like
            The commutator of a and b
    """

    return a @ b - b @ a


# define some useful functions
def hsk(H, ss, sc_off, k=(0, 0, 0)):
    """Speed up Hk and Sk generation.

    Calculates the Hamiltonian and the Overlap matrix at a given k point. It is faster that the sisl version.

    Args:
        H : np.array_like
            Hamiltonian in spin box form
        ss : np.array_like
            Overlap matrix in spin box form
        sc_off : list
            supercell indexes of the Hamiltonian
        k : tuple, optional
            The k point where the matrices are set up. Defaults to (0, 0, 0)

    Returns:
        np.array_like
            Hamiltonian at the given k point
        np.array_like
            Overlap matrix at the given k point
    """

    # this two conversion lines are from the sisl source
    k = np.asarray(k, np.float64)
    k.shape = (-1,)

    # this generates the list of phases
    phases = np.exp(-1j * 2 * np.pi * k @ sc_off.T)

    # phases applied to the hamiltonian
    HK = np.einsum("abc,a->bc", H, phases)
    SK = np.einsum("abc,a->bc", ss, phases)

    return HK, SK


def make_kset(dirs="xyz", NUMK=20):
    """Simple k-grid generator to sample the Brillouin zone.

    Depending on the value of the dirs
    argument k sampling in 1,2 or 3 dimensions is generated.
    If dirs argument does not contain either of x, y or z
    a kset of a single k-pont at the origin is returned. The total number of k points is the NUMK**(dimensions)


    Args:
        dirs : str, optional
            Directions of the k points in the Brillouin zone. They are the three lattice vectors. Defaults to "xyz"
        NUMK : int, optional
            The number of k points in a direction. Defaults to 20

    Returns:
        np.array_like
            An array of k points that uniformly sample the Brillouin zone in the given directions
    """

    # if there is no xyz in dirs return the Gamma point
    if not (sum([d in dirs for d in "xyz"])):
        return np.array([[0, 0, 0]])

    kran = len(dirs) * [np.linspace(0, 1, NUMK, endpoint=False)]
    mg = np.meshgrid(*kran)
    dirsdict = dict()

    for d in enumerate(dirs):
        dirsdict[d[1]] = mg[d[0]].flatten()
    for d in "xyz":
        if not (d in dirs):
            dirsdict[d] = 0 * dirsdict[dirs[0]]
    kset = np.array([dirsdict[d] for d in "xyz"]).T

    return kset


def make_contour(emin=-20, emax=0.0, enum=42, p=150):
    """A more sophisticated contour generator.

    Calculates the parameters for the complex contour integral. It uses the
    Legendre-Gauss quadrature method.   It returns a class that contains
    the information for the contour integral.

    Args:
        emin : int, optional
            Energy minimum of the contour. Defaults to -20
        emax : float, optional
            Energy maximum of the contour. Defaults to 0.0, so the Fermi level
        enum : int, optional
            Number of sample points along the contour. Defaults to 42
        p : int, optional
            Shape parameter that describes the distribution of the sample points. Defaults to 150

    Returns:
        ccont
            Contains all the information for the contour integral
    """
    x, wl = roots_legendre(enum)
    R = (emax - emin) / 2  # radius
    z0 = (emax + emin) / 2  # center point
    y1 = -np.log(1 + np.pi * p)  # lower bound
    y2 = 0  # upper bound

    y = (y2 - y1) / 2 * x + (y2 + y1) / 2
    phi = (np.exp(-y) - 1) / p  # angle parameter
    ze = z0 + R * np.exp(1j * phi)  # complex points for path
    we = -(y2 - y1) / 2 * np.exp(-y) / p * 1j * (ze - z0) * wl

    # just an empty container class
    class ccont:
        pass

    cont = ccont()
    cont.R = R
    cont.z0 = z0
    cont.ze = ze
    cont.we = we
    cont.enum = enum

    return cont


def tau_u(u):
    """Pauli matrix in direction u.

    Returns the vector u in the basis of the Pauli matrices.

    Args:
        u : list or np.array_like
            The direction

    Returns:
        np.array_like
            Arbitrary direction in the base of the Pauli matrices
    """

    # u is force to be of unit length
    u = u / np.linalg.norm(u)

    return u[0] * TAU_X + u[1] * TAU_Y + u[2] * TAU_Z


#
def crossM(u):
    """Definition for the cross-product matrix.

    It acts as a cross product with vector u.

    Args:
        u : list or np.array_like
            The second vector in the cross product

    Returns:
        np.array_like
            The matrix that represents teh cross product with a vector
    """
    return np.array([[0, -u[2], u[1]], [u[2], 0, -u[0]], [-u[1], u[0], 0]])


def RotM(theta, u, eps=1e-10):
    """Definition of rotation matrix with angle theta around direction u.


    Args:
        theta : float
            The angle of rotation
        u : np.array_like
            The rotation axis
        eps : float, optional
            Cutoff for small elements in the resulting matrix. Defaults to 1e-10

    Returns:
        np.array_like
            The rotation matrix
    """

    u = u / np.linalg.norm(u)

    M = (
        np.cos(theta) * np.eye(3)
        + np.sin(theta) * crossM(u)
        + (1 - np.cos(theta)) * np.outer(u, u)
    )

    # kill off small numbers
    M[abs(M) < eps] = 0.0
    return M


def RotMa2b(a, b, eps=1e-10):
    """Definition of rotation matrix rotating unit vector a to unit vector b.

    Function returns array R such that R @ a = b holds.

    Args:
        a : np.array_like
            First vector
        b : np.array_like
            Second vector
        eps : float, optional
            Cutoff for small elements in the resulting matrix. Defaults to 1e-10

    Returns:
        np.array_like
            The rotation matrix with the above property
    """

    v = np.cross(a, b)
    c = a @ b
    M = np.eye(3) + crossM(v) + crossM(v) @ crossM(v) / (1 + c)

    # kill off small numbers
    M[abs(M) < eps] = 0.0
    return M


def read_siesta_emin(eigfile):
    """It reads the lowest energy level from the siesta run.

    It uses the .EIG file from siesta that contains the eigenvalues.

    Args:
        eigfile : str
            The path to the .EIG file

    Returns:
        float
            The energy minimum
    """

    # read the file
    eigs = eigSileSiesta(eigfile).read_data()

    return eigs.min()


def int_de_ke(traced, we):
    """It numerically integrates the traced matrix.

    It is a wrapper from numpy.trapz and it contains the
    relevant constants to calculate the energy integral from
    equation 93 or 96.

    Args:
        traced : np.array_like
            The trace of a matrix or a matrix product
        we : float
            The weight of a point on the contour

    Returns:
        float
            The energy calculated from the integral formula
    """

    return np.trapz(-1 / np.pi * np.imag(traced * we))


def onsite_projection(matrix, idx1, idx2):
    """It produces the slices of a matrix for the on site projection.

    The slicing is along the last two axes as these contains the orbital indexing.

    Args:
        matrix : (..., :, :) np.array_like
            Some matrix
        idx : np.array_like
            The indexes of the orbitals

    Returns:
        np.array_like
            Reduced matrix based on the projection
    """

    return matrix[..., idx1, :][..., idx2]


def calc_Vu(H, Tu):
    """Calculates the local perturbation in case of a spin rotation.

    Args:
        H : (NO, NO) np.array_like
            Hamiltonian
        Tu : (NO, NO) array_like
            Rotation around u

    Returns:
        Vu1 : (NO, NO) np.array_like
            First order perturbed matrix
        Vu2 : (NO, NO) np.array_like
            Second order perturbed matrix
    """

    Vu1 = 1j / 2 * commutator(H, Tu)  # equation 100
    Vu2 = 1 / 8 * commutator(commutator(Tu, H), Tu)  # equation 100

    return Vu1, Vu2


def build_hh_ss(dh):
    """It builds the Hamiltonian and Overlap matrix from the sisl.dh class.

    It restructures the data in the SPIN BOX representation, where NS is
    the number of supercells and NO is the number of orbitals.

    Args:
        dh : sisl.physics.Hamiltonian
            Hamiltonian read in by sisl

    Returns:
        hh : (NS, NO, NO) np.array_like
            Hamiltonian in SPIN BOX representation
        ss : (NS, NO, NO) np.array_like
            Overlap matrix in SPIN BOX representation
    """

    NO = dh.no  # shorthand for number of orbitals in the unit cell

    # preprocessing Hamiltonian and overlap matrix elements
    h11 = dh.tocsr(dh.M11r)
    h11 += dh.tocsr(dh.M11i) * 1.0j
    h11 = h11.toarray().reshape(NO, dh.n_s, NO).transpose(0, 2, 1).astype("complex128")

    h22 = dh.tocsr(dh.M22r)
    h22 += dh.tocsr(dh.M22i) * 1.0j
    h22 = h22.toarray().reshape(NO, dh.n_s, NO).transpose(0, 2, 1).astype("complex128")

    h12 = dh.tocsr(dh.M12r)
    h12 += dh.tocsr(dh.M12i) * 1.0j
    h12 = h12.toarray().reshape(NO, dh.n_s, NO).transpose(0, 2, 1).astype("complex128")

    h21 = dh.tocsr(dh.M21r)
    h21 += dh.tocsr(dh.M21i) * 1.0j
    h21 = h21.toarray().reshape(NO, dh.n_s, NO).transpose(0, 2, 1).astype("complex128")

    sov = (
        dh.tocsr(dh.S_idx)
        .toarray()
        .reshape(NO, dh.n_s, NO)
        .transpose(0, 2, 1)
        .astype("complex128")
    )

    # Reorganization of Hamiltonian and overlap matrix elements to SPIN BOX representation
    U = np.vstack(
        [np.kron(np.eye(NO, dtype=int), [1, 0]), np.kron(np.eye(NO, dtype=int), [0, 1])]
    )
    # This is the permutation that transforms ud1ud2 to u12d12
    # That is this transforms FROM SPIN BOX to ORBITAL BOX => U
    # the inverse transformation is U.T u12d12 to ud1ud2
    # That is FROM ORBITAL BOX to SPIN BOX => U.T

    # From now on everything is in SPIN BOX!!
    hh, ss = np.array(
        [
            U.T
            @ np.block([[h11[:, :, i], h12[:, :, i]], [h21[:, :, i], h22[:, :, i]]])
            @ U
            for i in range(dh.lattice.nsc.prod())
        ]
    ), np.array(
        [
            U.T
            @ np.block(
                [[sov[:, :, i], sov[:, :, i] * 0], [sov[:, :, i] * 0, sov[:, :, i]]]
            )
            @ U
            for i in range(dh.lattice.nsc.prod())
        ]
    )

    return hh, ss


def setup_pairs_and_magnetic_entities(
    magnetic_entities, pairs, dh, simulation_parameters
):
    """It creates the complete structure of the dictionaries and fills some basic data.

    It creates orbital indexes, spin box indexes, coordinates and tags for magnetic entities.
    Furthermore it creates the structures for the energies, the perturbed potentials and
    the Greens function calculation. It dose the same for the pairs.

    Args:
        pairs : dict
            Contains the initial pair information
        magnetic_entities : dict
            Contains the initial magnetic entity information
        dh : sisl.physics.Hamiltonian
            Hamiltonian read in by sisl
        simulation_parameters : dict
            A set of parameters from the simulation

    Returns:
        pairs : dict
            Contains the initial information and the complete structure
        magnetic_entities : dict
            Contains the initial information and the complete structure
    """

    # for every site we have to store 3 Greens function (and the associated _tmp-s) in the 3 reference directions
    for mag_ent in magnetic_entities:
        parsed = parse_magnetic_entity(dh, **mag_ent)  # parse orbital indexes
        mag_ent["orbital_indices"] = parsed
        mag_ent["spin_box_indices"] = blow_up_orbindx(
            parsed
        )  # calculate spin box indexes
        # if orbital is not set use all
        if "l" not in mag_ent.keys():
            mag_ent["l"] = "all"

        # tag creation for one atom
        if isinstance(mag_ent["atom"], int):
            mag_ent["tags"] = [
                f"[{mag_ent['atom']}]{dh.atoms[mag_ent['atom']].tag}({mag_ent['l']})"
            ]
            mag_ent["xyz"] = [dh.xyz[mag_ent["atom"]]]
        # tag creation for more atoms
        if isinstance(mag_ent["atom"], list):
            mag_ent["tags"] = []
            mag_ent["xyz"] = []
            # iterate over atoms
            for atom_idx in mag_ent["atom"]:
                mag_ent["tags"].append(
                    f"[{atom_idx}]{dh.atoms[atom_idx].tag}({mag_ent['l']})"
                )
                mag_ent["xyz"].append(dh.xyz[atom_idx])

        # calculate size for Greens function generation
        spin_box_shape = len(mag_ent["spin_box_indices"])

        # we will store the second order energy derivations here
        mag_ent["energies"] = []

        # These will be the perturbed potentials from eq. 100
        mag_ent["Vu1"] = []  # so they are independent in memory
        mag_ent["Vu2"] = []

        mag_ent["Gii"] = []  # Greens function
        mag_ent["Gii_tmp"] = []  # Greens function for parallelization
        for _ in simulation_parameters["ref_xcf_orientations"]:
            # Rotations for every quantization axis
            mag_ent["Vu1"].append([])
            mag_ent["Vu2"].append([])
            # Greens functions for every quantization axis
            mag_ent["Gii"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape, spin_box_shape),
                    dtype="complex128",
                )
            )
            mag_ent["Gii_tmp"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape, spin_box_shape),
                    dtype="complex128",
                )
            )

    # for every site we have to store 2x3 Greens function (and the associated _tmp-s)
    # in the 3 reference directions, because G_ij and G_ji are both needed
    for pair in pairs:
        # calculate distance
        xyz_ai = magnetic_entities[pair["ai"]]["xyz"]
        xyz_aj = magnetic_entities[pair["aj"]]["xyz"]
        xyz_aj = xyz_aj + pair["Ruc"] @ simulation_parameters["cell"]
        pair["dist"] = np.linalg.norm(xyz_ai - xyz_aj)

        # calculate size for Greens function generation
        spin_box_shape_i = len(magnetic_entities[pair["ai"]]["spin_box_indices"])
        spin_box_shape_j = len(magnetic_entities[pair["aj"]]["spin_box_indices"])
        # tag generation
        pair["tags"] = []
        for mag_ent in [magnetic_entities[pair["ai"]], magnetic_entities[pair["aj"]]]:
            tag = ""
            # get atoms of magnetic entity
            atoms_idx = mag_ent["atom"]
            orbitals = mag_ent["l"]

            # if magnetic entity contains one atoms
            if isinstance(atoms_idx, int):
                tag += f"[{atoms_idx}]{dh.atoms[atoms_idx].tag}({orbitals})"

            # if magnetic entity contains more than one atoms
            if isinstance(atoms_idx, list):
                # iterate over atoms
                atom_group = "{"
                for atom_idx in atoms_idx:
                    atom_group += f"[{atom_idx}]{dh.atoms[atom_idx].tag}({orbitals})--"
                # end {} of the atoms in the magnetic entity
                tag += atom_group[:-2] + "}"
            pair["tags"].append(tag)
        pair["energies"] = []  # we will store the second order energy derivations here

        pair["Gij"] = []  # Greens function
        pair["Gji"] = []
        pair["Gij_tmp"] = []  # Greens function for parallelization
        pair["Gji_tmp"] = []
        for _ in simulation_parameters["ref_xcf_orientations"]:
            # Greens functions for every quantization axis
            pair["Gij"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape_i, spin_box_shape_j),
                    dtype="complex128",
                )
            )
            pair["Gij_tmp"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape_i, spin_box_shape_j),
                    dtype="complex128",
                )
            )
            pair["Gji"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape_j, spin_box_shape_i),
                    dtype="complex128",
                )
            )
            pair["Gji_tmp"].append(
                np.zeros(
                    (simulation_parameters["eset"], spin_box_shape_j, spin_box_shape_i),
                    dtype="complex128",
                )
            )

    return pairs, magnetic_entities


def parallel_Gk(HK, SK, eran, eset):
    """Calculates the Greens function by inversion.

    It calculates the Greens function on all the energy levels at the same time.

    Args:
        HK : (NO, NO), np.array_like
            Hamiltonian at a given k point
        SK : (NO, NO), np.array_like
            Overlap Matrix at a given k point
        eran : (eset) np.array_like
            Energy sample along the contour
        eset : int
            Number of energy samples along the contour

    Returns:
        Gk : (eset, NO, NO), np.array_like
            Green's function at a given k point
    """

    # Calculates the Greens function on all the energy levels
    return inv(SK * eran.reshape(eset, 1, 1) - HK)


def sequential_GK(HK, SK, eran, eset):
    """Calculates the Greens function by inversion.

    It calculates sequentially over the energy levels.

    Args:
        HK : (NO, NO), np.array_like
            Hamiltonian at a given k point
        SK : (NO, NO), np.array_like
            Overlap Matrix at a given k point
        eran : (eset) np.array_like
            Energy sample along the contour
        eset : int
            Number of energy samples along the contour

    Returns:
        Gk : (eset, NO, NO), np.array_like
            Green's function at a given k point
    """

    # creates an empty holder
    Gk = np.zeros(shape=(eset, HK.shape[0], HK.shape[1]), dtype="complex128")
    # fills the holder sequentially by the Greens function on a given energy
    for j in range(eset):
        Gk[j] = inv(SK * eran[j] - HK)

    return Gk


def remove_clutter_for_save(pairs, magnetic_entities):
    """Removes unimportant data from the dictionaries.

    It is used before saving to throw away data that
    is not needed for post processing.

    Args:
        pairs : dict
            Contains all the pair information
        magnetic_entities : dict
            Contains all the magnetic entity information

    Returns:
        pairs : dict
            Contains all the reduced pair information
        magnetic_entities : dict
            Contains all the reduced magnetic entity information
    """

    # remove clutter from magnetic entities and pair information
    for pair in pairs:
        del pair["Gij"]
        del pair["Gij_tmp"]
        del pair["Gji"]
        del pair["Gji_tmp"]
    for mag_ent in magnetic_entities:
        del mag_ent["Gii"]
        del mag_ent["Gii_tmp"]
        del mag_ent["Vu1"]
        del mag_ent["Vu2"]

    return pairs, magnetic_entities


def blow_up_orbindx(orb_indices):
    """Function to blow up orbital indices to make SPIN BOX indices.

    Args:
        orb_indices : np.array_like
            These are the indices in ORBITAL BOX

    Returns:
        orb_indices : np.array_like
            These are the indices in SPIN BOX
    """

    orb_indices = np.array([[2 * o, 2 * o + 1] for o in orb_indices]).flatten()

    return orb_indices


def spin_tracer(M):
    """Spin tracer utility.

    This takes an operator with the orbital-spin sequence:
    orbital 1 up,
    orbital 1 down,
    orbital 2 up,
    orbital 2 down,
    that is in the SPIN-BOX representation,
    and extracts orbital dependent Pauli traces.


    Args:
        M : np.array_like
            Traceable matrix

    Returns:
        dict
            It contains the traced matrix with "x", "y", "z" and "c"
    """
    M11 = M[0::2, 0::2]
    M12 = M[0::2, 1::2]
    M21 = M[1::2, 0::2]
    M22 = M[1::2, 1::2]

    M_o = dict()
    M_o["x"] = M12 + M21
    M_o["y"] = 1j * (M12 - M21)
    M_o["z"] = M11 - M22
    M_o["c"] = M11 + M22

    return M_o


def parse_magnetic_entity(dh, atom=None, l=None, **kwargs):
    """Function to define orbital indexes of a given magnetic entity.

    Args:
        dh : sisl.physics.Hamiltonian
            Hamiltonian from sisl
        atom : integer or list of integers, optional
            Defining atom (or atoms) in the unit cell forming the magnetic entity. Defaults to None
        l : integer, optional
            Defining the angular momentum channel. Defaults to None

    Returns:
        list
            The orbital indexes of the given magnetic entity
    """
    # case where we deal with more than one atom defining the magnetic entity
    if type(atom) == list:
        dat = []
        for a in atom:
            a_orb_idx = dh.geometry.a2o(a, all=True)
            if (
                type(l) == int
            ):  # if specified we restrict to given l angular momentum channel inside each atom
                a_orb_idx = a_orb_idx[[o.l == l for o in dh.geometry.atoms[a].orbitals]]
            dat.append(a_orb_idx)
        orbital_indeces = np.hstack(dat)
    # case where we deal with a singel atom magnetic entity
    elif type(atom) == int:
        orbital_indeces = dh.geometry.a2o(atom, all=True)
        if (
            type(l) == int
        ):  # if specified we restrict to given l angular momentum channel
            orbital_indeces = orbital_indeces[
                [o.l == l for o in dh.geometry.atoms[atom].orbitals]
            ]

    return orbital_indeces  # numpy array containing integers labeling orbitals associated to a magnetic entity.


def calculate_anisotropy_tensor(mag_ent):
    """Calculates the renormalized anisotropy tensor from the energies.

    Args:
        mag_ent : dict
            An element from the magnetic entities

    Returns:
        K : np.array_like
            elements of the anisotropy tensor
        consistency_check : float
            Absolute value of the difference from the consistency check
    """

    # get the energies
    energies = mag_ent["energies"]

    K = np.zeros((3, 3))

    # calculate the diagonal tensor elements
    K[0, 0] = energies[1, 1] - energies[1, 0]
    K[1, 1] = energies[0, 1] - energies[0, 0]
    K[2, 2] = 0

    # calculate the off-diagonal tensor elements
    K[0, 1] = (energies[2, 0] + energies[2, 1]) / 2 - energies[2, 2]
    K[1, 0] = K[0, 1]
    K[0, 2] = (energies[1, 0] + energies[1, 1]) / 2 - energies[1, 2]
    K[2, 0] = K[0, 2]
    K[1, 2] = (energies[0, 0] + energies[0, 1]) / 2 - energies[0, 2]
    K[2, 1] = K[1, 2]

    # perform consistency check
    calculated_diff = K[1, 1] - K[0, 0]
    expected_diff = energies[2, 0] - energies[2, 1]
    consistency_check = abs(calculated_diff - expected_diff)

    return K, consistency_check


def calculate_exchange_tensor(pair):
    """Calculates the exchange tensor from the energies.

    It produces the isotropic exchange, the relevant elements
    from the Dzyaloshinskii-Morilla (Dm) tensor, the symmetric-anisotropy
    and the complete exchange tensor.

    Args:
        pair : dict
            An element from the pairs

    Returns:
        J_iso : float
            Isotropic exchange (Tr[J] / 3)
        J_S : np.array_like
            Symmetric-anisotropy (J_S = J - J_iso * I ––> Jxx, Jyy, Jxy, Jxz, Jyz)
        D : np.array_like
            DM elements (Dx, Dy, Dz)
        J : np.array_like
            Complete exchange tensor flattened (Jxx, Jxy, Jxz, Jyx, Jyy, Jyz, Jzx, Jzy, Jzz)
    """

    # energies from rotations
    energies = pair["energies"]
    # Initialize output arrays
    J = np.zeros((3, 3))
    D = np.zeros(3)

    # J matrix calculations
    # J(1,1) = mean([DEij(2,2,2), DEij(2,2,3)])
    J[0, 0] = np.mean([energies[1, 3], energies[2, 3]])

    # J(1,2) = -mean([DEij(1,2,3), DEij(2,1,3)])
    J[0, 1] = -np.mean([energies[2, 1], energies[2, 2]])
    J[1, 0] = J[0, 1]

    # J(1,3) = -mean([DEij(1,2,2), DEij(2,1,2)])
    J[0, 2] = -np.mean([energies[1, 1], energies[1, 2]])
    J[2, 0] = J[0, 2]

    # J(2,2) = mean([DEij(2,2,1), DEij(1,1,3)])
    J[1, 1] = np.mean([energies[0, 3], energies[2, 0]])

    # J(2,3) = -mean([DEij(1,2,1), DEij(2,1,1)])
    J[1, 2] = -np.mean([energies[0, 1], energies[0, 2]])
    J[2, 1] = J[1, 2]

    # J(3,3) = mean([DEij(1,1,1), DEij(1,1,2)])
    J[2, 2] = np.mean([energies[0, 0], energies[1, 0]])

    # D vector calculations
    # D(1) = mean([DEij(1,2,1), -DEij(2,1,1)])
    D[0] = np.mean([energies[0, 1], -energies[0, 2]])

    # D(2) = mean([DEij(2,1,2), -DEij(1,2,2)])
    D[1] = np.mean([energies[1, 2], -energies[1, 1]])

    # D(3) = mean([DEij(1,2,3), -DEij(2,1,3)])
    D[2] = np.mean([energies[2, 1], -energies[2, 2]])

    J_iso = np.trace(J) / 3
    # based on the grogu output pdf
    # traceless symmetric exchange matrix:
    # Jxx, Jyy, Jxy, Jxz, Jyz
    J_S = np.array([J[0, 0] - J_iso, J[1, 1] - J_iso, J[0, 1], J[0, 2], J[1, 2]])

    return J_iso, J_S, D, J


def read_fdf(path):
    """It reads the simulation parameters, magnetic entities and pairs from the fdf.

    Args:
        path : string
            The path to the .fdf file

    Returns:
        fdf_arguments : dict
            The read input arguments from the fdf file
        magnetic_entities : list
            It contains the dictionaries associated with the magnetic entities
        pairs : dict
            It contains the dictionaries associated with the pair information
    """

    # read fdf file
    fdf = fdfSileSiesta(path)
    fdf_arguments = dict()

    InputFile = fdf.get("InputFile")
    if InputFile is not None:
        fdf_arguments["infile"] = InputFile

    OutputFile = fdf.get("OutputFile")
    if OutputFile is not None:
        fdf_arguments["outfile"] = OutputFile

    ScfXcfOrientation = fdf.get("ScfXcfOrientation")
    if ScfXcfOrientation is not None:
        fdf_arguments["scf_xcf_orientation"] = np.array(ScfXcfOrientation)

    XCF_Rotation = fdf.get("XCF_Rotation")
    if XCF_Rotation is not None:
        rotations = []
        # iterate over rows
        for rot in XCF_Rotation:
            # convert row to dictionary
            dat = np.array(rot.split()[:9], dtype=float)
            o = dat[:3]
            vw = dat[3:].reshape(2, 3)
            rotations.append(dict(o=o, vw=vw))
        fdf_arguments["ref_xcf_orientations"] = rotations

    Kset = fdf.get("INTEGRAL.Kset")
    if Kset is not None:
        fdf_arguments["kset"] = Kset

    Kdirs = fdf.get("INTEGRAL.Kdirs")
    if Kdirs is not None:
        fdf_arguments["kdirs"] = Kdirs

    # This is permitted because it means automatic Ebot definition
    fdf_arguments["ebot"] = fdf.get("INTEGRAL.Ebot")

    Eset = fdf.get("INTEGRAL.Eset")
    if Eset is not None:
        fdf_arguments["eset"] = Eset

    Esetp = fdf.get("INTEGRAL.Esetp")
    if Esetp is not None:
        fdf_arguments["esetp"] = Esetp

    ParallelSolver = fdf.get("GREEN.ParallelSolver")
    if ParallelSolver is not None:
        fdf_arguments["parallel_solver_for_Gk"] = ParallelSolver

    PadawanMode = fdf.get("PadawanMode")
    if PadawanMode is not None:
        fdf_arguments["padawan_mode"] = PadawanMode

    Pairs = fdf.get("Pairs")
    if Pairs is not None:
        pairs = []
        # iterate over rows
        for fdf_pair in Pairs:
            # convert data
            dat = np.array(fdf_pair.split()[:5], dtype=int)
            # create pair dictionary
            my_pair = dict(ai=dat[0], aj=dat[1], Ruc=np.array(dat[2:]))
            pairs.append(my_pair)

    MagneticEntities = fdf.get("MagneticEntities")
    if MagneticEntities is not None:
        magnetic_entities = []
        for mag_ent in MagneticEntities:
            row = mag_ent.split()
            dat = []
            for string in row:
                if string.find("#") != -1:
                    break
                dat.append(string)
            if dat[0] in {"Cluster", "cluster"}:
                magnetic_entities.append(dict(atom=[int(_) for _ in dat[1:]]))
                continue
            elif dat[0] in {"AtomShell", "Atomshell", "atomShell", "atomshell"}:
                magnetic_entities.append(
                    dict(atom=int(dat[1]), l=[int(_) for _ in dat[2:]])
                )
                continue
            elif dat[0] in {"AtomOrbital", "Atomorbital", "tomOrbital", "atomorbital"}:
                magnetic_entities.append(
                    dict(atom=int(dat[1]), orb=[int(_) for _ in dat[2:]])
                )
                continue
            elif dat[0] in {"Orbitals", "orbitals"}:
                magnetic_entities.append(dict(orb=[int(_) for _ in dat[1:]]))
                continue
            else:
                raise Exception("Unrecognizable magnetic entity in .fdf!")

    return fdf_arguments, magnetic_entities, pairs


def process_input_args(
    default_arguments,
    fdf_arguments,
    command_line_arguments,
    accepted_inputs=ACCEPTED_INPUTS,
):
    """It returns the final simulation parameters based on the inputs.

    The merging is done in the order of priority:
    1. command line arguments
    2. fdf arguments
    3. default arguments

    Args:
        default_arguments : dict
            Default arguments from grogupy
        fdf_arguments : dict
            Arguments read from the fdf input file
        command_line_arguments : dict
            Arguments from the command line

    Returns:
        dict
            The final simulation parameters
    """

    # iterate over fdf_arguments and update default arguments
    for key, value in fdf_arguments.values():
        if value is not None and key in accepted_inputs:
            default_arguments[key] = value

    # iterate over command_line_arguments and update default arguments
    for key, value in command_line_arguments.values():
        if value is not None and key in accepted_inputs:
            default_arguments[key] = value

    return default_arguments


def save_pickle(outfile, data):
    """Saves the data in the outfile with pickle.

    Args:
        outfile : str
            Path to outfile
        data : dict
            Contains the data
    """

    # save dictionary
    with open(outfile, "wb") as output_file:
        dump(data, output_file)


def load_pickle(infile):
    """Loads the data from the infile with pickle.

    Args:
        infile : str
            Path to infile

    Returns:
        data : dict
            A dictionary of data
    """

    # open and read file
    with open(infile, "wb") as input_file:
        data = load(data, input_file)

    return data


def print_job_description(simulation_parameters):
    """It prints the parameters and the description of the job.


    Args:
        simulation_parameters : dict
            It contains the simulations parameters
    """

    print(
        "================================================================================================================================================================"
    )
    print("Input file: ")
    print(simulation_parameters["infile"])
    print("Output file: ")
    print(simulation_parameters["outfile"])
    print(
        "Number of nodes in the parallel cluster: ",
        simulation_parameters["parallel_size"],
    )
    if simulation_parameters["parallel_solver_for_Gk"]:
        print("solver used for Greens function calculation: parallel")
    else:
        print("solver used for Greens function calculation: sequential")
    print(
        "================================================================================================================================================================"
    )
    print("Cell [Ang]: ")
    print(simulation_parameters["cell"])
    print(
        "================================================================================================================================================================"
    )
    print("DFT axis: ")
    print(simulation_parameters["scf_xcf_orientation"])
    print("Quantization axis and perpendicular rotation directions:")
    for ref in simulation_parameters["ref_xcf_orientations"]:
        print(ref["o"], " --» ", ref["vw"])
    print(
        "================================================================================================================================================================"
    )
    print("Parameters for the contour integral:")
    print("Number of k points: ", simulation_parameters["kset"])
    print("k point directions: ", simulation_parameters["kdirs"])
    if simulation_parameters["automatic_ebot"]:
        print(
            "Ebot: ",
            simulation_parameters["ebot"],
            "        WARNING: This was automatically determined!",
        )
    else:
        print("Ebot: ", simulation_parameters["ebot"])
    print("Eset: ", simulation_parameters["eset"])
    print("Esetp: ", simulation_parameters["esetp"])
    print(
        "================================================================================================================================================================"
    )


def print_parameters(simulation_parameters):
    """It prints the simulation parameters for the grogu out.

    Args:
        simulation_parameters : dict
            It contains the simulations parameters
    """

    print(
        "================================================================================================================================================================"
    )
    print("Input file: ")
    print(simulation_parameters["infile"])
    print("Output file: ")
    print(simulation_parameters["outfile"])
    print(
        "Number of nodes in the parallel cluster: ",
        simulation_parameters["parallel_size"],
    )
    print(
        "================================================================================================================================================================"
    )
    print("Cell [Ang]: ")
    print(simulation_parameters["cell"])
    print(
        "================================================================================================================================================================"
    )
    print("DFT axis: ")
    print(simulation_parameters["scf_xcf_orientation"])
    print("Quantization axis and perpendicular rotation directions:")
    for ref in simulation_parameters["ref_xcf_orientations"]:
        print(ref["o"], " --» ", ref["vw"])
    print(
        "================================================================================================================================================================"
    )
    print("Parameters for the contour integral:")
    print("Number of k points: ", simulation_parameters["kset"])
    print("k point directions: ", simulation_parameters["kdirs"])
    print("Ebot: ", simulation_parameters["ebot"])
    print("Eset: ", simulation_parameters["eset"])
    print("Esetp: ", simulation_parameters["esetp"])
    print(
        "================================================================================================================================================================"
    )


def print_atoms_and_pairs(magnetic_entities, pairs):
    """It prints the pair and magnetic entity information for the grogu out.

    Args:
        magnetic_entities : dict
            It contains the data on the magnetic entities
        pairs : dict
            It contains the data on the pairs
    """

    print("Atomic information: ")
    print(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    print(
        "[atom index]Element(orbitals)        x [Ang]        y [Ang]        z [Ang]        Sx        Sy        Sz        Q        Lx        Ly        Lz        Jx        Jy        Jz"
    )
    print(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    # iterate over magnetic entities
    for mag_ent in magnetic_entities:
        # iterate over atoms
        for tag, xyz in zip(mag_ent["tags"], mag_ent["xyz"]):
            # coordinates and tag
            print(f"{tag}        {xyz[0]}        {xyz[1]}        {xyz[2]}")
        print("")

    print(
        "================================================================================================================================================================"
    )
    print("Anisotropy [meV]")
    print(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    print("Magnetic entity        x [Ang]        y [Ang]        z [Ang]")
    print(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    # iterate over magnetic entities
    for mag_ent in magnetic_entities:
        # iterate over atoms
        for tag, xyz in zip(mag_ent["tags"], mag_ent["xyz"]):
            # coordinates and tag
            print(f"{tag}        {xyz[0]}        {xyz[1]}        {xyz[2]}")
            print("Consistency check: ", mag_ent["K_consistency"])
            print("K:        # Kxx, Kxy, Kxz, Kyx, Kyy, Kyz, Kzx, Kzy, Kzz")
            print(mag_ent["K"])
        print("")

    print(
        "================================================================================================================================================================"
    )
    print("Exchange [meV]")
    print(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    print("Magnetic entity1        Magnetic entity2       [i  j  k]       d [Ang]")
    print(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    # iterate over pairs
    for pair in pairs:
        # print pair parameters
        print(
            f"{pair['tags'][0]}        {pair['tags'][1]}        {pair['Ruc']}        d [Ang] {pair['dist']}"
        )
        # print magnetic parameters
        print("Isotropic: ", pair["J_iso"], "        # Tr[J] / 3")
        print("")
        print("DMI: ", pair["D"], "        # Dx, Dy, Dz")
        print("")
        print(
            "Symmetric-anisotropy: ",
            pair["J_S"],
            "        # J_S = J - J_iso * I ––> Jxx, Jyy, Jxy, Jxz, Jyz",
        )
        print("")
        print("J:        # Jxx, Jxy, Jxz, Jyx, Jyy, Jyz, Jzx, Jzy, Jzz")
        print(pair["J"])
        print(
            "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
        )

    print(
        "================================================================================================================================================================"
    )


def print_runtime_information(times):
    """It prints the runtime information for the grogu out.

    Args:
        times : dict
            It contains the runtime data
    """

    print("Runtime information: ")
    print(f"Total runtime: {times['end_time'] - times['start_time']} s")
    print(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    print(f"Initial setup: {times['setup_time'] - times['start_time']} s")
    print(
        f"Hamiltonian conversion and XC field extraction: {times['H_and_XCF_time'] - times['setup_time']:.3f} s"
    )
    print(
        f"Pair and site datastructure creatrions: {times['site_and_pair_dictionaries_time'] - times['H_and_XCF_time']:.3f} s"
    )
    print(
        f"k set cration and distribution: {times['k_set_time'] - times['site_and_pair_dictionaries_time']:.3f} s"
    )
    print(
        f"Rotating XC potential: {times['reference_rotations_time'] - times['k_set_time']:.3f} s"
    )
    print(
        f"Greens function inversion: {times['green_function_inversion_time'] - times['reference_rotations_time']:.3f} s"
    )
    print(
        f"Calculate energies and magnetic components: {times['end_time'] - times['green_function_inversion_time']:.3f} s"
    )


def parse_command_line():
    """This function can read input from the command line."""

    parser = ArgumentParser()

    parser.add_argument(
        "-i",
        "--input",
        dest="infile",
        default=None,
        type=str,
        help="Input file name",
        required=True,
    )
    parser.add_argument(
        "-o",
        "--output",
        dest="outfile",
        default=None,
        type=str,
        help="Output file name",
    )
    parser.add_argument(
        "--kset",
        dest="kset",
        default=None,
        type=int,
        help="k-space resolution of calculation",
    )
    parser.add_argument(
        "--kdirs",
        dest="kdirs",
        default=None,
        type=str,
        help="Definition of k-space dimensionality",
    )
    parser.add_argument(
        "--ebot",
        dest="ebot",
        default=None,
        type=float,
        help="Bottom energy of the contour",
    )
    parser.add_argument(
        "--eset",
        dest="eset",
        default=None,
        type=int,
        help="Number of energy points on the contour",
    )
    parser.add_argument(
        "--eset-p",
        dest="esetp",
        default=None,
        type=int,
        help="Parameter tuning the distribution on the contour",
    )

    parser.add_argument(
        "--parallel-green",
        dest="parallel_solver_for_Gk",
        default=None,
        type=bool,
        help="Whether to use the parallel or sequential solver for Greens function",
    )
    parser.add_argument(
        "--padawan-mode",
        dest="padawan_mode",
        default=None,
        type=bool,
        help="If it is on it turns on extra helpful information for new users",
    )

    cmd_line_args = parser.parse_args()

    return cmd_line_args


def main(simulation_parameters, magnetic_entities, pairs):
    # runtime information
    times = dict()
    times["start_time"] = timer()

    # MPI parameters
    comm = MPI.COMM_WORLD
    size = comm.Get_size()
    rank = comm.Get_rank()
    root_node = 0

    if rank == root_node:
        # include parallel size in simulation parameters
        simulation_parameters["parallel_size"] = size

        # check versions for debugging
        try:
            print("sisl version: ", sisl.__version__)
        except:
            print("sisl version unknown.")
        try:
            print("numpy version: ", np.__version__)
        except:
            print("numpy version unknown.")

    # rename outfile
    if not simulation_parameters["outfile"].endswith(".pickle"):
        simulation_parameters["outfile"] += ".pickle"

    # if ebot is not given put it 1 eV under the smallest energy
    if simulation_parameters["ebot"] is None:
        try:
            eigfile = simulation_parameters["infile"][:-3] + "EIG"
            simulation_parameters["ebot"] = read_siesta_emin(eigfile) - 1
            simulation_parameters["automatic_ebot"] = True
        except:
            print("Could not determine ebot.")
            print("Parameter was not given and .EIG file was not found.")
    else:
        simulation_parameters["automatic_ebot"] = False

    # add third direction for off-diagonal anisotropy
    for orientation in simulation_parameters["ref_xcf_orientations"]:
        o1 = orientation["vw"][0]
        o2 = orientation["vw"][1]
        orientation["vw"].append((o1 + o2) / np.sqrt(2))

    # read sile
    fdf = sisl.get_sile(simulation_parameters["infile"])

    # read in hamiltonian
    dh = fdf.read_hamiltonian()

    # read unit cell vectors
    simulation_parameters["cell"] = fdf.read_geometry().cell

    # unit cell index
    uc_in_sc_idx = dh.lattice.sc_index([0, 0, 0])

    if rank == root_node:
        print("\n\n\n\n\n")
        print(
            "#################################################################### JOB INFORMATION ###########################################################################"
        )
        print_job_description(simulation_parameters)
        print(
            "################################################################################################################################################################"
        )
        print("\n\n\n\n\n")

        times["setup_time"] = timer()
        print(f"Setup done. Elapsed time: {times['setup_time']} s")
        print(
            "================================================================================================================================================================"
        )

    NO = dh.no  # shorthand for number of orbitals in the unit cell

    # reformat Hamltonian and Overlap matrix for manipulations
    hh, ss = build_hh_ss(dh)

    # symmetrizing Hamiltonian and Overlap matrix to make them hermitian
    for i in range(dh.lattice.sc_off.shape[0]):
        j = dh.lattice.sc_index(-dh.lattice.sc_off[i])
        h1, h1d = hh[i], hh[j]
        hh[i], hh[j] = (h1 + h1d.T.conj()) / 2, (h1d + h1.T.conj()) / 2
        s1, s1d = ss[i], ss[j]
        ss[i], ss[j] = (s1 + s1d.T.conj()) / 2, (s1d + s1.T.conj()) / 2

    # identifying TRS and TRB parts of the Hamiltonian
    TAUY = np.kron(np.eye(NO), TAU_Y)
    hTR = np.array([TAUY @ hh[i].conj() @ TAUY for i in range(dh.lattice.nsc.prod())])
    hTRS = (hh + hTR) / 2
    hTRB = (hh - hTR) / 2

    # extracting the exchange field
    traced = [spin_tracer(hTRB[i]) for i in range(dh.lattice.nsc.prod())]  # equation 77
    XCF = np.array(
        [
            np.array([f["x"] / 2 for f in traced]),
            np.array([f["y"] / 2 for f in traced]),
            np.array([f["z"] / 2 for f in traced]),
        ]
    )

    # check if exchange field has scalar part
    max_xcfs = abs(np.array(np.array([f["c"] / 2 for f in traced]))).max()
    if max_xcfs > 1e-12:
        warnings.warn(
            f"Exchange field has non negligible scalar part. Largest value is {max_xcfs}"
        )

    if rank == root_node:
        times["H_and_XCF_time"] = timer()
        print(
            f"Hamiltonian and exchange field rotated. Elapsed time: {times['H_and_XCF_time']} s"
        )
        print(
            "================================================================================================================================================================"
        )

    # initialize pairs and magnetic entities based on input information
    pairs, magnetic_entities = setup_pairs_and_magnetic_entities(
        magnetic_entities, pairs, dh, simulation_parameters
    )

    if rank == root_node:
        times["site_and_pair_dictionaries_time"] = timer()
        print(
            f"Site and pair dictionaries created. Elapsed time: {times['site_and_pair_dictionaries_time']} s"
        )
        print(
            "================================================================================================================================================================"
        )

    # generate k space sampling
    kset = make_kset(
        dirs=simulation_parameters["kdirs"], NUMK=simulation_parameters["kset"]
    )

    # generate weights for k points
    wkset = np.ones(len(kset)) / len(kset)

    # split the k points based on MPI size
    kpcs = np.array_split(kset, size)

    # use progress bar if available
    if rank == root_node and tqdm_imported:
        kpcs[root_node] = tqdm(kpcs[root_node], desc="k loop")

    if rank == root_node:
        times["k_set_time"] = timer()
        print(f"k set created. Elapsed time: {times['k_set_time']} s")
        print(
            "================================================================================================================================================================"
        )

    # this will contain the three Hamiltonian in the
    # reference directions needed to calculate the energy
    # variations upon rotation
    hamiltonians = []

    # iterate over the reference directions (quantization axes)
    for i, orient in enumerate(simulation_parameters["ref_xcf_orientations"]):
        # obtain rotated exchange field and Hamiltonian
        R = RotMa2b(simulation_parameters["scf_xcf_orientation"], orient["o"])
        rot_XCF = np.einsum("ij,jklm->iklm", R, XCF)
        rot_H_XCF = sum(
            [np.kron(rot_XCF[i], tau) for i, tau in enumerate([TAU_X, TAU_Y, TAU_Z])]
        )
        rot_H_XCF_uc = rot_H_XCF[uc_in_sc_idx]

        # obtain total Hamiltonian with the rotated exchange field
        rot_H = hTRS + rot_H_XCF  # equation 76

        # store the relevant information of the Hamiltonian
        hamiltonians.append(dict(orient=orient["o"], H=rot_H))

        # these are the rotations (for now) perpendicular to the quantization axis
        for u in orient["vw"]:
            # section 2.H
            Tu = np.kron(np.eye(NO, dtype=int), tau_u(u))
            Vu1, Vu2 = calc_Vu(rot_H_XCF_uc, Tu)

            for mag_ent in magnetic_entities:
                idx = mag_ent["spin_box_indices"]
                # fill up the perturbed potentials (for now) based on the on-site projections
                mag_ent["Vu1"][i].append(onsite_projection(Vu1, idx, idx))
                mag_ent["Vu2"][i].append(onsite_projection(Vu2, idx, idx))

    if rank == root_node:
        times["reference_rotations_time"] = timer()
        print(
            f"Rotations done perpendicular to quantization axis. Elapsed time: {times['reference_rotations_time']} s"
        )
        print(
            "================================================================================================================================================================"
        )

    # provide helpful information to estimate the runtime and memory
    # requirements of the Greens function calculations
    if rank == root_node:
        print("Starting matrix inversions.")
        if simulation_parameters["padawan_mode"]:
            print("Padawan mode: ")
            print(f"Total number of k points: {kset.shape[0]}")
            print(
                f"Number of energy samples per k point: {simulation_parameters['eset']}"
            )
            print(f"Total number of directions: {len(hamiltonians)}")
            print(
                f"Total number of matrix inversions: {kset.shape[0] * len(hamiltonians) * simulation_parameters['eset']}"
            )
            print(
                f"The shape of the Hamiltonian and the Greens function is {NO}x{NO}={NO*NO}"
            )
            # https://stackoverflow.com/questions/70746660/how-to-predict-memory-requirement-for-np-linalg-inv
            # memory is O(64 n**2) for complex matrices
            memory_size = getsizeof(hamiltonians[0]["H"].base) / 1024
            print(
                f"Memory taken by a single Hamiltonian is: {getsizeof(hamiltonians[0]['H'].base) / 1024} KB"
            )
            print(f"Expected memory usage per matrix inversion: {memory_size * 32} KB")
            print(
                f"Expected memory usage per k point for parallel inversion: {memory_size * len(hamiltonians) * simulation_parameters['eset'] * 32} KB"
            )
            print(
                f"Expected memory usage on root node: {len(np.array_split(kset, size)[0]) * memory_size * len(hamiltonians) * simulation_parameters['eset'] * 32 / 1024} MB"
            )
        print(
            "================================================================================================================================================================"
        )

    # MPI barrier
    comm.Barrier()

    # make energy contour
    cont = make_contour(
        emin=simulation_parameters["ebot"],
        enum=simulation_parameters["eset"],
        p=simulation_parameters["esetp"],
    )
    eran = cont.ze

    # sampling the integrand on the contour and the BZ
    for k in kpcs[rank]:
        # weight of k point in BZ integral
        wk = wkset[rank]

        # iterate over reference directions
        for i, hamiltonian_orientation in enumerate(hamiltonians):
            # calculate Hamiltonian and Overlap matrix in a given k point
            H = hamiltonian_orientation["H"]
            HK, SK = hsk(H, ss, dh.sc_off, k)

            if simulation_parameters["parallel_solver_for_Gk"]:
                Gk = parallel_Gk(HK, SK, eran, simulation_parameters["eset"])
            else:
                # solve Greens function sequentially for the energies, because of memory bound
                Gk = sequential_GK(HK, SK, eran, simulation_parameters["eset"])

            # store the Greens function slice of the magnetic entities
            for mag_ent in magnetic_entities:
                idx = mag_ent["spin_box_indices"]
                mag_ent["Gii_tmp"][i] += onsite_projection(Gk, idx, idx) * wk

            for pair in pairs:
                # add phase shift based on the cell difference
                phase = np.exp(1j * 2 * np.pi * k @ pair["Ruc"].T)

                # get the pair orbital sizes from the magnetic entities
                ai = magnetic_entities[pair["ai"]]["spin_box_indices"]
                aj = magnetic_entities[pair["aj"]]["spin_box_indices"]

                # store the Greens function slice of the magnetic entities
                pair["Gij_tmp"][i] += onsite_projection(Gk, ai, aj) * phase * wk
                pair["Gji_tmp"][i] += onsite_projection(Gk, aj, ai) / phase * wk

    # summ reduce partial results of mpi nodes
    for i in range(len(hamiltonians)):
        for mag_ent in magnetic_entities:
            comm.Reduce(mag_ent["Gii_tmp"][i], mag_ent["Gii"][i], root=root_node)

        for pair in pairs:
            comm.Reduce(pair["Gij_tmp"][i], pair["Gij"][i], root=root_node)
            comm.Reduce(pair["Gji_tmp"][i], pair["Gji"][i], root=root_node)

    if rank == root_node:
        times["green_function_inversion_time"] = timer()
        print(
            f"Calculated Greens functions. Elapsed time: {times['green_function_inversion_time']} s"
        )
        print(
            "================================================================================================================================================================"
        )

    if rank == root_node:
        # iterate over the magnetic entities
        for tracker, mag_ent in enumerate(magnetic_entities):
            # iterate over the quantization axes
            for i, Gii in enumerate(mag_ent["Gii"]):
                storage = []
                # iterate over the first and second order local perturbations
                for Vu1, Vu2 in zip(mag_ent["Vu1"][i], mag_ent["Vu2"][i]):
                    # The Szunyogh-Lichtenstein formula
                    traced = np.trace(
                        (Vu2 @ Gii + 0.5 * Gii @ Vu1 @ Gii), axis1=1, axis2=2
                    )  # this is the on site projection
                    # evaluation of the contour integral
                    storage.append(int_de_ke(traced, cont.we))
                # fill up the magnetic entities dictionary with the energies
                magnetic_entities[tracker]["energies"].append(storage)
            # convert to np array
            magnetic_entities[tracker]["energies"] = np.array(
                magnetic_entities[tracker]["energies"]
            )

        # iterate over the pairs
        for tracker, pair in enumerate(pairs):
            # iterate over the quantization axes
            for i, (Gij, Gji) in enumerate(zip(pair["Gij"], pair["Gji"])):
                site_i = magnetic_entities[pair["ai"]]
                site_j = magnetic_entities[pair["aj"]]
                storage = []
                # iterate over the first order local perturbations in all possible orientations for the two sites
                for Vui in site_i["Vu1"][i]:
                    for Vuj in site_j["Vu1"][i]:
                        # The Szunyogh-Lichtenstein formula
                        traced = np.trace(
                            (Vui @ Gij @ Vuj @ Gji), axis1=1, axis2=2
                        )  # this is the on site projection
                        # evaluation of the contour integral
                        storage.append(int_de_ke(traced, cont.we))
                # fill up the pairs dictionary with the energies
                pairs[tracker]["energies"].append(storage)
            # convert to np array
            pairs[tracker]["energies"] = np.array(pairs[tracker]["energies"])

        # calculate magnetic parameters
        for mag_ent in magnetic_entities:
            K, consistency = calculate_anisotropy_tensor(mag_ent)
            mag_ent["K"] = K * sisl.unit_convert("eV", "meV")
            mag_ent["K_consistency"] = consistency

        for pair in pairs:
            J_iso, J_S, D, J = calculate_exchange_tensor(pair)
            pair["J_iso"] = J_iso * sisl.unit_convert("eV", "meV")
            pair["J_S"] = J_S * sisl.unit_convert("eV", "meV")
            pair["D"] = D * sisl.unit_convert("eV", "meV")
            pair["J"] = J * sisl.unit_convert("eV", "meV")

        times["end_time"] = timer()
        print("\n\n\n\n\n")
        print(
            "##################################################################### GROGU OUTPUT #############################################################################"
        )

        print_parameters(simulation_parameters)
        print_atoms_and_pairs(magnetic_entities, pairs)
        print(
            "################################################################################################################################################################"
        )
        print_runtime_information(times)
        print("")

        # remove unwanted stuff before saving
        pairs, magnetic_entities = remove_clutter_for_save(pairs, magnetic_entities)
        # create output dictionary with all the relevant data
        results = dict(
            parameters=simulation_parameters,
            magnetic_entities=magnetic_entities,
            pairs=pairs,
            runtime=times,
        )

        # save results
        save_pickle(simulation_parameters["outfile"], results)

        print("\n\n\n\n\n")


if __name__ == "__main__":
    # loading parameters
    # it is not clear how to give grogu.fdf path...
    # command_line_arguments = parse_command_line()
    # fdf_arguments, magnetic_entities, pairs = read_fdf(command_line_arguments["infile"])

    # right now we do not use any of these input, but it shows
    # the order of priority when processing arguments
    default_arguments = False
    fdf_arguments = False
    command_line_arguments = False
    # simulation_parameters = process_input_args(
    #    default_arguments, fdf_arguments, command_line_arguments, ACCEPTED_INPUTS
    # )

    ####################################################################################################
    #                                   This is the input file for now                                 #
    ####################################################################################################
    magnetic_entities = [
        dict(atom=3, l=2),
        dict(atom=4, l=2),
        dict(atom=5, l=2),
    ]
    pairs = [
        dict(ai=0, aj=1, Ruc=np.array([0, 0, 0])),
        dict(ai=0, aj=2, Ruc=np.array([0, 0, 0])),
        dict(ai=1, aj=2, Ruc=np.array([0, 0, 0])),
        dict(ai=0, aj=2, Ruc=np.array([-1, -1, 0])),
        dict(ai=1, aj=2, Ruc=np.array([-1, -1, 0])),
        dict(ai=0, aj=2, Ruc=np.array([-1, 0, 0])),
        dict(ai=1, aj=2, Ruc=np.array([-1, 0, 0])),
        dict(ai=1, aj=2, Ruc=np.array([-2, 0, 0])),
        dict(ai=1, aj=2, Ruc=np.array([-3, 0, 0])),
    ]
    simulation_parameters = dict()
    simulation_parameters["infile"] = ""
    simulation_parameters["outfile"] = "./"
    simulation_parameters["scf_xcf_orientation"] = np.array([0, 0, 1])
    simulation_parameters["ref_xcf_orientations"] = [
        dict(o=np.array([1, 0, 0]), vw=[np.array([0, 1, 0]), np.array([0, 0, 1])]),
        dict(o=np.array([0, 1, 0]), vw=[np.array([1, 0, 0]), np.array([0, 0, 1])]),
        dict(o=np.array([0, 0, 1]), vw=[np.array([1, 0, 0]), np.array([0, 1, 0])]),
    ]
    simulation_parameters["kset"] = 0
    simulation_parameters["kdirs"] = ""
    simulation_parameters["ebot"] = None
    simulation_parameters["eset"] = 0
    simulation_parameters["esetp"] = 0
    simulation_parameters["parallel_solver_for_Gk"] = False
    simulation_parameters["padawan_mode"] = True
    ####################################################################################################
    #                                   This is the input file for now                                 #
    ####################################################################################################

    main(simulation_parameters, magnetic_entities, pairs)
