# Source code for grogupy.utilities

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

    import numpy as np
    from scipy.special import roots_legendre
    from sisl.io.siesta import eigSileSiesta

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

    def commutator(a, b):
        """Shorthand for commutator.

        Commutator of two matrices in the mathematical sense.

        Args:
            a : np.array_like
                The first matrix
            b : np.array_like
                The second matrix

        Returns:
            np.array_like
                The commutator of a and b
        """

        return a @ b - b @ a

    # define some useful functions

    def hsk(H, ss, sc_off, k=(0, 0, 0)):
        """Speed up Hk and Sk generation.

        Calculates the Hamiltonian and the Overlap matrix at a given k point. It is faster that the sisl version.

        Args:
            H : np.array_like
                Hamiltonian in spin box form
            ss : np.array_like
                Overlap matrix in spin box form
            sc_off : list
                supercell indexes of the Hamiltonian
            k : tuple, optional
                The k point where the matrices are set up. Defaults to (0, 0, 0)

        Returns:
            np.array_like
                Hamiltonian at the given k point
            np.array_like
                Overlap matrix at the given k point
        """

        # this two conversion lines are from the sisl source
        k = np.asarray(k, np.float64)
        k.shape = (-1,)

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

        # phases applied to the hamiltonian
        HK = np.einsum("abc,a->bc", H, phases)
        SK = np.einsum("abc,a->bc", ss, phases)

        return HK, SK

    def make_kset(dirs="xyz", NUMK=20):
        """Simple k-grid generator to sample the Brillouin zone.

        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 returned. The total number of k points is the NUMK**(dimensions)

        Args:
            dirs : str, optional
                Directions of the k points in the Brillouin zone. They are the three lattice vectors. Defaults to "xyz"
            NUMK : int, optional
                The number of k points in a direction. Defaults to 20

        Returns:
            np.array_like
                An array of k points that uniformly sample the Brillouin zone in the given directions
        """

        # if there is no xyz in dirs return the Gamma point
        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 make_contour(emin=-20, emax=0.0, enum=42, p=150):
        """A more sophisticated contour generator.

        Calculates the parameters for the complex contour integral. It uses the
        Legendre-Gauss quadrature method.   It returns a class that contains
        the information for the contour integral.

        Args:
            emin : int, optional
                Energy minimum of the contour. Defaults to -20
            emax : float, optional
                Energy maximum of the contour. Defaults to 0.0, so the Fermi level
            enum : int, optional
                Number of sample points along the contour. Defaults to 42
            p : int, optional
                Shape parameter that describes the distribution of the sample points. Defaults to 150

        Returns:
            ccont
                Contains all the information for the contour integral
        """
        x, wl = roots_legendre(enum)
        R = (emax - emin) / 2  # radius
        z0 = (emax + emin) / 2  # center point
        y1 = -np.log(1 + np.pi * p)  # lower bound
        y2 = 0  # upper bound

        y = (y2 - y1) / 2 * x + (y2 + y1) / 2
        phi = (np.exp(-y) - 1) / p  # angle parameter
        ze = z0 + R * np.exp(1j * phi)  # complex points for path
        we = -(y2 - y1) / 2 * np.exp(-y) / p * 1j * (ze - z0) * wl

        # just an empty container class
        class ccont:
            pass

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

        return cont

    def tau_u(u):
        """Pauli matrix in direction u.

        Returns the vector u in the basis of the Pauli matrices.

        Args:
            u : list or np.array_like
                The direction

        Returns:
            np.array_like
                Arbitrary direction in the base of the Pauli matrices
        """

        # u is force to be of unit length
        u = u / np.linalg.norm(u)

        return u[0] * tau_x + u[1] * tau_y + u[2] * tau_z

    #

    def crossM(u):
        """Definition for the cross-product matrix.

        It acts as a cross product with vector u.

        Args:
            u : list or np.array_like
                The second vector in the cross product

        Returns:
            np.array_like
                The matrix that represents teh cross product with a vector
        """
        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.

        Args:
            theta : float
                The angle of rotation
            u : np.array_like
                The rotation axis
            eps : float, optional
                Cutoff for small elements in the resulting matrix. Defaults to 1e-10

        Returns:
            np.array_like
                The rotation matrix
        """

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

        # kill off small numbers
        M[abs(M) < eps] = 0.0
        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.

        Args:
            a : np.array_like
                First vector
            b : np.array_like
                Second vector
            eps : float, optional
                Cutoff for small elements in the resulting matrix. Defaults to 1e-10

        Returns:
            np.array_like
                The rotation matrix with the above property
        """

        v = np.cross(a, b)
        c = a @ b
        M = np.eye(3) + crossM(v) + crossM(v) @ crossM(v) / (1 + c)

        # kill off small numbers
        M[abs(M) < eps] = 0.0
        return M

    def read_siesta_emin(eigfile):
        """It reads the lowest energy level from the siesta run.

        It uses the .EIG file from siesta that contains the eigenvalues.

        Args:
            eigfile : str
                The path to the .EIG file

        Returns:
            float
                The energy minimum
        """

        # read the file
        eigs = eigSileSiesta(eigfile).read_data()

        return eigs.min()

    def int_de_ke(traced, we):
        """It numerically integrates the traced matrix.

        It is a wrapper from numpy.trapz and it contains the
        relevant constants to calculate the energy integral from
        equation 93 or 96.

        Args:
            traced : np.array_like
                The trace of a matrix or a matrix product
            we : float
                The weight of a point on the contour

        Returns:
            float
                The energy calculated from the integral formula
        """

        return np.trapz(-1 / np.pi * np.imag(traced * we))