# Source code for grogupy.io

    # 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.

    """Docstring in io.
    """

    from pickle import dump, load

    import numpy as np
    from sisl.io import fdfSileSiesta

    # 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 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("Anisotropy diag: ", 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"
        )