# 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 globals import TAU_0, TAU_X, TAU_Y, TAU_Z
from scipy.special import roots_legendre
from sisl.io.siesta import eigSileSiesta


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