# Source code for grogupy.magnetism

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

        It uses the grogu convention for output.

        Args:
            mag_ent : dict
                An element from the magnetic entities

        Returns:
            K : np.array_like
                elements of the anisotropy tensor
        """

        # get the energies
        energies = mag_ent["energies"]

        # calculate the diagonal tensor elements
        Kxx = energies[1, 1] - energies[1, 0]
        Kyy = energies[0, 1] - energies[0, 0]
        Kzz = 0

        # perform consistency check
        calculated_diff = Kyy - Kxx
        expected_diff = energies[2, 0] - energies[2, 1]
        consistency_check = abs(calculated_diff - expected_diff)

        return Kxx, Kyy, Kzz, 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