
import numpy as np
import pyccl as ccl
from tqdm import tqdm
from scipy import interpolate
import warnings
import copy
from itertools import product

from ..utils.Tabulate import _set_parameter
from ..utils.misc     import destory_Pk

__all__ = ['BaryonificationClass', 'Baryonification3D', 'Baryonification2D']

class BaryonificationClass(object):
    """
    Base class for implementing displacement function models.

    It takes in various input profiles and cosmological parameters to calculate the 
    displacement of particles/cells that is needed to convertfrom one matter distribution to another. 
    The class provides methods to set up interpolation tables for quick calculations of mass displacements.

    Parameters
    ----------
    DMO : object
        An instance of a dark matter-only profile, see `DarkMatterOnly`.
    DMB : object
        An instance of a dark matter-baryon profile, see `DarkMatterBaryon`.
    cosmo : object
        A CCL cosmology instance containing the cosmological parameters used for calculations.
    epsilon_max : float, optional
        The maximum displacement factor for the mass profile, in units of halo radius. Default is 20.
    mass_def : object, optional
        Mass definition object from CCL, default is `MassDef(200, 'critical')`.

    Notes
    -----
    The `BaryonificationClass` generates a displacement function that specifies 
    how baryonic processes alter the distribution of matter.

    **Key Methods and Workflow:**

    1. **`displacement()`**: Main method, which provides the computed displacements at a given radius
       from a given halo, at a given scale factor. It checks if the interpolation table is set up and 
       reads out the displacement using `_readout()`. It verifies that all required parameters are 
       provided and calculates the displacement function based on the inputs.

    2. **`setup_interpolator()`**: This method constructs interpolation tables for the displacement 
       function over a range of redshifts, masses, and radii. It allows for efficient computation 
       of mass displacements by precomputing and storing results. The method iterates over possible 
       values of input parameters, checks for validity, and constructs interpolators using the 
       `PchipInterpolator`.

       .. math::

           d(r, M, a) = f(\\log(1 + z), \\log(M), \\log(r), \\text{other parameters})

       This function ensures that the computed profiles adhere to constraints such as monotonically 
       increasing mass profiles and differences between DMO and DMB profiles asymptotically converging
       to zero on large-enough scales.

    3. **`get_masses()`**: This is an abstract method that must be implemented in subclasses. It is 
       responsible for calculating the mass profiles given a model, radii, mass, scale factor, and 
       mass definition. This method is expected to be overridden to provide specific mass calculations 
       based on the profile models.

    
    4. **`_readout()`**: This helper method reads out the displacement from the precomputed 
       interpolation table. It ensures that displacements are set to zero for radii beyond 
       `epsilon_max` times the halo radius to avoid unphysical values.


    **Normalization and Cutoff Handling:**

    - The `cutoff` value for the DMO and DMB profiles (see `SchneiderProfiles`) is set to 1 Gpc by default, which assumes that the 
      profiles are negligible beyond this scale. This helps to prevent numerical divergences during 
      FFT calculations while ensuring asymptotic behavior at large scales.
    - Only the real-space cutoff is modified here to prevent divergence; the projected cutoff remains as specified by the user. 

    **Warnings:**

    - The class issues warnings if the mass profiles are nearly constant over most of the radius range, 
      which suggests potential numerical issues or negative densities. Users are advised to adjust FFT 
      precision parameters such as `padding_lo_fftlog`, `padding_hi_fftlog`, or `n_per_decade` 
      if such warnings occur.
    """


    def __init__(self, DMO, DMB, cosmo, epsilon_max = 20, mass_def = ccl.halos.massdef.MassDef(200, 'critical')):
        
        self.DMO = DMO
        self.DMB = DMB
        
        #Set cutoff to 1 Gpc for calculation, assuming profiles are negligible beyond that
        #Smaller cutoffs result in asymptotic value problems at large scales
        #Larger cutoffs lead to numerical divergence during FFTLogs
        #The user supplied cutoffs will be places when implementing cutoffs in data
        #NOTE: We have not altered the PROJECTED cutoff, only the real cutoff.
        #Projected cutoff must be specified to user input at all times.
        self.DMO.set_parameter('cutoff', 1000)
        self.DMB.set_parameter('cutoff', 1000)
        
        self.cosmo       = cosmo #CCL cosmology instance
        self.epsilon_max = epsilon_max
        self.mass_def    = mass_def


    def get_masses(self, model, r, M, a):
        """
        Abstract method for calculating mass profiles.

        This method is intended to be overridden by subclasses to provide specific mass profile 
        calculations. It should return the mass profile for a given model, radii, mass, scale factor, 
        and mass definition.

        Parameters
        ----------
        model : object
            The model instance used to calculate the mass profile (e.g., DMO or DMB).
        r : array_like
            Radii at which to evaluate the mass profile, in comoving Mpc.
        M : array_like
            Halo mass or array of halo masses, in solar masses.
        a : float
            Scale factor, related to redshift by `a = 1 / (1 + z)`.

        Raises
        ------
        NotImplementedError
            This method must be implemented in subclasses of `BaryonificationClass`.
        """

        raise NotImplementedError("Implement a get_masses() method first")


    def setup_interpolator(self, 
                           z_min = 1e-2, z_max = 5, N_samples_z = 30, z_linear_sampling = False, 
                           M_min = 1e12, M_max = 1e16, N_samples_Mass = 30, 
                           R_min = 1e-3, R_max = 1e2, N_samples_R = 100, 
                           other_params = {}, verbose = True):
        
        """
        Sets up interpolation tables for the displacement function.

        This method constructs interpolation tables over a range of redshifts, halo masses, 
        and radii. It precomputes the displacement function values to facilitate efficient 
        calculations of mass displacements due to baryonic processes. User can compute it
        as a function of other inputs to `SchneiderProfiles` using the `other_params`
        function argument.

        Parameters
        ----------
        z_min : float, optional
            Minimum redshift for the interpolation. Default is 1e-2.
        z_max : float, optional
            Maximum redshift for the interpolation. Default is 5.
        N_samples_z : int, optional
            Number of redshift samples during tabulation. Default is 30.
        z_linear_sampling : bool, optional
            If True, use linear sampling for redshift; otherwise, use logarithmic sampling. Default is False.
            Useful if z_min = 0, as log spacing fails in that case.
        M_min : float, optional
            Minimum halo mass for the interpolation, in solar masses. Default is 1e12.
        M_max : float, optional
            Maximum halo mass for the interpolation, in solar masses. Default is 1e16.
        N_samples_Mass : int, optional
            Number of mass samples. Default is 30.
        R_min : float, optional
            Minimum radius for the interpolation, in comoving Mpc. Default is 1e-3.
        R_max : float, optional
            Maximum radius for the interpolation, in comoving Mpc. Default is 1e2.
        N_samples_R : int, optional
            Number of radius samples. Default is 100.
        other_params : dict, optional
            Additional parameters for model customization. To be provided in the format `{key : [list-like of vals]}`. 
            The default is an empty dictionary.
        verbose : bool, optional
            If True, display progress information using `tqdm`. Default is True.


        Notes
        -----
        - Ensures that mass profiles are monotonic and valid across the specified parameter ranges.
        - Issues warnings if profiles are nearly constant over radius or if the interpolation 
        fails for specific halo masses. Warnings tend to be raised for the lowest masses,
        especially when using pixel window convolutions.
        """

        M_range  = np.geomspace(M_min, M_max, N_samples_Mass)
        r        = np.geomspace(R_min, R_max, N_samples_R)
        z_range  = np.linspace(z_min, z_max, N_samples_z) if z_linear_sampling else np.geomspace(z_min, z_max, N_samples_z)
        p_keys   = list(other_params.keys()); setattr(self, 'p_keys', p_keys)
        d_interp = np.zeros([z_range.size, M_range.size, r.size] + [other_params[k].size for k in p_keys])
        
        #If other_params is empty then iterator will be empty and the code still works fine
        iterator = [p for p in product(*[np.arange(other_params[k].size) for k in p_keys])]
        
        with tqdm(total = d_interp.size//(M_range.size*r.size), desc = 'Building Table', disable = not verbose) as pbar:
            for j in range(z_range.size):
                
                for c in iterator:
                    
                    #Modify the model input params so that they are run with the right parameters
                    for k_i in range(len(p_keys)):
                        _set_parameter(self.DMO, p_keys[k_i], other_params[p_keys[k_i]][c[k_i]])
                        _set_parameter(self.DMB, p_keys[k_i], other_params[p_keys[k_i]][c[k_i]])
                    
                    
                    M_DMO = self.get_masses(self.DMO, r, M_range, 1/(1 + z_range[j]))
                    M_DMB = self.get_masses(self.DMB, r, M_range, 1/(1 + z_range[j]))
                    
                    for i in range(M_range.size):
                        ln_DMB    = np.log(M_DMB[i])
                        ln_DMO    = np.log(M_DMO[i])
                        
                        #Require mass to always increase w/ radius
                        #And remove pts of DMO = DMB, improves large-scale convergence
                        #And require at least 1e-6 difference else the interpolator breaks :/
                        
                        min_diff  = -np.inf
                        diff_mask = np.ones_like(ln_DMB).astype(bool)
                        iterate   = 0
                        while (min_diff < 1e-5) & (diff_mask.sum() > 5):
                            
                            new_mask  = ( (np.diff(ln_DMB[diff_mask], prepend = 0) > 1e-5) & 
                                          (np.diff(ln_DMO[diff_mask], prepend = 0) > 1e-5) & 
                                          (np.abs(ln_DMB - ln_DMO)[diff_mask] > 1e-6) 
                                        )
                            
                            diff_mask[diff_mask] = new_mask
                            diff_mask[0] = True
                            
                            iterate += 1
                            
                            if iterate > 30:
                                diff_mask  = np.zeros_like(diff_mask).astype(bool) #Set everything to False and skip the building step next
                                warn_text  = (f"Mass profile of log10(M) = {np.log10(M_range[i])} is nearly constant over radius. " 
                                              "Suggests density is negative or zero for most of the range. If using convolutions,"
                                              "consider changing the fft precision params in the CCL profile:"
                                              "padding_lo_fftlog, padding_hi_fftlog, or n_per_decade")
                                warnings.warn(warn_text, UserWarning)
                                break
                                
                            if diff_mask.sum() < 5: 
                                warn_text  = (f"Mass profile of log10(M) = {np.log10(M_range[i])} is nearly constant over radius. " 
                                              "Or it is broken. Less than 5 datapoints are usable.")
                                warnings.warn(warn_text, UserWarning)
                                break
                            
                            min_diff  = np.min([np.min(np.diff(ln_DMB[diff_mask], prepend = 0)[1:]),
                                                np.min(np.diff(ln_DMO[diff_mask], prepend = 0)[1:])
                                               ])                                                          
                            
                        #If we have enough usable mass values, then proceed as usual
                        #This generally breaks for very small halos, where projection
                        #can be catastrophicall broken (eg. only negative densities)
                        if diff_mask.sum() > 5:
                                   
                            interp_DMB = interpolate.PchipInterpolator(ln_DMB[diff_mask], np.log(r)[diff_mask], extrapolate = False)
                            interp_DMO = interpolate.PchipInterpolator(np.log(r)[diff_mask], ln_DMO[diff_mask], extrapolate = False)

                            offset = np.exp(interp_DMB(interp_DMO(np.log(r)))) - r
                            offset = np.where(np.isfinite(offset), offset, 0)
                        
                        #If broken, then these halos contribute nothing to the displacement function.
                        #Just provide a warning saying this is happening
                        else:
                            offset = np.zeros_like(r)
                            warn_text = (f"Displacement function for halo with log10(M) = {np.log10(M_range[i])} failed to compute." 
                                         "Defaulting to d = 0. Consider changing the fft precision params in the CCL profile:"
                                         "padding_lo_fftlog, padding_hi_fftlog, or n_per_decade")
                            warnings.warn(warn_text, UserWarning)
                        
                        #Build a custom index into the array
                        index = tuple([j, i, slice(None)] + list(c))
                        d_interp[index] = offset
                            
                    pbar.update(1)


        input_grid = tuple([np.log(1 + z_range), np.log(M_range), np.log(r)] + [other_params[k] for k in p_keys])

        self.raw_input_d = d_interp
        self.raw_input_z_range = np.log(1 + z_range)
        self.raw_input_M_range = np.log(M_range)
        self.raw_input_r_range = np.log(r)
        for k in other_params.keys(): setattr(self, 'raw_input_%s_range' % k, other_params[k]) #Save other raw inputs too
            
        self.interp_d = interpolate.RegularGridInterpolator(input_grid, d_interp, bounds_error = False, fill_value = np.nan)    

        #Once all tabulation is done, we don't need to keep P(k) calculations in cosmology object.
        #This is good because the Pk class is not pickleable, so by destorying it here we
        #are able to keep this class pickleable.
        self.cosmo = destory_Pk(self.cosmo)

    
    def _readout(self, r, M, a, **kwargs):

        """
        Read out the displacement from the interpolation table.

        This method retrieves the displacement function values from the precomputed 
        interpolation table for given radii, halo masses, and scale factor. It sets 
        displacement values to zero beyond `r > R * epsilon_max`, where `R` is the
        halo radius.

        Parameters
        ----------
        r : array_like
            Radii at which to evaluate the displacement, in comoving Mpc.
        M : array_like
            Halo mass or array of halo masses, in solar masses.
        a : float
            Scale factor, related to redshift by `a = 1 / (1 + z)`.
        **kwargs : dict
            Additional parameters required by the interpolation table.

        Returns
        -------
        displ : ndarray
            Displacement values corresponding to the input radii and masses.

        Notes
        -----
        - Ensures that displacements are set to zero for radii beyond `epsilon_max` 
        times the halo radius to avoid unphysical values.
        """
        
        table = self.interp_d #The interpolation table to use
        r_use = np.atleast_1d(r)
        M_use = np.atleast_1d(M)
        a_use = np.atleast_1d(a)
        z_use = 1/a_use - 1
        
        displ = np.zeros([M_use.size, r_use.size])
        empty = np.ones_like(r_use)
        z_in  = np.log(1/a)*empty #This is log(1 + z)
        r_in  = np.log(r_use)
        k_in  = [kwargs[k] * empty for k in kwargs.keys()]
        
        for i in range(M_use.size):
            M_in  = np.log(M_use[i])*empty
            p_in  = tuple([z_in, M_in, r_in] + k_in)
            displ[i] = table(p_in)
            
            R        = self.mass_def.get_radius(self.cosmo, np.atleast_1d(M), a)/a #in comoving Mpc
            inside   = (r < self.epsilon_max*R)
            displ[i] = np.where(inside, displ, 0) #Set large-scale displacements to 0
            
        #Handle dimensions so input dimensions are mirrored in the output
        if np.ndim(r) == 0:
            displ = np.squeeze(displ, axis=-1)
        if np.ndim(M) == 0:
            displ = np.squeeze(displ, axis=0)
            
        return displ

    
    def displacement(self, r, M, a, **kwargs):
        """
        Return the displacement needed to convert the matter distribution
        from the dmo profiles to the dmb profiles, for the given radii, halo 
        masses, and scale factor. It never does a calculation on-the-fly and
        instead reads out a precomputed table.

        Parameters
        ----------
        r : array_like
            Radii at which to evaluate the displacement, in comoving Mpc.
        M : array_like
            Halo mass or array of halo masses, in solar masses.
        a : float
            Scale factor, related to redshift by `a = 1 / (1 + z)`.
        **kwargs : dict
            Additional parameters required by the interpolation table.

        Returns
        -------
        displ : ndarray
            Displacement values corresponding to the input radii and masses.
            In comoving Mpc.

        Raises
        ------
        NameError
            If the interpolation table has not been set up using `setup_interpolator()`.
        AssertionError
            If required parameters are not provided in `kwargs`.
        """
        
        if not hasattr(self, 'interp_d'):
            raise NameError("No Table created. Run setup_interpolator() method first")
            
        for k in self.p_keys:
            assert k in kwargs.keys(), "Need to provide %s as input into `displacement'. Table was built with this." % k
        
        return self._readout(r, M, a, **kwargs)



class Baryonification3D(BaryonificationClass):
    """
    Class implementing a 3D baryonification model. It extends 
    `BaryonificationClass` to provide the displacement function 
    in three dimensions.

    Inherits from
    -------------
    BaryonificationClass : Base class for baryonification models.

    Notes
    -----
    The `Baryonification3D` class is used to compute the 3D enclosed mass profiles.
    It integrates the density profile to obtain the cumulative 
    enclosed mass as a function of radius:

    .. math::

        M_{\\text{enc}}(r) = 4\\pi \\int_0^r \\rho(r') r'^2 \\, d\\ln r'

    the displacement function is then,

    .. math::

        \Delta d (r) = M_{\rm dmb}^{-1}(M_{\rm dmo}(r)) - r,

    
    Methods
    -------
    displacement(r, M, a, **kwargs)
        Compute the displacement function for a given mass, radii, and scale factor
    get_masses(model, r, M, a)
        Computes the enclosed mass profile for a given model, radii, halo mass, and scale factor.
    """

    def get_masses(self, model, r, M, a):
        """
        Computes the enclosed mass profile for a given model.

        This method calculates the cumulative enclosed mass profile for a specified 
        model (e.g., dark matter-only or dark matter-baryon), radii, halo mass, and 
        scale factor.

        Parameters
        ----------
        model : object
            The model instance used to calculate the mass profile (e.g., DMO or DMB).
        r : array_like
            Radii at which to evaluate the mass profile, in comoving Mpc.
        M : float or array_like
            Halo mass or array of halo masses, in solar masses.
        a : float
            Scale factor, related to redshift by `a = 1 / (1 + z)`.

        Returns
        -------
        M_f : ndarray
            Enclosed mass profile corresponding to the input radii and halo mass. 
            The output shape is (len(M), len(r)), unless M is a scalar, in which 
            case the shape is (len(r),). Units of solar masses.

        Notes
        -----
        - The method adjusts the integration range to ensure that the minimum and 
        maximum radii do not disrupt the integral, adding a 20% buffer.
        - Negative densities are set to zero to avoid unphysical results.
        - The enclosed mass \( M_{\\text{enc}}(r) \) is calculated using:

        .. math::

            M_{\\text{enc}}(r) = 4\\pi \\int_0^r \\rho(r') r'^2 \\, d\\ln r'

        - `PchipInterpolator` is used to interpolate the mass profile, ensuring smoothness 
        and avoiding artifacts due to Fourier space ringing.
        - If `M` is a scalar, the output is squeezed to remove extra dimensions.

        Examples
        --------
        Compute the enclosed mass profile for a given model:

        >>> baryon_model = Baryonification3D(DMO=dark_matter_profile, DMB=dmb_profile, cosmo=my_cosmology)
        >>> r = np.logspace(-2, 1, 50)  # Radii in comoving Mpc
        >>> M = 1e14  # Halo mass in solar masses
        >>> a = 0.8  # Scale factor corresponding to redshift z
        >>> mass_profile = baryon_model.get_masses(baryon_model.DMO, r, M, a)
        """
        
        #Make sure the min/max does not mess up the integral
        #Adding some 20% buffer just in case
        r_min = np.min([np.min(r), 1e-6])
        r_max = np.max([np.max(r), 1000])
        r_int = np.geomspace(r_min/1.2, r_max*1.2, 500)
        
        dlnr  = np.log(r_int[1]/r_int[0])
        rho   = model.real(self.cosmo, r_int, M, a)
        rho   = np.where(rho < 0, 0, rho) #Enforce non-zero densities
        
        if isinstance(M, (float, int) ): rho = rho[None, :]
            
        M_enc = np.cumsum(4*np.pi*r_int**3 * rho * dlnr, axis = -1)
        lnr   = np.log(r)
        
        M_f   = np.zeros([M_enc.shape[0], r.size])
        
        #Remove datapoints in profile where rho == 0 and then just interpolate
        #across them. This helps deal with ringing profiles due to 
        #fourier space issues, where profile could go negative sometimes
        for M_i in range(M_enc.shape[0]):
            Mask     = (rho[M_i] > 0) & (np.isfinite(M[M_i])) #Keep only finite points, and ones with increasing density
            M_f[M_i] = np.exp( interpolate.PchipInterpolator(np.log(r_int)[Mask], np.log(M_enc[M_i])[Mask], extrapolate = False)(lnr) )
        
        if isinstance(M, (float, int) ): M_f = np.squeeze(M_f, axis = 0)
            
        return M_f


class Baryonification2D(BaryonificationClass):

    """
    Class implementing a 2D baryonification model. It extends 
    `BaryonificationClass` to provide the displacement function 
    in two dimensions.

    Inherits from
    -------------
    BaryonificationClass : Base class for baryonification models.

    Notes
    -----
    The `Baryonification2D` class is used to compute the 2D mass profiles. 
    It integrates the projected surface density profile, \( \Sigma(r) \), to obtain 
    the cumulative enclosed mass as a function of radius. 

    The enclosed mass \( M_{\\text{enc}}(r) \) is calculated by integrating the projected 
    surface density profile \( \Sigma(r) \) over circular annuli:

    .. math::

        M_{\\text{enc}}(r) = 2\\pi \\int_0^r \Sigma(r') r' \, d\\ln r'

    the displacement function is then

    .. math::

        \Delta d (r_{\\rm p}) = M_{\\rm DMB, p}^{-1}(M_{\\rm DMO, p}(r_{\\rm p})) - r_{\\rm p},

    where :math:`r_{\\rm p}` is the projected radius.

    """

    def get_masses(self, model, r, M, a):
        """
        Computes the enclosed mass profile for a given model using 2D projection.

        This method calculates the cumulative enclosed mass profile for a specified 
        model (e.g., dark matter-only or dark matter-baryon), radii, halo mass, and 
        scale factor.

        Parameters
        ----------
        model : object
            The model instance used to calculate the mass profile (e.g., DMO or DMB).
        r : array_like
            Radii at which to evaluate the mass profile, in comoving Mpc.
        M : float or array_like
            Halo mass or array of halo masses, in solar masses.
        a : float
            Scale factor, related to redshift by `a = 1 / (1 + z)`.

        Returns
        -------
        M_f : ndarray
            Enclosed mass profile corresponding to the input radii and halo mass. 
            The output shape is (len(M), len(r)), unless M is a scalar, in which 
            case the shape is (len(r),). In units of solar masses.

        Notes
        -----
        - The method adjusts the integration range to ensure that the minimum and 
        maximum radii do not disrupt the integral, adding a 20% buffer.
        - Negative surface densities are set to zero to avoid unphysical results.
        - The enclosed mass \( M_{\\text{enc}}(r) \) is calculated using:

        .. math::

            M_{\\text{enc}}(r) = 2\\pi \\int_0^r \\Sigma(r') r' \\, d\\ln r'

        - `PchipInterpolator` is used to interpolate the mass profile, ensuring smoothness 
        and avoiding artifacts due to Fourier space ringing.
        - If `M` is a scalar, the output is squeezed to remove extra dimensions.

        Examples
        --------
        Compute the enclosed mass profile for a given model using 2D projection:

        >>> baryon_model = Baryonification2D(DMO=dark_matter_profile, DMB=dmb_profile, cosmo=my_cosmology)
        >>> r = np.logspace(-2, 1, 50)  # Radii in comoving Mpc
        >>> M = 1e14  # Halo mass in solar masses
        >>> a = 0.5  # Scale factor corresponding to redshift z
        >>> mass_profile = baryon_model.get_masses(baryon_model.DMO, r, M, a)
        """
        
        #Make sure the min/max does not mess up the integral
        #Adding some 20% buffer just in case
        r_min = np.min([np.min(r), 1e-6])
        r_max = np.max([np.max(r), 1000])
        r_int = np.geomspace(r_min/1.5, r_max*1.5, 500)
        
        #The scale fac. is used in Sigma cause the projection in ccl is
        #done in comoving coords not physical coords
        dlnr  = np.log(r_int[1]/r_int[0])
        Sigma = model.projected(self.cosmo, r_int, M, a) * a 
        Sigma = np.where(Sigma < 0, 0, Sigma) #Enforce non-zero densities
        
        if isinstance(M, (float, int) ): Sigma = Sigma[None, :]
        
        M_enc = np.cumsum(2*np.pi*r_int**2 * Sigma * dlnr, axis = -1)
        lnr   = np.log(r)
        
        
        M_f  = np.zeros([M_enc.shape[0], r.size])
        #Remove datapoints in profile where Sigma == 0 and then just interpolate
        #across them. This helps deal with ringing profiles due to 
        #fourier space issues, where profile could go negative sometimes
        for M_i in range(M_enc.shape[0]):
            Mask     = (Sigma[M_i] > 0) & (np.isfinite(M_enc[M_i])) #Keep only finite points, and ones with increasing density
            M_f[M_i] = np.exp( interpolate.PchipInterpolator(np.log(r_int)[Mask], np.log(M_enc[M_i])[Mask], extrapolate = False)(lnr) )
        
        if isinstance(M, (float, int) ): M_f = np.squeeze(M_f, axis = 0)
            
        return M_f