# Copyright (c) [2024] [Daniel Pozsar]
#
# 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 itertools import permutations, product
from pprint import pprint

import numpy as np
from scipy.special import roots_legendre
from sisl import unit_convert

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


# define some useful functions
def hsk(H, ss, sc_off, k=(0, 0, 0)):
    """
    One way to speed up Hk and Sk generation
    """
    k = np.asarray(k, np.float64)  # this two conversion lines
    k.shape = (-1,)  # are from the sisl source

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

    HK = np.einsum("abc,a->bc", H, phases)
    SK = np.einsum("abc,a->bc", ss, phases)

    return HK, SK


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

    x, wl = roots_legendre(enum)
    R = (emax - emin) / 2
    z0 = (emax + emin) / 2
    y1 = -np.log(1 + np.pi * p)
    y2 = 0

    y = (y2 - y1) / 2 * x + (y2 + y1) / 2
    phi = (np.exp(-y) - 1) / p
    ze = z0 + R * np.exp(1j * phi)
    we = -(y2 - y1) / 2 * np.exp(-y) / p * 1j * (ze - z0) * wl

    class ccont:
        # just an empty container class
        pass

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

    return cont


def make_kset(dirs="xyz", NUMK=20):
    """
    Simple k-grid generator. 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 returend.
    """
    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 commutator(a, b):
    "Shorthand for commutator"
    return a @ b - b @ a


def tau_u(u):
    """
    Pauli matrix in direction u.
    """
    u = u / np.linalg.norm(u)  # u is force to be of unit length
    return u[0] * tau_x + u[1] * tau_y + u[2] * tau_z


#
def crossM(u):
    """
    Definition for the cross-product matrix.
    Acting as a cross product with vector u.
    """
    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.
    """
    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)
    )
    M[abs(M) < eps] = 0.0  # kill off small numbers
    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.
    """
    v = np.cross(a, b)
    c = a @ b
    M = np.eye(3) + crossM(v) + crossM(v) @ crossM(v) / (1 + c)
    M[abs(M) < eps] = 0.0  # kill off small numbers
    return M


def spin_tracer(M):
    """
    Spin tracer utility.
    This akes 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.
    """

    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 indeces of a given magnetic entity.
    dh: a sisl Hamiltonian object
    atom: an integer or list of integers, defining atom (or atoms) in the unicell forming the magnetic entity
    l: integer, defining the angular momentum channel
    """
    # 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 blow_up_orbindx(orb_indices):
    """
    Function to blow up orbital indeces to make SPIN BOX indices.
    """
    return np.array([[2 * o, 2 * o + 1] for o in orb_indices]).flatten()


def calculate_exchange_tensor(pair):
    o1, o2, o3 = pair["energies"]  # o1=x, o2=y, o3=z
    # 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])]),

    J_ii = np.array([o2[-1], o1[0], o1[-1]])  # xx, yy, zz
    J_S = -0.5 * np.array([o3[1] + o3[2], o2[1] + o2[1], o1[1] + o1[2]])  # yz, zx, xy
    D = 0.5 * np.array([o1[1] - o1[2], o2[2] - o2[1], o3[1] - o3[2]])  # x, y, z
    return J_ii.sum() / 3, np.concatenate([J_ii[:2] - J_ii.sum() / 3, J_S]).flatten(), D


def print_pair_atomic_indices(pair, magnetic_entities, dh):
    atomic_indices = ""
    # iterate over the two magnetic entities in a pair
    for mag_ent in [magnetic_entities[pair["ai"]], magnetic_entities[pair["aj"]]]:
        # get atoms of magnetic entity
        atoms_idx = mag_ent["atom"]
        # if orbital is not set use all
        if "l" not in mag_ent.keys():
            mag_ent["l"] = "all"
        orbitals = mag_ent["l"]

        # if magnetic entity contains one atoms
        if isinstance(atoms_idx, int):
            atomic_indices += 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
            atomic_indices += atom_group[:-2] + "}"
        # separate magnetic entities
        atomic_indices += "        "
    return atomic_indices


def print_output(simulation_parameters, magnetic_entities, pairs, dh, times):
    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 simulation_parameters["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(
        "[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:
        # get atoms of magnetic entity
        atoms_idx = mag_ent["atom"]
        # if orbital is not set use all
        if "l" not in mag_ent.keys():
            mag_ent["l"] = "all"
        orbitals = mag_ent["l"]

        # if magnetic entity contains one atom
        if isinstance(atoms_idx, int):
            # coordinates and tag
            x, y, z = dh.xyz[atoms_idx]
            print(
                f"[{atoms_idx}]{dh.atoms[atoms_idx].tag}({orbitals})        {x}        {y}        {z}"
            )

        # if magnetic entity contains more than one atoms
        if isinstance(atoms_idx, list):
            # iterate over atoms
            for atom_idx in atoms_idx:
                # coordinates and tag
                x, y, z = dh.xyz[atom_idx]
                print(
                    f"[{atom_idx}]{dh.atoms[atom_idx].tag}({orbitals})        {x}        {y}        {z}"
                )
        print("")

    print(
        "================================================================================================================================================================"
    )
    print("Exchange [meV]")
    print(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    print("Magnetic entity1        Magnetic entity2       [i  j  k]       d [Ang]")
    print(
        "----------------------------------------------------------------------------------------------------------------------------------------------------------------"
    )
    # iterate over pairs
    for pair in pairs:
        # calculate magnetic parameters
        J_iso, J_S, D = calculate_exchange_tensor(pair)
        J_iso = J_iso * unit_convert("eV", "meV")
        J_S = J_S * unit_convert("eV", "meV")
        D = D * unit_convert("eV", "meV")

        # print pair parameters
        print(
            print_pair_atomic_indices(pair, magnetic_entities, dh),
            f"    {pair['Ruc']}    d [Ang] Not yet.",
        )
        # print magnetic parameters
        print("Isotropic: ", J_iso)
        print("DMI: ", D)
        print("Symmetric-anisotropy: ", J_S)
        print("Energies for debugging: ")
        pprint(np.array(pair["energies"]))
        print(
            "J_ii for debugging: (check if this is the same as in calculate_exchange_tensor)"
        )
        o1, o2, o3 = pair["energies"]
        pprint(np.array([o2[-1], o3[0], o1[0]]))
        print("Test J_xx = E(y,z) = E(z,y)")
        print(o2[-1], o3[-1])
        print("")

    print(
        "================================================================================================================================================================"
    )
    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"
    )