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

import numpy as np


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
