import warnings
from sys import stdout
from timeit import default_timer as timer

import numpy as np
import sisl
from mpi4py import MPI
from numpy.linalg import inv
from tqdm import tqdm
from useful import *


def main():
    start_time = timer()

    # this cell mimicks an input file
    fdf = sisl.get_sile("../../lat3_791/Fe3GeTe2.fdf")
    # this information needs to be given at the input!!
    scf_xcf_orientation = np.array([0, 0, 1])  # z
    # list of reference directions for around which we calculate the derivatives
    # o is the quantization axis, v and w are two axes perpendicular to it
    # at this moment the user has to supply o,v,w on the input.
    # we can have some default for this
    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])]),
    ]

    # human readable definition of magnetic entities
    # magnetic_entities = [
    #    dict(atom=0, ),
    #    dict(atom=1, ),
    #    dict(atom=2, ),
    #    dict(atom=3, l=2),
    #    dict(atom=4, l=2),
    #    dict(atom=5, l=2),
    # ]
    # pairs = [
    #    dict(ai=3, aj=4, Ruc=np.array([0, 0, 0])),  # isotropic should be -82 meV
    #    dict(ai=3, aj=5, Ruc=np.array([0, 0, 0])),  # these should all be around -41.9 in the isotropic part
    #    dict(ai=4, aj=5, Ruc=np.array([0, 0, 0])),
    #    dict(ai=3, aj=0, Ruc=np.array([0, 0, 0])),
    #    dict(ai=3, aj=1, Ruc=np.array([0, 0, 0])),
    #    dict(ai=3, aj=2, Ruc=np.array([0, 0, 0])),
    # ]
    # magnetic_entities = [
    #    dict(atom=3, l=2),
    #    dict(atom=4, l=2),
    #    dict(atom=5, l=2),
    # ]

    # pair information
    # pairs = [
    #    dict(ai=0, aj=1, Ruc=np.array([0, 0, 0])),  # isotropic should be -82 meV
    #    dict(
    #        ai=0, aj=2, Ruc=np.array([0, 0, 0])
    #    ),  # these should all be around -41.9 in the isotropic part
    #    dict(ai=1, aj=2, Ruc=np.array([0, 0, 0])),
    #    dict(ai=0, aj=1, Ruc=np.array([-1, 0, 0])),
    #    dict(ai=0, aj=2, Ruc=np.array([-1, 0, 0])),
    #    dict(ai=0, aj=1, Ruc=np.array([1, 0, 0])),
    #    dict(ai=0, aj=2, Ruc=np.array([1, 0, 0])),
    #    dict(ai=0, aj=1, Ruc=np.array([0, -1, 0])),
    #    dict(ai=0, aj=2, Ruc=np.array([0, -1, 0])),
    #    dict(ai=0, aj=1, Ruc=np.array([0, 1, 0])),
    #    dict(ai=0, aj=2, Ruc=np.array([0, 1, 0])),
    #    dict(ai=1, aj=2, Ruc=np.array([-1, 0, 0])),
    # ]
    magnetic_entities = [
        dict(atom=3, l=2),
        dict(atom=4, l=2),
        dict(atom=5, l=2),
        dict(
            atom=[3, 4],
        ),
    ]

    # pair information
    pairs = [
        dict(ai=0, aj=1, Ruc=np.array([0, 0, 0])),  # isotropic should be -82 meV
        dict(
            ai=0, aj=2, Ruc=np.array([0, 0, 0])
        ),  # these should all be around -41.9 in the isotropic part
        dict(ai=1, aj=2, Ruc=np.array([0, 0, 0])),
        dict(ai=0, aj=2, Ruc=np.array([-1, 0, 0])),
        dict(ai=1, aj=2, Ruc=np.array([-1, 0, 0])),
    ]

    # Brilloun zone sampling and Green function contour integral
    kset = 100
    kdirs = "xy"
    ebot = -30
    eset = 100
    esetp = 1000

    # MPI parameters
    comm = MPI.COMM_WORLD
    size = comm.Get_size()
    rank = comm.Get_rank()
    root_node = 0
    if rank == root_node:
        print("Number of nodes in the parallel cluster: ", size)

    simulation_parameters = dict(
        path="Not yet specified.",
        scf_xcf_orientation=scf_xcf_orientation,
        ref_xcf_orientations=ref_xcf_orientations,
        kset=kset,
        kdirs=kdirs,
        ebot=ebot,
        eset=eset,
        esetp=esetp,
        parallel_size=size,
    )

    # digestion of the input
    # read in hamiltonian
    dh = fdf.read_hamiltonian()
    try:
        simulation_parameters["geom"] = fdf.read_geometry()
    except:
        print("Error reading geometry.")

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

    setup_time = timer()

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

    # 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"] for f in traced]),
            np.array([f["y"] for f in traced]),
            np.array([f["z"] for f in traced]),
        ]
    )  # equation 77

    # Check if exchange field has scalar part
    max_xcfs = abs(np.array(np.array([f["c"] 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}"
        )

    H_and_XCF_time = timer()

    # for every site we have to store 3 Greens function (and the associated _tmp-s) in the 3 reference directions
    for i, mag_ent in enumerate(magnetic_entities):
        parsed = parse_magnetic_entity(dh, **mag_ent)  # parse orbital indexes
        magnetic_entities[i]["orbital_indeces"] = parsed
        magnetic_entities[i]["spin_box_indeces"] = blow_up_orbindx(
            parsed
        )  # calculate spin box indexes
        spin_box_shape = len(
            mag_ent["spin_box_indeces"]
        )  # calculate size for Greens function generation

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

        mag_ent["Gii"] = []  # Greens function
        mag_ent["Gii_tmp"] = []  # Greens function for parallelization
        mag_ent["Vu1"] = [
            list([]) for _ in range(len(ref_xcf_orientations))
        ]  # These will be the perturbed potentials from eq. 100
        mag_ent["Vu2"] = [list([]) for _ in range(len(ref_xcf_orientations))]
        for i in ref_xcf_orientations:
            mag_ent["Gii"].append(
                np.zeros((eset, spin_box_shape, spin_box_shape), dtype="complex128")
            )  # Greens functions for every quantization axis
            mag_ent["Gii_tmp"].append(
                np.zeros((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:
        spin_box_shape_i, spin_box_shape_j = len(
            magnetic_entities[pair["ai"]]["spin_box_indeces"]
        ), len(
            magnetic_entities[pair["aj"]]["spin_box_indeces"]
        )  # calculate size for Greens function generation

        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 i in ref_xcf_orientations:
            pair["Gij"].append(
                np.zeros((eset, spin_box_shape_i, spin_box_shape_j), dtype="complex128")
            )
            pair["Gij_tmp"].append(
                np.zeros((eset, spin_box_shape_i, spin_box_shape_j), dtype="complex128")
            )  # Greens functions for every quantization axis
            pair["Gji"].append(
                np.zeros((eset, spin_box_shape_j, spin_box_shape_i), dtype="complex128")
            )
            pair["Gji_tmp"].append(
                np.zeros((eset, spin_box_shape_j, spin_box_shape_i), dtype="complex128")
            )

    site_and_pair_dictionaries_time = timer()

    kset = make_kset(dirs=kdirs, NUMK=kset)  # generate k space sampling
    wkset = np.ones(len(kset)) / len(kset)  # generate weights for k points
    kpcs = np.array_split(kset, size)  # split the k points based on MPI size
    kpcs[root_node] = tqdm(kpcs[root_node], desc="k loop", file=stdout)

    k_set_time = timer()

    # this will contain all the data needed to calculate the energy variations upon rotation
    hamiltonians = []

    # iterate over the reference directions (quantization axes)
    for i, orient in enumerate(ref_xcf_orientations):
        # obtain rotated exchange field
        R = RotMa2b(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

        hamiltonians.append(
            dict(orient=orient["o"], H=rot_H, rotations=[])
        )  # store orientation and rotated Hamiltonian

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

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

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

    reference_rotations_time = timer()

    if rank == root_node:
        print("Number of magnetic entities being calculated: ", len(magnetic_entities))
        print(
            "We have to calculate the Greens function for three reference direction and we are going to calculate 15 energy integrals per site."
        )
        print(f"The shape of the Hamiltonian and the Greens function is {NO}x{NO}.")
    comm.Barrier()
    # ----------------------------------------------------------------------

    # make energy contour
    # we are working in eV now  !
    # and sisil shifts E_F to 0 !
    cont = make_contour(emin=ebot, enum=eset, p=esetp)
    eran = cont.ze

    # ----------------------------------------------------------------------
    # sampling the integrand on the contour and the BZ
    for k in kpcs[rank]:
        wk = wkset[rank]  # weight of k point in BZ integral
        for i, hamiltonian_orientation in enumerate(
            hamiltonians
        ):  # iterate over reference directions
            # calculate Greens function
            H = hamiltonian_orientation["H"]
            HK, SK = hsk(H, ss, dh.sc_off, k)
            Gk = inv(SK * eran.reshape(eset, 1, 1) - HK)

            # store the Greens function slice of the magnetic entities (for now) based on the on-site projections
            for mag_ent in magnetic_entities:
                mag_ent["Gii_tmp"][i] += (
                    Gk[:, mag_ent["spin_box_indeces"]][..., mag_ent["spin_box_indeces"]]
                    * 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_indeces"]
                aj = magnetic_entities[pair["aj"]]["spin_box_indeces"]

                # store the Greens function slice of the magnetic entities (for now) based on the on-site projections
                pair["Gij_tmp"][i] += Gk[:, ai][..., aj] * phase * wk
                pair["Gji_tmp"][i] += 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)

    green_function_inversion_time = timer()

    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
                    )
                    # evaluation of the contour integral
                    storage.append(np.trapz(-1 / np.pi * np.imag(traced * cont.we)))

                # fill up the magnetic entities dictionary with the energies
                mag_ent["energies"].append(storage)

        # 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)
                        # evaluation of the contour integral
                        storage.append(np.trapz(-1 / np.pi * np.imag(traced * cont.we)))

                # fill up the pairs dictionary with the energies
                pairs[tracker]["energies"].append(storage)

        end_time = timer()

        print(
            "############################### GROGU OUTPUT ###################################"
        )
        print(
            "================================================================================"
        )
        print("Input file: ")
        print(simulation_parameters["path"])
        print(
            "Number of nodes in the parallel cluster: ",
            simulation_parameters["parallel_size"],
        )
        print(
            "================================================================================"
        )
        try:
            print("Cell [Ang]: ")
            print(simulation_parameters["geom"].cell)
        except:
            print("Geometry could not be read.")
        print(
            "================================================================================"
        )
        print("DFT axis: ")
        print(simulation_parameters["scf_xcf_orientation"])
        print("Quantization axis and perpendicular rotation directions:")
        for ref in ref_xcf_orientations:
            print(ref["o"], " --» ", ref["vw"])
        print(
            "================================================================================"
        )
        print("number of k points: ", simulation_parameters["kset"])
        print("k point directions: ", simulation_parameters["kdirs"])
        print(
            "================================================================================"
        )
        print("Parameters for the contour integral:")
        print("Ebot: ", simulation_parameters["ebot"])
        print("Eset: ", simulation_parameters["eset"])
        print("Esetp: ", simulation_parameters["esetp"])
        print(
            "================================================================================"
        )
        print("Atomic informations: ")
        print("")
        print("")
        print("Not yet specified.")
        print("")
        print("")
        print(
            "================================================================================"
        )
        print("Exchange [meV]")
        print(
            "--------------------------------------------------------------------------------"
        )
        print("Atom1    Atom2   [i  j  k]   d [Ang]")
        print(
            "--------------------------------------------------------------------------------"
        )
        for pair in pairs:
            J_iso, J_S, D = calculate_exchange_tensor(pair)
            J_iso = J_iso * sisl.unit_convert("eV", "meV")
            J_S = J_S * sisl.unit_convert("eV", "meV")
            D = D * sisl.unit_convert("eV", "meV")

            print(print_atomic_indices(pair, magnetic_entities, dh))
            print("Isotropic: ", J_iso)
            print("DMI: ", D)
            print("Symmetric-anisotropy: ", J_S)
            print("")

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


if __name__ == "__main__":
    main()