import numpy as np
from scipy.interpolate import interp1d
from scipy.linalg import block_diag, cholesky

def band_truncate_3d(M3d, n):
    """ 
    Truncate the bands of a 3D matrix M3d along the first two dimensions.

    Parameters
    ----------
    M3d : np.ndarray
        A 3D numpy array of shape (n, n, m) representing the spectral density matrices for `m` frequency bins.
    n : int
        The number of bands to keep in the truncation.  

    Returns
    -------
    M3d : np.ndarray
        A 3D numpy array of the same shape as `M3d`, where each slice is band-truncated.
        
    Notes
    -----
    Docstring is generated by GitHub Copilot.
    
    """
    for k in range(M3d.shape[2]):
        M3d[:,:,k] = band_truncate(M3d[:,:,k], n=n)
    
    return M3d


def band_truncate(M, n=1):
    """
    Truncate the bands of a 2D matrix M along the first two dimensions.
    
    Parameters
    ----------
    M : np.ndarray
        A 2D numpy array of shape (n, n) representing the spectral density matrix.
    n : int
        The number of bands to keep in the truncation.

    Returns
    -------
    M : np.ndarray
        A 2D numpy array of the same shape as `M`, where each slice is band-truncated.

    Notes
    -----
    Docstring is generated by GitHub Copilot.
    
    """
    if n is None:
        n = M.shape[0]
        
    Mt = M*0

    for ix in range(M.shape[0]):
        n1 = ix-n if (ix-n)>=0 else 0
        n2 = ix+n+1 if (ix+n)<=M.shape[0] else M.shape[0]

        Mt[ix, n1:n2] = M[ix, n1:n2]
    
    return Mt

def fft_time(omega, t0=0, n_fft=None):
    """
    Generate a time vector for FFT based on the frequency vector `omega`.
    
    Parameters
    ----------
    omega : np.ndarray
        A 1D numpy array representing the frequency vector in rad/s.
    t0 : float, optional

        The starting time for the time vector. Default is 0.
    n_fft : int, optional

        The number of points in the FFT. If not provided, it is set to the length of `omega`.

    Returns
    -------
    t : np.ndarray
        A 1D numpy array representing the time vector corresponding to the frequency vector `omega`.

    Notes
    -----
    Docstring is generated by GitHub Copilot.
    
    """
    
    if n_fft is None:
        n_fft = len(omega)
    domega = omega[1] - omega[0] 
    t = np.linspace(t0, np.pi*2/domega, n_fft)

    return t   

def spectrum_to_process(S, reg_factor=None, zero_limit=None):
    """
    Convert a spectral density matrix to its Cholesky factor process.
    This function takes a 3D array representing a spectral density matrix `S` and computes its 
    Cholesky factor for each frequency slice. Optional regularization can be applied to ensure positive definiteness, 
    and a zero limit can be set to skip decomposition for near-zero matrices.
   
    Parameters
    ----------
    S : np.ndarray
        A 3D array of shape (n, n, m) representing the spectral density matrices for `m` frequency bins.
    reg_factor : float, optional
        Regularization factor to be added to the diagonal of each frequency slice of `S` to ensure positive definiteness.
    zero_limit : float, optional
        Threshold below which the Cholesky decomposition is skipped for a frequency slice.

    Returns
    -------
    B : np.ndarray
        A 3D array of the same shape as `S`, where each slice is the lower-triangular Cholesky factor of the corresponding slice in `S`.

    Raises
    ------
    ValueError
        If regularization causes the norm of `S` to increase by more than 10%.

    Notes
    -----
    - The function modifies `S` in-place if regularization is applied.
    - Cholesky decomposition is only performed for slices where the maximum absolute value exceeds `zero_limit`.

    Docstring is generated by GitHub Copilot.
    """

    B = S*0     # copy S to chol factor
    norm_pre = np.linalg.norm(S)
    if reg_factor is not None:
        for k in range(S.shape[2]):
            S[:,:,k] = S[:,:,k] + np.linalg.norm(S[:,:,k]) * np.eye(S.shape[0]) * reg_factor
        
    if (np.linalg.norm(S)-norm_pre)/norm_pre>0.1:
        raise ValueError('Regularization causes significant increase in norm of S (10%). Please adjust.')
    
    # Cholesky decomposition
    for k in range(B.shape[2]):
        if zero_limit is None or np.max(np.max(np.abs(S[:,:,k])))>zero_limit:
            B[:,:,k] = cholesky(S[:, :, k], lower=True)
    return B
        

def simulate_mdof(S, omega, fs=None, tmax=None, reg_factor=None, zero_limit=1e-12, 
                  phase_angles=None, return_phase_angles=False, interpolation_kind='linear', 
                  component_scaling=None, print_status=False):
    
    '''
    Simulate time series from given cross-spectral density matrix, using FFT.

    Parameters
    ----------------
    S : float
        cross-spectral density matrix (Ndofs x Ndofs x Nfreqs) as complex numpy array
    omega : float
        numpy array defining frequencies in rad/s
    fs : float, optional
        sampling frequency in Hz; if not given, the value is defined as two times the maximum of `omega`
    tmax : float, optional
        duration of realization in seconds
    reg_factor : float, optional
        to help the Cholesky decomposition to achieve a numerical solution, 
        a diagonal matrix with the norm of all frequency components of the matrix S scaled by the given factor
        - if no value is given (--> None imposed), no regularization is conducted; 
        used as input to function `spectrum_to_process` which decomposes the spectral density 
    zero_limit : float, optional
        frequency components where the norm of S is below this limit is set to zero
        if no value is given (--> None imposed), no zero limit is applied; 
        used as input to function `spectrum_to_process` which decomposes the spectral density 
    phase_angles : float, optional
        if previously conducted simulations are to be recreated, they can by inputting the phase angles
        (to get the phase angles from simulation, see `return_phase_angles` input parameter)
    return_phase_angles : bool, default=False
        whether or not to return phase angle enabling recreation of the simulation
    interpolation_kind : optional, default='linear'
        interpolation type used for interpolation (to ensure that we get defined duration and sampling) prior to ifft
    component_scaling : float, optional
        values to use for scaling of components prior to decomposition (and thereafter rescale afterwards);
        can help if some components are much smaller than others; size should match number of DOFs in S
    print_status : False, optional
        whether or not to print status messages

    Returns
    ----------------
    p : float
        time history 
    t : float
        numpy array with time axis values corresponding to `p`
    alpha : float
        phase angles, only returned if `return_phase_angles` is True
    '''
    
    if omega[0] !=0:
        omega = np.insert(omega, 0, 0)
        S = np.dstack([0*S[:,:,0], S])
    
    if component_scaling is not None:
        c = np.tile(np.array(component_scaling), int(np.round(S.shape[0]/len(component_scaling))))
        scale_mat = c[np.newaxis,:].T @ c[np.newaxis, :]
        for k in range(S.shape[2]):
            S[:,:,k] = scale_mat * S[:,:,k]

    B = spectrum_to_process(S, reg_factor=reg_factor, zero_limit=zero_limit)  
    B = interp1d(omega, B, axis=2, fill_value=0, kind=interpolation_kind, bounds_error=False)

    n_dofs = S.shape[0]

    if fs is None:
        omega_max = omega[-1]
    else:
        omega_max = fs*(np.pi*2)
        if omega_max>omega[-1] and print_status:
            print('fs too high (larger than largest omega value in input) - zero padding automatically enforced')
        
    if tmax is None:
        domega = omega[1]-omega[0]
    else:
        domega = np.pi*2/tmax

    omega = np.arange(omega[0], omega_max+domega, domega)
    
    # Summation
    p = np.zeros([n_dofs, len(omega)])    
    Bi = B(omega)

    if phase_angles is None:
        alpha = 2*np.pi * np.random.rand(n_dofs, len(omega))
    else:
        alpha = phase_angles*1
        
    for j in range(n_dofs):   
        for m in range(j+1):
            p[j, :] = p[j, :] + (np.sqrt(2*domega) * 
                     np.real(np.fft.fft(Bi[j, m, :] * np.exp(1j*alpha[m, :]))))
        
    if component_scaling is not None:
        p = p/np.tile(c[:, np.newaxis], reps=p.shape[1])
        
    t = fft_time(omega)
    
    if return_phase_angles:
        return p, t, alpha
    else:
        return p, t


def simulate_mdof_direct(S, omega, reg_factor=None):
    """
    Simulate time series from given cross-spectral density matrix, using direct summation.

    Parameters
    ----------------
    S : float
        cross-spectral density matrix (Ndofs x Ndofs x Nfreqs) as complex numpy array
    omega : float
        numpy array defining frequencies in rad/s
    reg_factor : float, optional
        to help the Cholesky decomposition to achieve a numerical solution, 
        a diagonal matrix with the norm of all frequency components of the matrix S scaled by the given factor
        - if no value is given (--> None imposed), no regularization is conducted; 
        used as input to function `spectrum_to_process` which decomposes the spectral density

    Returns
    ----------------
    p : float
        time history
    t : float
        numpy array with time axis values corresponding to `p`      
    
    """
    B = spectrum_to_process(S, reg_factor)
    
    # Summation
    domega = omega[1] - omega[0]
    
    t = fft_time(omega)
    p = np.zeros([B.shape[0], len(t)])
    for j in range(B.shape[0]): 
        
        # Eq. 19 computed below for each j
        for m in range(j+1):
            for k, omega_k in enumerate(omega):
                phi = 2*np.pi * np.random.rand()
                p[j, :] = p[j, :] + np.sqrt(2*domega) * np.real(B[j, m, k] * np.exp(1j*(omega_k*t + phi)))
    
    return p, t
    