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

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