import numpy as np
import numpy.typing as npt
from numpy.fft import fft, ifft, fftfreq
from scipy import constants as const
from scipy import sparse
from scipy import linalg


def discrete_hamiltonian(
    dx: tuple[float, ...] | float, potential: npt.NDArray, mass: float
) -> npt.NDArray:
    """Derive the hamiltonian for an arbitrary potential using the
    discretized laplace operator. The function is implemented for 1 to
    3 dimensions.

    :param dx: Step size used for the coordinates. If the potential is
        one-dimensional, this can simply be the step size of the
        coordinate used in :param:`potential`.
        For larger dimensions, this parameter must be a tuple with the
        same size as the dimension.
    :param potential: A (multidimensional) potential.
    :param mass: The mass of the particle, used for the kinetic term.
    """
    ndim = potential.ndim
    if np.shape(dx) != (ndim,):
        if ndim != 1 or np.shape(dx) != ():
            raise ValueError(f"Expected list of {ndim} values for 'dx'")
        else:
            dx = np.array([dx])
    grid_points = potential.shape
    laplace_d2_1 = sparse.diags(
        [1, -2, 1], [-1, 0, 1], shape=(grid_points[0], grid_points[0])
    ) / (dx[0] ** 2)

    if ndim == 1:
        laplace = laplace_d2_1
    elif ndim >= 2:
        laplace_d2_2 = sparse.diags(
            [1, -2, 1], [-1, 0, 1], shape=(grid_points[1], grid_points[1])
        ) / (dx[1] ** 2)
        if ndim == 2:
            # laplace = sparse.kron(laplace_d2_1, eye_2) + sparse.kron(
            #     eye_1, laplace_d2_2
            # )
            laplace = sparse.kronsum(laplace_d2_1, laplace_d2_2)
        if ndim == 3:
            laplace_d2_3 = sparse.diags(
                [1, -2, 1], [-1, 0, 1], shape=(grid_points[2], grid_points[2])
            ) / (dx[2] ** 2)
            laplace = sparse.kronsum(
                sparse.kronsum(laplace_d2_1, laplace_d2_2), laplace_d2_3
            )
    else:
        raise NotImplementedError(f"Number of dimensions not implemented: '{ndim}'")
    t = -(const.hbar**2) / (2 * mass) * laplace
    v = sparse.diags([potential.flatten()], [0])
    return t + v

def functional_hamiltonian(
    dx: tuple[float, ...] | float, potential: npt.NDArray, mass: float
) -> npt.NDArray:
    """Returns a scipy.sparse.linalg LinearOperator for the hamiltonian of
    an arbitrary potential using the discretized laplace operator.
    The function is implemented for 1 to 3 dimensions.

    :param dx: Step size used for the coordinates. If the potential is
        one-dimensional, this can simply be the step size of the
        coordinate used in :param:`potential`.
        For larger dimensions, this parameter must be a tuple with the
        same size as the dimension.
    :param potential: A (multidimensional) potential.
    :param mass: The mass of the particle, used for the kinetic term.
    """
    ndim = potential.ndim
    if np.shape(dx) != (ndim,):
        if ndim != 1 or np.shape(dx) != ():
            raise ValueError(f"Expected list of {ndim} values for 'dx'")
        else:
            dx = np.array([dx])
    grid_points = potential.shape

    if ndim != 3:
        raise Exception(f"Number of dimensions ({ndim}) is not implemented.")
    
    kx = 2*np.pi*fftfreq(grid_points[0], dx[0])
    ky = 2*np.pi*fftfreq(grid_points[1], dx[1])
    kz = 2*np.pi*fftfreq(grid_points[2], dx[2])
    KX, KY, KZ = np.meshgrid(kx, ky, kz, indexing='ij')

    def mv(psi):
        """Function for action of hamiltonian on wavefunction"""
        psi_nonflat = psi.reshape(grid_points)
        #calculate the laplacian with a fourier transform
        laplace_psi = (ifft(-KX**2*fft(psi_nonflat, axis = 0), axis = 0) + 
                        ifft(-KY**2*fft(psi_nonflat, axis = 1), axis = 1) + 
                        ifft(-KZ**2*fft(psi_nonflat, axis = 2), axis = 2))
        v = sparse.diags([potential.flatten()], [0])
        return -(const.hbar**2) / (2 * mass) *laplace_psi.flatten() + v*psi

    len_psi = np.prod(grid_points)
    return sparse.linalg.LinearOperator((len_psi,len_psi), matvec=mv)

def nstationary_solution(
    dx: tuple[float, ...] | float, potential: npt.NDArray, mass: float, k: int = 1,
    method="sparse") -> tuple[npt.NDArray, npt.NDArray]:
    """Derive the stationary solution for a discrete potential of 1 to
    3 dimensions numerically.

    :param dx: Step size used for the coordinates. If the potential is
        one-dimensional, this can simply be the step size of the
        coordinate used in :param:`potential`.
        For larger dimensions, this parameter must be a tuple with the
        same size as the dimension.
    :param potential: A (multidimensional) potential.
    :param mass: The mass of the particle, used for the kinetic term.
    :param k: Number of eigenvalues and eigenstates to return.
    :returns: Tuple of eigenvalues and eigenvectors. Note that the
        eigenvectors are transposed so that they are of shape
        ``(k, np.size(potential))``
    """
    e_min = np.min(potential)
    if method == "matrix_free":
        eigenvals, eigenvects = sparse.linalg.eigsh(
            functional_hamiltonian(dx, potential, mass), k=k, which="LM", sigma=e_min)
    elif method == "sparse":
        eigenvals, eigenvects = sparse.linalg.eigsh(
            discrete_hamiltonian(dx, potential, mass), k=k, which="LM", sigma=e_min)
    else:
            raise ValueError(f'Method "{method}" does not exist.')
        
    
    return eigenvals, eigenvects.T

def nstationary_solution_export(path,
    dx: tuple[float, ...] | float, potential: npt.NDArray, mass: float, k: int = 1
) -> tuple[npt.NDArray, npt.NDArray]:
    """Derive the stationary solution for a discrete potential of 1 to
    3 dimensions numerically.

    :param dx: Step size used for the coordinates. If the potential is
        one-dimensional, this can simply be the step size of the
        coordinate used in :param:`potential`.
        For larger dimensions, this parameter must be a tuple with the
        same size as the dimension.
    :param potential: A (multidimensional) potential.
    :param mass: The mass of the particle, used for the kinetic term.
    :param k: Number of eigenvalues and eigenstates to return.
    :returns: Tuple of eigenvalues and eigenvectors. Note that the
        eigenvectors are transposed so that they are of shape
        ``(k, np.size(potential))``
    """
    e_min = np.min(potential)
    sparse.save_npz(path + r"\hamiltonian.npz",discrete_hamiltonian(dx, potential, mass))
    return e_min

def nstationary_solution_mod(
    dx: tuple[float, ...] | float, potential: npt.NDArray, mass: float, k: int = 1
) -> tuple[npt.NDArray, npt.NDArray]:
    """
    Modified function to only compute solutions in the local minimum.

    Derive the stationary solution for a discrete potential of 1 to
    3 dimensions numerically.

    :param dx: Step size used for the coordinates. If the potential is
        one-dimensional, this can simply be the step size of the
        coordinate used in :param:`potential`.
        For larger dimensions, this parameter must be a tuple with the
        same size as the dimension.
    :param potential: A (multidimensional) potential.
    :param mass: The mass of the particle, used for the kinetic term.
    :param k: Number of eigenvalues and eigenstates to return.
    :returns: Tuple of eigenvalues and eigenvectors. Note that the
        eigenvectors are transposed so that they are of shape
        ``(k, np.size(potential))``
    """
    e_min = np.min(potential)
    eigenvals, eigenvects = linalg.eigh(
        discrete_hamiltonian(dx, potential, mass),subset_by_index=(0,k-1),subset_by_value=(e_min,np.inf)
    )
    return eigenvals, eigenvects.T