"""
Eigenvalues and eigenmodes functions
"""
import numpy as np
from lapy import Solver
from scipy.special import sph_harm

def calc_eigs(fem, num_modes):
    """Calculate the eigenvalues and eigenmodes of a surface.

    Parameters
    ----------
    fem : lapy compatible finite element mesh
        Finite element mesh
    num_modes : int
        Number of eigenmodes to be calculated

    Returns
    ------
    evals : array (num_modes x 1)
        Eigenvalues
    emodes : array (number of surface points x num_modes)
        Eigenmodes
    """
    
    evals, emodes = fem.eigs(k=num_modes)
    
    return evals, emodes

def _get_eigengroups(eigs):
    """
    Helper function to find eigengroups
    """
    print("""IMPORTANT: EIGENMODES MUST NOT BE TRUNCATED AT FIRST NON-ZERO MODE
    FOR THIS FUNCTION TO WORK""")
    lam = eigs.shape[1] # number of eigenmodes
    l = np.floor((lam-1)/2).astype(int)    
    if lam == 1:
        return np.asarray([0])
    if lam == 2:
        groups = [np.zeros(1).astype(int)]
        groups.append(np.ones(1).astype(int))
        return groups
    
    groups = []
    ii = 0
    i = 0
    for g in range(l+1):
        ii += 2*g+1
        if ii >= lam:
            groups.append(np.arange(i,lam))
            return groups
        groups.append(np.arange(i,ii))
        i = ii

def reconstruct_data(coeffs, eigenmodes, pv=100.0):
    """
    Reconstruct a dataset of `n_vertices` given a set of eigenmodes and coeffs
    conditioned on data using ordinary least squares (OLS)

    Parameters
    ----------
    coeffs : np.ndarray of shape (M,)
        Coefficients output from fitting OLS
    eigenmodes : np.ndarray of shape (n_vertices, M)
        Eigenmodes of `n_vertices` by number of eigenvalues M
    pv : float (default 100), rounds up to nearest index
        Percentage of `eigenmodes` to use for reconstruction (default 100%)

    Returns
    -------
    new_data : np.ndarray of (n_vertices,)
        Reconstructed data

    """
    
    if pv == 0.0:
        raise ValueError("Percentage of modes to use for reconstruction must be greater than 1")
        
    _, M = eigenmodes.shape
    n = int(np.ceil(M * pv / 100))
    
    eigenmodes = eigenmodes[:,:n]
    coeffs = coeffs[:n].reshape(-1, 1)
    
    new_data = coeffs.T @ eigenmodes.T
    
    return new_data.squeeze()
    
    
def eigen_decomposition(data, eigenmodes, method='matrix'):
    """
    Decompose data using eigenmodes and calculate the coefficient of 
    contribution of each vector
    
    Parameters:
    -----------
    data : np.ndarray of shape (n_vertices, 3)
        N = number of vertices, P = columns of independent data
    eigenmodes : np.ndarray of shape (n_vertices, M)
        N = number of vertices, M = number of eigenmodes
    method : string
        method of calculation of coefficients: 'matrix', 'matrix_separate', 
        'regression'
    
    Returns:
    -------
    coeffs : numpy array of shape (N, 3)
     coefficient values
    
    """
    
    if data.ndim > 1:
        N, P = data.shape
    else:
        P = 1
    
    _, M = eigenmodes.shape
    
    if method == 'matrix':
        print("Using matrix decomposition to reconstruct data")
        coeffs = np.linalg.solve((eigenmodes.T @ eigenmodes), (eigenmodes.T @ data))
            
    elif method == 'regression':
        print("Using regression method to reconstruct data")
        coeffs = np.zeros((M, P))
        if P > 1:
            for p in range(P):
                coeffs[:, p] = np.linalg.lstsq(eigenmodes, data[:, p], rcond=None)[0]
        else:
            coeffs = np.linalg.lstsq(eigenmodes, data, rcond=None)[0]
            
    else:
        raise ValueError("Accepted methods for decomposition are 'matrix', and 'regression'")
                
    return coeffs
    
    
def compute_axes_ellipsoid(eigenvalues):
    """
    Compute the axes of an ellipsoid given the eigenmodes.
    """    
    return np.sqrt(eigenvalues)
    

def transform_to_spheroid(eigenvalues, eigenmodes):
    """
    Transform the eigenmodes to a spheroid space
    """
    ellipsoid_axes = compute_axes_ellipsoid(eigenvalues)
    #ellipsoid_axes = ellipsoid_axes.reshape(-1, 1)
    
    spheroid_eigenmodes = eigenmodes / ellipsoid_axes
    
    return spheroid_eigenmodes
    
    
def transform_to_ellipsoid(eigenvalues, eigenmodes):
    """
    Transform the eigenmodes in spheroid space back to ellipsoid by stretching
    """
    
    ellipsoid_axes = compute_axes_ellipsoid(eigenvalues)
    
    ellipsoid_eigenmodes = eigenmodes * ellipsoid_axes
    
    return ellipsoid_eigenmodes


def resample_spheroid(spheroid_eigenmodes, angle):
    """
    Resample the N-D hypersphere generated by the N orthogonal unit modes

    """
    # ensure the eigenmodes are normalized on the unit hypersphere
    spheroid_eigenmodes = spheroid_eigenmodes #/ np.linalg.norm(spheroid_eigenmodes, axis=0)
    
    # length of group to determine evenness or oddness of dimensions
    #dims = spheroid_eigenmodes.shape[0]
    
    # initialize the new points p
    p = spheroid_eigenmodes
    
    # compute the coordinates for the new points
    for i in range(1, spheroid_eigenmodes.shape[1]):
        p[:, i] *= np.cos(angle)
    
    # Compute the coordinates for new points p
    # print("Computing the coordinates for each dimension, multiplying by single angle per group")
    # if dims % 2 == 0: # EVEN
    #     p *= np.cos(angle)
    # else: #dims % 2 == 1 ODD
    #     p *= np.sin(angle)
    
    # ensure that the unit modes are orthonormal (QR decomposition)
    new_modes = p
    
    return new_modes

def _cartesian_to_spherical(x, y, z):
    """
    Convert Cartesian coordinates to spherical coordinates.
    
    Parameters:
    x, y, z: Cartesian coordinates
    
    Returns:
    r: radius
    theta: azimuthal angle in radians
    phi: polar angle in radians
    """
    r = np.sqrt(x**2 + y**2 + z**2)  # radius
    theta = np.arctan2(y, x)  # azimuthal angle
    phi = np.arccos(z / r)  # polar angle

    # Adjust theta to be in range [0, 2pi)
    theta = theta if theta >= 0 else theta + 2 * np.pi

    return r, theta, phi

def compute_spherical_harmonics(l, m, theta, phi):
    """
    Function to compute the spherical harmonics Y_{lm}(theta, phi)
    
    Parameters:
    l: int - The degree of the spherical harmonic, should be >= 0.
    m: int - The order of the spherical harmonic, should be -l <= m <= l.
    theta: float - The azimuthal (longitudinal) coordinate, 0 <= theta < 2*pi.
    phi: float - The polar (colatitudinal) coordinate, 0 <= phi <= pi.

    Returns:
    The value of the spherical harmonic Y_{lm}(theta, phi)
    """
    # Compute the spherical harmonics
    Y_lm = sph_harm(m, l, theta, phi)
    return Y_lm