<commit_before>#! /usr/bin/env python

"""
Module with functions for correcting bad pixels in cubes.
"""


__author__ = '<NAME>, <NAME>'
__all__ = ['frame_fix_badpix_isolated',
           'cube_fix_badpix_isolated',
           'cube_fix_badpix_annuli',
           'cube_fix_badpix_clump',
           'cube_fix_badpix_ifs',
           'cube_fix_badpix_with_kernel']

import numpy as np
from skimage.draw import disk, ellipse
from scipy.ndimage import median_filter
from astropy.stats import sigma_clipped_stats
from ..stats import sigma_filter
from ..var import frame_center, get_annulus_segments, frame_filter_lowpass
from ..stats import clip_array
from ..config import timing, time_ini, Progressbar
from .rescaling import find_scal_vector, frame_rescaling

import warnings
try:
    from numba import njit
    no_numba = False
except ImportError:
    msg = "Numba python bindings are missing."
    warnings.warn(msg, ImportWarning)
    no_numba = True


def frame_fix_badpix_isolated(array, bpm_mask=None, sigma_clip=3, num_neig=5,
                              size=5, protect_mask=0, cxy=None, mad=False,
                              ignore_nan=True, verbose=True, full_output=False):
    """ Corrects the bad pixels, marked in the bad pixel mask. The bad pixel is
    replaced by the median of the adjacent pixels. This function is very fast
    but works only with isolated (sparse) pixels.

    Parameters
    ----------
    array : numpy ndarray
        Input 2d array.
    bpm_mask : numpy ndarray, optional
        Input bad pixel map. Zeros frame where the bad pixels have a value of
        1. If None is provided a bad pixel map will be created using
        sigma clip statistics.
    sigma_clip : int, optional
        In case no bad pixel mask is provided all the pixels above and below
        sigma_clip*STDDEV will be marked as bad.
    num_neig : int, optional
        The side of the square window around each pixel where the sigma clipped
        statistics are calculated (STDDEV and MEDIAN). If the value is equal to
        0 then the statistics are computed in the whole frame.
    size : odd int, optional
        The size the box (size x size) of adjacent pixels for the median
        filter.
    protect_mask : int or float, optional
        If larger than 0, radius of a circular aperture (at the center of the
        frames) in which no bad pixels will be identified. This can be useful
        to protect the star and vicinity.
    cxy: None or tuple
        If protect_mask is True, this is the location of the star centroid in
        the images. If None, assumes the star is already centered. If a tuple,
        the location of the star is assumed to be the same in all frames of the
        cube.
    mad : {False, True}, bool optional
        If True, the median absolute deviation will be used instead of the
        standard deviation.
    ignore_nan: bool, optional
        Whether to not consider NaN values as bad pixels. If False, will also
        correct them.
    verbose : bool, optional
        If True additional information will be printed.
    full_output: bool, {False,True}, optional
        Whether to return as well the cube of bad pixel maps and the cube of
        defined annuli.

    Return
    ------
    frame : numpy ndarray
        Frame with bad pixels corrected.
    bpm_mask: 2d array
        The bad pixel map
    """
    if array.ndim != 2:
        raise TypeError('Array is not a 2d array or single frame')
    if size % 2 == 0:
        raise TypeError('Size of the median blur kernel must be an odd integer')

    if bpm_mask is not None:
        bpm_mask = bpm_mask.astype('bool')

    if verbose:
        start = time_ini()

    if num_neig > 0:
        neigh = True
    else:
        neigh = False

    frame = array.copy()
    if cxy is None:
        cy, cx = frame_center(frame)
    else:
        cx, cy = cxy

    if bpm_mask is None:
        ori_nan_mask = np.where(np.isnan(frame))
        ind = clip_array(frame, sigma_clip, sigma_clip, neighbor=neigh,
                         num_neighbor=num_neig, mad=mad)
        bpm_mask = np.zeros_like(frame)
        bpm_mask[ind] = 1
        if ignore_nan:
            bpm_mask[ori_nan_mask] = 0
        if protect_mask:
            cir = disk((cy, cx), protect_mask, shape=bpm_mask.shape)
            bpm_mask[cir] = 0
        bpm_mask = bpm_mask.astype('bool')

    smoothed = median_filter(frame, size, mode='mirror')
    frame[np.where(bpm_mask)] = smoothed[np.where(bpm_mask)]
    array_out = frame
    count_bp = np.sum(bpm_mask)

    if verbose:
        msg = "/nDone replacing {} bad pixels using the median of neighbors"
        print(msg.format(count_bp))
        timing(start)

    if full_output:
        return array_out, bpm_mask
    else:
        return array_out


def cube_fix_badpix_isolated(array, bpm_mask=None, sigma_clip=3, num_neig=5,
                             size=5, frame_by_frame=False, protect_mask=0,
                             cxy=None, mad=False, ignore_nan=True, verbose=True,
                             full_output=False):
    """ Corrects the bad pixels, marked in the bad pixel mask. The bad pixel is
    replaced by the median of the adjacent pixels. This function is very fast
    but works only with isolated (sparse) pixels.

    Parameters
    ----------
    array : numpy ndarray
        Input 3d array.
    bpm_mask : numpy ndarray, optional
        Input bad pixel map. Zeros frame where the bad pixels have a value of 1.
        If None is provided a bad pixel map will be created per frame using
        sigma clip statistics.
    sigma_clip : int, optional
        In case no bad pixel mask is provided all the pixels above and below
        sigma_clip*STDDEV will be marked as bad.
    num_neig : int, optional
        The side of the square window around each pixel where the sigma clipped
        statistics are calculated (STDDEV and MEDIAN). If the value is equal to
        0 then the statistics are computed in the whole frame.
    size : odd int, optional
        The size the box (size x size) of adjacent pixels for the median filter.
    frame_by_frame: bool, optional
        Whether to correct bad pixels frame by frame in the cube. By default it
        is set to False; the bad pixels are computed on the mean frame of the
        stack (faster but not necessarily optimal).
    protect_mask : int or float, optional
        If larger than 0, radius of a circular aperture (at the center of the
        frames) in which no bad pixels will be identified. This can be useful
        to protect the star and vicinity.
    cxy: None, tuple or 2d numpy ndarray
        If protect_mask is True, this is the location of the star centroid in
        the images. If None, assumes the star is already centered. If a tuple,
        the location of the star is assumed to be the same in all frames of the
        cube. If a (n_frames x 2) ndarray, it should contain the xy location of
        the star in each frame.
    mad : {False, True}, bool optional
        If True, the median absolute deviation will be used instead of the
        standard deviation.
    ignore_nan: bool, optional
        Whether to not consider NaN values as bad pixels. If False, will also
        correct them.
    verbose : bool, optional
        If True additional information will be printed.
    full_output: bool, {False,True}, optional
        Whether to return as well the cube of bad pixel maps and the cube of
        defined annuli.

    Return
    ------
    array_out : numpy ndarray
        Cube with bad pixels corrected.
    bpm_mask: 2d or 3d array [if full_output is True]
        The bad pixel map or the cube of bpix maps
    """
    if array.ndim != 3:
        raise TypeError('Array is not a 3d array or cube')
    if size % 2 == 0:
        raise TypeError('Size of the median blur kernel must be an odd integer')

    if bpm_mask is not None:
        bpm_mask = bpm_mask.astype('bool')

    if verbose:
        start = time_ini()

    if num_neig > 0:
        neigh = True
    else:
        neigh = False

    nz = array.shape[0]

    if cxy is None:
        cy, cx = frame_center(array[0])
    elif isinstance(cxy, tuple):
        cx, cy = cxy
    elif isinstance(cxy, np.ndarray):
        if cxy.shape[0] != nz or cxy.shape[1] != 2 or cxy.ndim != 2:
            raise ValueError("cxy does not have right shape")
        elif not frame_by_frame:
            msg = "cxy must be a tuple or None if not in frame_by_frame mode"
            raise ValueError(msg)
        else:
            cx = cxy[:, 0]
            cy = cxy[:, 1]

    array_out = array.copy()
    final_bpm = np.zeros_like(array_out, dtype=bool)
    n_frames = array.shape[0]
    count_bp = 0
    if frame_by_frame:
        if np.isscalar(cx):
            cx = [cx]*nz
            cy = [cy]*nz
        if bpm_mask is not None:
            if bpm_mask.ndim == 2:
                bpm_mask = [bpm_mask]*n_frames
                bpm_mask = np.array(bpm_mask)
        for i in Progressbar(range(n_frames), desc="processing frames"):
            if bpm_mask is not None:
                bpm_mask_tmp = bpm_mask[i]
            else:
                bpm_mask_tmp = None
            res = frame_fix_badpix_isolated(array[i], bpm_mask=bpm_mask_tmp,
                                            sigma_clip=sigma_clip,
                                            num_neig=num_neig, size=size,
                                            protect_mask=protect_mask,
                                            verbose=False, cxy=(cx[i], cy[i]),
                                            ignore_nan=ignore_nan,
                                            full_output=True)
            array_out[i] = res[0]
            final_bpm[i] = res[1]
        count_bp = np.sum(final_bpm)
    else:
        if bpm_mask is None:
            ori_nan_mask = np.where(np.isnan(np.nanmean(array, axis=0)))
            ind = clip_array(np.nanmean(array, axis=0), sigma_clip, sigma_clip,
                             neighbor=neigh, num_neighbor=num_neig, mad=mad)
            final_bpm = np.zeros_like(array[0], dtype=bool)
            final_bpm[ind] = 1
            if ignore_nan:
                final_bpm[ori_nan_mask] = 0
            if protect_mask:
                cir = disk((cy, cx), protect_mask, shape=final_bpm.shape)
                final_bpm[cir] = 0
            final_bpm = final_bpm.astype('bool')
        else:
            if bpm_mask.ndim == 3:
                final_bpm = np.median(bpm_mask, axis=0)
            else:
                final_bpm = bpm_mask.copy()

        for i in Progressbar(range(n_frames), desc="processing frames"):
            frame = array_out[i]
            smoothed = median_filter(frame, size, mode='mirror')
            frame[np.where(final_bpm)] = smoothed[np.where(final_bpm)]
            if verbose:
                count_bp += np.sum(final_bpm)

    if verbose:
        msg = "Done replacing {:.0f} bad pixels using the median of neighbors"
        print(msg.format(count_bp))
        if not frame_by_frame:
            msg = "(i.e. {:.0f} static bad pixels per channel))"
            print(msg.format(count_bp/n_frames))
        timing(start)

    if full_output:
        return array_out, final_bpm
    else:
        return array_out


def cube_fix_badpix_annuli(array, fwhm, cy=None, cx=None, sig=5.,
                           protect_mask=0, r_in_std=50, r_out_std=None,
                           verbose=True, half_res_y=False, min_thr=None,
                           max_thr=None, full_output=False):
    """
    Function to correct the bad pixels annulus per annulus (centered on the
    provided location of the star), in an input frame or cube.
    This function is faster than bp_clump_removal; hence to be prefered in all
    cases where there is only one bright source with circularly symmetric PSF.
    The bad pixel values are replaced by: ann_median + ann_stddev*random_gauss;
    where ann_median is the median of the annulus, ann_stddev is the standard
    deviation in the annulus, and random_gauss is a random factor picked from a
    gaussian distribution centered on 0 and with variance 1.

    Parameters
    ----------
    array : 3D or 2D array
        Input 3d cube or 2d image.
    fwhm: float or 1D array
        Vector containing the full width half maximum of the PSF in pixels, for
        each channel (cube_like); or single value (frame_like)
    cy, cx : None, float or 1D array, optional
        If None: will use the barycentre of the image found by
        photutils.centroid_com()
        If floats: coordinates of the center, assumed to be the same in all
        frames if the input is a cube.
        If 1D arrays: they must be the same length as the 0th dimension of the
        input cube.
    sig: Float scalar, optional
        Number of stddev above or below the median of the pixels in the same
        annulus, to consider a pixel as bad.
    protect_mask : int or float, optional
        If larger than 0, radius of a circular aperture (at the center of the
        frames) in which no bad pixels will be identified. This can be useful
        to protect the star and vicinity.
    r_in_std: float or None, optional
        Inner radius (in pixels) of the annulus used for the calculation of the
        standard deviation of the background noise - used as a min threshold
        when identifying bad pixels. Default: 50.
    r_out_std: float or None, optional
        Outer radius in pixels of the annulus used for the calculation of the
        standard deviation of the background noise - used as a min threshold
        when identifying bad pixels. If set to None, the default will be to
        consider the largest annulus starting at r_in_std which fits within the
        frame.
    verbose: bool, {False, True}, optional
        Whether to print out the number of bad pixels in each frame.
    half_res_y: bool, {True,False}, optional
        Whether the input data have only half the angular resolution vertically
        compared to horizontally (e.g. SINFONI data).
        The algorithm will then correct the bad pixels every other row.
    min_thr, max_thr: {None,float}, optional
        Any pixel whose value is lower (resp. larger) than this threshold will
        be automatically considered bad and hence sigma_filtered. If None, it
        is not used.
    full_output: bool, {False,True}, optional
        Whether to return as well the cube of bad pixel maps and the cube of
        defined annuli.

    Returns:
    --------
    obj_tmp: 2d or 3d array
        The bad pixel corrected frame/cube.
    If full_output is set to True, it returns as well:
    bpix_map: 2d or 3d array
        The bad pixel map or the cube of bpix maps
    ann_frame_cumul: 2 or 3d array
        The cube of defined annuli
    """

    obj_tmp = array.copy()
    ndims = obj_tmp.ndim
    assert ndims == 2 or ndims == 3, "Object is not two or three dimensional.\n"

    # thresholds
    if min_thr is None:
        min_thr = np.amin(obj_tmp)-1
    if max_thr is None:
        max_thr = np.amax(obj_tmp)-1

    def bp_removal_2d(obj_tmp, cy, cx, fwhm, sig, protect_mask, r_in_std,
                      r_out_std, verbose):

        n_x = obj_tmp.shape[1]
        n_y = obj_tmp.shape[0]

        # Squash the frame if twice less resolved vertically than horizontally
        if half_res_y:
            if n_y % 2 != 0:
                msg = 'The input frames do not have of an even number of rows. '
                msg2 = 'Hence, you should not use option half_res_y = True'
                raise ValueError(msg+msg2)
            n_y = int(n_y/2)
            cy = int(cy/2)
            frame = obj_tmp.copy()
            obj_tmp = np.zeros([n_y, n_x])
            for yy in range(n_y):
                obj_tmp[yy] = frame[2*yy]

        # 1/ Stddev of background
        if r_in_std or r_out_std:
            r_in_std = min(r_in_std*fwhm, cx-2, cy-2, n_x-cx-2, n_y-cy-2)
            if r_out_std:
                r_out_std *= fwhm
            else:
                r_out_std = min(n_y-(cy+r_in_std), cy-r_in_std,
                                n_x-(cx+r_in_std), cx-r_in_std)
            ceny, cenx = frame_center(obj_tmp)
            width = max(2, r_out_std-r_in_std)
            obj_tmp_crop = get_annulus_segments(obj_tmp, r_in_std, width,
                                                mode="val")
        else:
            obj_tmp_crop = obj_tmp
        _, _, stddev = sigma_clipped_stats(obj_tmp_crop, sigma=2.5)

        # 2/ Define each annulus, its median and stddev

        ymax = max(cy, n_y-cy)
        xmax = max(cx, n_x-cx)
        if half_res_y:
            ymax *= 2
        rmax = np.sqrt(ymax**2+xmax**2)
        # the annuli definition is optimized for Airy rings
        ann_width = max(1.5, 0.5*fwhm)  # 0.61*fwhm
        nrad = int(rmax/ann_width)+1
        d_bord_max = max(n_y-cy, cy, n_x-cx, cx)
        if half_res_y:
            d_bord_max = max(2*(n_y-cy), 2*cy, n_x-cx, cx)

        big_ell_frame = np.zeros_like(obj_tmp)
        sma_ell_frame = np.zeros_like(obj_tmp)
        ann_frame_cumul = np.zeros_like(obj_tmp)
        n_neig = np.zeros(nrad, dtype=np.int16)
        med_neig = np.zeros(nrad)
        std_neig = np.zeros(nrad)
        neighbours = np.zeros([nrad, n_y*n_x])

        for rr in range(nrad):
            if rr > int(d_bord_max/ann_width):
                # just to merge farthest annuli with very few elements
                rr_big = nrad
                rr_sma = int(d_bord_max/ann_width)
            else:
                rr_big = rr
                rr_sma = rr
            if half_res_y:
                big_ell_idx = ellipse(r=cy, c=cx,
                                      r_radius=((rr_big+1)*ann_width)/2,
                                      c_radius=(rr_big+1)*ann_width,
                                      shape=(n_y, n_x))
                if rr != 0:
                    small_ell_idx = ellipse(r=cy, c=cx,
                                            r_radius=(rr_sma*ann_width)/2,
                                            c_radius=rr_sma*ann_width,
                                            shape=(n_y, n_x))
            else:
                big_ell_idx = disk((cy, cx), radius=(rr_big+1)*ann_width,
                                   shape=(n_y, n_x))
                if rr != 0:
                    small_ell_idx = disk((cy, cx), radius=rr_sma*ann_width,
                                         shape=(n_y, n_x))
            big_ell_frame[big_ell_idx] = 1
            if rr != 0:
                sma_ell_frame[small_ell_idx] = 1
            ann_frame = big_ell_frame - sma_ell_frame
            n_neig[rr] = ann_frame[np.where(ann_frame)].shape[0]
            neighbours[rr, :n_neig[rr]] = obj_tmp[np.where(ann_frame)]
            ann_frame_cumul[np.where(ann_frame)] = rr

            # We delete iteratively max and min outliers in each annulus,
            # so that the annuli median and stddev are not corrupted by bpixs
            neigh = neighbours[rr, :n_neig[rr]]
            n_rm = 0
            n_pix_init = neigh.shape[0]
            while neigh.shape[0] >= np.amin(n_neig[rr]) and n_rm < n_pix_init/5:
                min_neigh = np.amin(neigh)
                if reject_outliers(neigh, min_neigh, m=5, stddev=stddev):
                    min_idx = np.argmin(neigh)
                    neigh = np.delete(neigh, min_idx)
                    n_rm += 1
                else:
                    max_neigh = np.amax(neigh)
                    if reject_outliers(neigh, max_neigh, m=5, stddev=stddev):
                        max_idx = np.argmax(neigh)
                        neigh = np.delete(neigh, max_idx)
                        n_rm += 1
                    else:
                        break
            n_neig[rr] = neigh.shape[0]
            neighbours[rr, :n_neig[rr]] = neigh
            neighbours[rr, n_neig[rr]:] = 0
            med_neig[rr] = np.median(neigh)
            std_neig[rr] = np.std(neigh)

        # 3/ Create a tuple-array with coordinates of a circle of radius 1.8*fwhm
        # centered on the provided coordinates of the star
        if protect_mask:
            if half_res_y:
                circl_new = ellipse(cy, cx, r_radius=protect_mask/2.,
                                    c_radius=protect_mask, shape=(n_y, n_x))
            else:
                circl_new = disk((cy, cx), radius=protect_mask,
                                 shape=(n_y, n_x))
        else:
            circl_new = []

        # 4/ Loop on all pixels to check bpix
        obj_tmp_corr = obj_tmp.copy()
        obj_tmp_corr, bpix_map = correct_ann_outliers(obj_tmp, ann_width, sig,
                                                      med_neig, std_neig, cy,
                                                      cx, min_thr, max_thr,
                                                      stddev, half_res_y)

        # 5/ Count bpix and uncorrect if within the circle
        nbpix_tot = np.sum(bpix_map)
        nbpix_tbc = nbpix_tot - np.sum(bpix_map[circl_new])
        bpix_map[circl_new] = 0
        obj_tmp_corr[circl_new] = obj_tmp[circl_new]
        if verbose:
            print(nbpix_tot, ' bpix in total, and ', nbpix_tbc, ' corrected.')

        # Unsquash all the frames
        if half_res_y:
            frame = obj_tmp_corr.copy()
            frame_bpix = bpix_map.copy()
            n_y = 2*n_y
            obj_tmp_corr = np.zeros([n_y, n_x])
            bpix_map = np.zeros([n_y, n_x])
            ann_frame = ann_frame_cumul.copy()
            ann_frame_cumul = np.zeros([n_y, n_x])
            for yy in range(n_y):
                obj_tmp_corr[yy] = frame[int(yy/2)]
                bpix_map[yy] = frame_bpix[int(yy/2)]
                ann_frame_cumul[yy] = ann_frame[int(yy/2)]

        return obj_tmp_corr, bpix_map, ann_frame_cumul

    if cy is None or cx is None:
        cy, cx = frame_center(obj_tmp)
    if ndims == 2:
        obj_tmp, bpix_map, ann_frame_cumul = bp_removal_2d(obj_tmp, cy, cx,
                                                           fwhm, sig,
                                                           protect_mask,
                                                           r_in_std, r_out_std,
                                                           verbose)
    if ndims == 3:
        n_z = obj_tmp.shape[0]
        bpix_map = np.zeros_like(obj_tmp)
        ann_frame_cumul = np.zeros_like(obj_tmp)
        if isinstance(fwhm, (int, float)):
            fwhm = [fwhm]*n_z
        if isinstance(cy, (float, int)) and isinstance(cx, (float, int)):
            cy = [cy]*n_z
            cx = [cx]*n_z
        for i in range(n_z):
            if verbose:
                print('************Frame # ', i, ' *************')
                print('centroid assumed at coords:', cx[i], cy[i])
            res_i = bp_removal_2d(obj_tmp[i], cy[i], cx[i], fwhm[i], sig,
                                  protect_mask, r_in_std, r_out_std, verbose)
            obj_tmp[i], bpix_map[i], ann_frame_cumul[i] = res_i

    if full_output:
        return obj_tmp, bpix_map, ann_frame_cumul
    else:
        return obj_tmp


def cube_fix_badpix_clump(array, bpm_mask=None, cy=None, cx=None, fwhm=4.,
                          sig=4., protect_mask=0, verbose=True,
                          half_res_y=False, min_thr=None, max_nit=15, mad=True,
                          full_output=False):
    """
    Function to identify and correct clumps of bad pixels. Very fast when a bad
    pixel map is provided. If a bad pixel map is not provided, the bad pixel
    clumps will be searched iteratively and replaced by the median of good
    neighbouring pixel values, when enough of them are available. The size of
    the box is set by the closest odd integer larger than fwhm (to avoid
    accidentally replacing point sources).



    Parameters
    ----------
    array : 3D or 2D array
        Input 3d cube or 2d image.
    bpix_map: 3D or 2D array, opt
        Input bad pixel array. Should have same dimenstions as array. If not
        provided, the algorithm will attempt to identify bad pixel clumps
        automatically.
    cy,cx : float or 1D array. opt
        Vector with approximate y and x coordinates of the star for each channel
        (cube_like), or single 2-elements vector (frame_like). These will be
        used if bpix_map is None and protect_mask set to True. If left to None,
        default values will correspond to the central pixel coordinates.
    fwhm: float or 1D array, opt
        Vector containing the full width half maximum of the PSF in pixels, for
        each channel (cube_like); or single value (frame_like). Shouod be
        provided if bpix map is None.
    sig: float, optional
        Value representing the number of "sigmas" above or below the "median" of
        the neighbouring pixel, to consider a pixel as bad. See details on
        parameter "m" of function reject_outlier.
    protect_mask : int or float, optional
        If larger than 0, radius of a circular aperture (at the center of the
        frames) in which no bad pixels will be identified. This can be useful
        to protect the star and vicinity.
    verbose: bool, {False,True}, optional
        Whether to print the number of bad pixels and number of iterations
        required for each frame.
    half_res_y: bool, {True,False}, optional
        Whether the input data has only half the angular resolution vertically
        compared to horizontally (e.g. the case of SINFONI data); in other words
        there are always 2 rows of pixels with exactly the same values.
        The algorithm will just consider every other row (hence making it
        twice faster), then apply the bad pixel correction on all rows.
    min_thr: float, tuple or None, opt
        If a float is provided, corresponds to a minimum absolute threshold
        below which pixels are not considered bad (can be used to avoid the
        identification of bad pixels within noise).
        If a tuple of 2 values, corresponds to the range of values within which
        not to consider a pixel as bad. (e.g. (-0.1, 10.)).
    max_nit: float, optional
        Maximum number of iterations on a frame to correct bpix. Typically, it
        should be set to less than ny/2 or nx/2. This is a mean of precaution in
        case the algorithm gets stuck with 2 neighbouring pixels considered bpix
        alternately on two consecutively iterations hence leading to an infinite
        loop (very very rare case).
    mad : {False, True}, bool optional
        If True, the median absolute deviation will be used instead of the
        standard deviation.
    full_output: bool, {False,True}, optional
        Whether to return as well the cube of bad pixel maps and the cube of
        defined annuli.

    Returns:
    --------
    obj_tmp: 2d or 3d array
        the bad pixel corrected frame/cube.
    If full_output is set to True, it returns as well:
    bpix_map: 2d or 3d array
        the bad pixel map or the cube of bpix maps
    """

    obj_tmp = array.copy()
    ndims = obj_tmp.ndim
    assert ndims == 2 or ndims == 3, "Object is not two or three dimensional.\n"

    if bpm_mask is not None:
        if bpm_mask.shape[-2:] != array.shape[-2:]:
            raise TypeError("Bad pixel map has wrong y/x dimensions.")

    def bp_removal_2d(obj_tmp, cy, cx, fwhm, sig, protect_mask, min_thr,
                      half_res_y, mad, verbose):
        n_x = obj_tmp.shape[1]
        n_y = obj_tmp.shape[0]

        if half_res_y:
            if n_y % 2 != 0:
                msg = 'The input frames do not have of an even number of rows. '
                msg2 = 'Hence, you should not use option half_res_y = True'
                raise ValueError(msg+msg2)
            n_y = int(n_y/2)
            frame = obj_tmp.copy()
            obj_tmp = np.zeros([n_y, n_x])
            for yy in range(n_y):
                obj_tmp[yy] = frame[2*yy]

        fwhm_round = int(round(fwhm))
        # This should reduce the chance to accidently correct a bright planet:
        if fwhm_round % 2 == 0:
            neighbor_box = max(3, fwhm_round+1)
        else:
            neighbor_box = max(3, fwhm_round)
        nneig = sum(np.arange(3, neighbor_box+2, 2))

        # 1/ Create a tuple-array with coordinates of a circle of radius 1.8*fwhm
        # centered on the approximate coordinates of the star
        if protect_mask:
            if half_res_y:
                circl_new = ellipse(int(cy/2), cx, r_radius=0.5*protect_mask,
                                    c_radius=protect_mask, shape=(n_y, n_x))
            else:
                circl_new = disk((cy, cx), radius=protect_mask,
                                  shape=(n_y, n_x))
        else:
            circl_new = []

        # 3/ Create a bad pixel map, by detecting them with clip_array
        bp = clip_array(obj_tmp, sig, sig, out_good=False, neighbor=True,
                        num_neighbor=neighbor_box, mad=mad,
                        half_res_y=half_res_y)
        bpix_map = np.zeros_like(obj_tmp)
        bpix_map[bp] = 1
        if min_thr is not None:
            if np.isscalar(min_thr):
                min_thr = (-min_thr, min_thr)
            elif not isinstance(min_thr, tuple):
                msg = "if provided, min_thr should be float or tuple"
                raise ValueError(msg)
            else:
                if len(min_thr) != 2:
                    msg = "if min_thr is a tuple, it should have 2 elements"
                    raise ValueError(msg)
            cond1 = obj_tmp > min_thr[0]
            cond2 = obj_tmp < min_thr[1]
            bpix_map[np.where(cond1 & cond2)] = 0
        nbpix_tot = np.sum(bpix_map)
        bpix_map[circl_new] = 0
        nbpix_tbc = np.sum(bpix_map)
        bpix_map_cumul = np.zeros_like(bpix_map)
        bpix_map_cumul[:] = bpix_map[:]
        nit = 0

        # 4/ Loop over the bpix correction with sigma_filter, until 0 bpix left
        while nbpix_tbc > 0 and nit < max_nit:
            nit = nit+1
            if verbose:
                print("Iteration {}: {} bpix in total, {} to be "
                      "corrected".format(nit, nbpix_tot, nbpix_tbc))
            obj_tmp = sigma_filter(obj_tmp, bpix_map, neighbor_box=neighbor_box,
                                   min_neighbors=nneig, half_res_y=half_res_y,
                                   verbose=verbose)
            bp = clip_array(obj_tmp, sig, sig, out_good=False, neighbor=True,
                            num_neighbor=neighbor_box, mad=mad,
                            half_res_y=half_res_y)
            bpix_map = np.zeros_like(obj_tmp)
            bpix_map[bp] = 1
            if min_thr is not None:
                cond1 = obj_tmp > min_thr[0]
                cond2 = obj_tmp < min_thr[1]
                bpix_map[np.where(cond1 & cond2)] = 0
            nbpix_tot = np.sum(bpix_map)
            bpix_map[circl_new] = 0
            nbpix_tbc = np.sum(bpix_map)
            bpix_map_cumul = bpix_map_cumul+bpix_map

        if verbose:
            print('All bad pixels are corrected.')

        if half_res_y:
            frame = obj_tmp.copy()
            frame_bpix = bpix_map_cumul.copy()
            n_y = 2*n_y
            obj_tmp = np.zeros([n_y, n_x])
            bpix_map_cumul = np.zeros([n_y, n_x])
            for yy in range(n_y):
                obj_tmp[yy] = frame[int(yy/2)]
                bpix_map_cumul[yy] = frame_bpix[int(yy/2)]

        return obj_tmp, bpix_map_cumul

    if ndims == 2:
        if bpm_mask is None:
            if (cy is None or cx is None) and protect_mask:
                cy, cx = frame_center(array)
            obj_tmp, bpix_map_cumul = bp_removal_2d(obj_tmp, cy, cx, fwhm, sig,
                                                    protect_mask, min_thr,
                                                    half_res_y, mad, verbose)
        else:
            fwhm_round = int(round(fwhm))
            fwhm_round = fwhm_round+1-(fwhm_round % 2)  # make it odd
            neighbor_box = max(3, fwhm_round)  # to not replace a companion
            nneig = sum(np.arange(3, neighbor_box+2, 2))
            obj_tmp = sigma_filter(obj_tmp, bpm_mask, neighbor_box, nneig,
                                   half_res_y, verbose)
            bpix_map_cumul = bpm_mask

    if ndims == 3:
        n_z = obj_tmp.shape[0]
        if bpm_mask is None:
            if cy is None or cx is None:
                cy, cx = frame_center(array)
                cy = [cy]*n_z
                cx = [cx]*n_z
            elif isinstance(cy, (float, int)) and isinstance(cx, (float, int)):
                cy = [cy]*n_z
                cx = [cx]*n_z
            if isinstance(fwhm, (float, int)):
                fwhm = [fwhm]*n_z
            bpix_map_cumul = np.zeros_like(obj_tmp)
            for i in range(n_z):
                if verbose:
                    print('************Frame # ', i, ' *************')
                obj_tmp[i], bpix_map_cumul[i] = bp_removal_2d(obj_tmp[i], cy[i],
                                                              cx[i], fwhm[i],
                                                              sig, protect_mask,
                                                              min_thr,
                                                              half_res_y, mad,
                                                              verbose)
        else:
            if isinstance(fwhm, (float, int)):
                fwhm_round = int(round(fwhm))
            else:
                fwhm_round = int(np.median(fwhm))
            fwhm_round = fwhm_round+1-(fwhm_round % 2)  # make it odd
            neighbor_box = max(3, fwhm_round)  # to not replace a companion
            nneig = sum(np.arange(3, neighbor_box+2, 2))
            for i in range(n_z):
                if verbose:
                    print('************Frame # ', i, ' *************')
                if bpm_mask.ndim == 3:
                    bpm = bpm_mask[i]
                else:
                    bpm = bpm_mask
                obj_tmp[i] = sigma_filter(obj_tmp[i], bpm, neighbor_box,
                                          nneig, half_res_y, verbose)
            bpix_map_cumul = bpm_mask

    if full_output:
        return obj_tmp, bpix_map_cumul
    else:
        return obj_tmp


def cube_fix_badpix_ifs(array, lbdas, fluxes=None, mask=None, cy=None, cx=None,
                        clumps=True, sigma_clip=3, num_neig=5, size=5,
                        protect_mask=0, mad=False, fwhm=4, min_thr=None,
                        max_nit=15, imlib='vip-fft', interpolation='lanczos4',
                        ignore_nan=True, verbose=True, full_output=False):
    """
    Function to identify and correct bad pixels in an IFS cube, leveraging on
    the radial expansion of the PSF with wavelength.
    Bad pixel identification is done with either the `cube_fix_badpix_isolated`
    or the `cube_fix_badpix_clump` function in PSF subtracted frames (through<commit_msg>Updated output bad pixel map when correcting for clumps to have only 0 or 1<commit_after>#! /usr/bin/env python

"""
Module with functions for correcting bad pixels in cubes.
"""


__author__ = '<NAME>, <NAME>'
__all__ = ['frame_fix_badpix_isolated',
           'cube_fix_badpix_isolated',
           'cube_fix_badpix_annuli',
           'cube_fix_badpix_clump',
           'cube_fix_badpix_ifs',
           'cube_fix_badpix_with_kernel']

import numpy as np
from skimage.draw import disk, ellipse
from scipy.ndimage import median_filter
from astropy.stats import sigma_clipped_stats
from ..stats import sigma_filter
from ..var import frame_center, get_annulus_segments, frame_filter_lowpass
from ..stats import clip_array
from ..config import timing, time_ini, Progressbar
from .rescaling import find_scal_vector, frame_rescaling

import warnings
try:
    from numba import njit
    no_numba = False
except ImportError:
    msg = "Numba python bindings are missing."
    warnings.warn(msg, ImportWarning)
    no_numba = True


def frame_fix_badpix_isolated(array, bpm_mask=None, sigma_clip=3, num_neig=5,
                              size=5, protect_mask=0, cxy=None, mad=False,
                              ignore_nan=True, verbose=True, full_output=False):
    """ Corrects the bad pixels, marked in the bad pixel mask. The bad pixel is
    replaced by the median of the adjacent pixels. This function is very fast
    but works only with isolated (sparse) pixels.

    Parameters
    ----------
    array : numpy ndarray
        Input 2d array.
    bpm_mask : numpy ndarray, optional
        Input bad pixel map. Zeros frame where the bad pixels have a value of
        1. If None is provided a bad pixel map will be created using
        sigma clip statistics.
    sigma_clip : int, optional
        In case no bad pixel mask is provided all the pixels above and below
        sigma_clip*STDDEV will be marked as bad.
    num_neig : int, optional
        The side of the square window around each pixel where the sigma clipped
        statistics are calculated (STDDEV and MEDIAN). If the value is equal to
        0 then the statistics are computed in the whole frame.
    size : odd int, optional
        The size the box (size x size) of adjacent pixels for the median
        filter.
    protect_mask : int or float, optional
        If larger than 0, radius of a circular aperture (at the center of the
        frames) in which no bad pixels will be identified. This can be useful
        to protect the star and vicinity.
    cxy: None or tuple
        If protect_mask is True, this is the location of the star centroid in
        the images. If None, assumes the star is already centered. If a tuple,
        the location of the star is assumed to be the same in all frames of the
        cube.
    mad : {False, True}, bool optional
        If True, the median absolute deviation will be used instead of the
        standard deviation.
    ignore_nan: bool, optional
        Whether to not consider NaN values as bad pixels. If False, will also
        correct them.
    verbose : bool, optional
        If True additional information will be printed.
    full_output: bool, {False,True}, optional
        Whether to return as well the cube of bad pixel maps and the cube of
        defined annuli.

    Return
    ------
    frame : numpy ndarray
        Frame with bad pixels corrected.
    bpm_mask: 2d array
        The bad pixel map
    """
    if array.ndim != 2:
        raise TypeError('Array is not a 2d array or single frame')
    if size % 2 == 0:
        raise TypeError('Size of the median blur kernel must be an odd integer')

    if bpm_mask is not None:
        bpm_mask = bpm_mask.astype('bool')

    if verbose:
        start = time_ini()

    if num_neig > 0:
        neigh = True
    else:
        neigh = False

    frame = array.copy()
    if cxy is None:
        cy, cx = frame_center(frame)
    else:
        cx, cy = cxy

    if bpm_mask is None:
        ori_nan_mask = np.where(np.isnan(frame))
        ind = clip_array(frame, sigma_clip, sigma_clip, neighbor=neigh,
                         num_neighbor=num_neig, mad=mad)
        bpm_mask = np.zeros_like(frame)
        bpm_mask[ind] = 1
        if ignore_nan:
            bpm_mask[ori_nan_mask] = 0
        if protect_mask:
            cir = disk((cy, cx), protect_mask, shape=bpm_mask.shape)
            bpm_mask[cir] = 0
        bpm_mask = bpm_mask.astype('bool')

    smoothed = median_filter(frame, size, mode='mirror')
    frame[np.where(bpm_mask)] = smoothed[np.where(bpm_mask)]
    array_out = frame
    count_bp = np.sum(bpm_mask)

    if verbose:
        msg = "/nDone replacing {} bad pixels using the median of neighbors"
        print(msg.format(count_bp))
        timing(start)

    if full_output:
        return array_out, bpm_mask
    else:
        return array_out


def cube_fix_badpix_isolated(array, bpm_mask=None, sigma_clip=3, num_neig=5,
                             size=5, frame_by_frame=False, protect_mask=0,
                             cxy=None, mad=False, ignore_nan=True, verbose=True,
                             full_output=False):
    """ Corrects the bad pixels, marked in the bad pixel mask. The bad pixel is
    replaced by the median of the adjacent pixels. This function is very fast
    but works only with isolated (sparse) pixels.

    Parameters
    ----------
    array : numpy ndarray
        Input 3d array.
    bpm_mask : numpy ndarray, optional
        Input bad pixel map. Zeros frame where the bad pixels have a value of 1.
        If None is provided a bad pixel map will be created per frame using
        sigma clip statistics.
    sigma_clip : int, optional
        In case no bad pixel mask is provided all the pixels above and below
        sigma_clip*STDDEV will be marked as bad.
    num_neig : int, optional
        The side of the square window around each pixel where the sigma clipped
        statistics are calculated (STDDEV and MEDIAN). If the value is equal to
        0 then the statistics are computed in the whole frame.
    size : odd int, optional
        The size the box (size x size) of adjacent pixels for the median filter.
    frame_by_frame: bool, optional
        Whether to correct bad pixels frame by frame in the cube. By default it
        is set to False; the bad pixels are computed on the mean frame of the
        stack (faster but not necessarily optimal).
    protect_mask : int or float, optional
        If larger than 0, radius of a circular aperture (at the center of the
        frames) in which no bad pixels will be identified. This can be useful
        to protect the star and vicinity.
    cxy: None, tuple or 2d numpy ndarray
        If protect_mask is True, this is the location of the star centroid in
        the images. If None, assumes the star is already centered. If a tuple,
        the location of the star is assumed to be the same in all frames of the
        cube. If a (n_frames x 2) ndarray, it should contain the xy location of
        the star in each frame.
    mad : {False, True}, bool optional
        If True, the median absolute deviation will be used instead of the
        standard deviation.
    ignore_nan: bool, optional
        Whether to not consider NaN values as bad pixels. If False, will also
        correct them.
    verbose : bool, optional
        If True additional information will be printed.
    full_output: bool, {False,True}, optional
        Whether to return as well the cube of bad pixel maps and the cube of
        defined annuli.

    Return
    ------
    array_out : numpy ndarray
        Cube with bad pixels corrected.
    bpm_mask: 2d or 3d array [if full_output is True]
        The bad pixel map or the cube of bpix maps
    """
    if array.ndim != 3:
        raise TypeError('Array is not a 3d array or cube')
    if size % 2 == 0:
        raise TypeError('Size of the median blur kernel must be an odd integer')

    if bpm_mask is not None:
        bpm_mask = bpm_mask.astype('bool')

    if verbose:
        start = time_ini()

    if num_neig > 0:
        neigh = True
    else:
        neigh = False

    nz = array.shape[0]

    if cxy is None:
        cy, cx = frame_center(array[0])
    elif isinstance(cxy, tuple):
        cx, cy = cxy
    elif isinstance(cxy, np.ndarray):
        if cxy.shape[0] != nz or cxy.shape[1] != 2 or cxy.ndim != 2:
            raise ValueError("cxy does not have right shape")
        elif not frame_by_frame:
            msg = "cxy must be a tuple or None if not in frame_by_frame mode"
            raise ValueError(msg)
        else:
            cx = cxy[:, 0]
            cy = cxy[:, 1]

    array_out = array.copy()
    final_bpm = np.zeros_like(array_out, dtype=bool)
    n_frames = array.shape[0]
    count_bp = 0
    if frame_by_frame:
        if np.isscalar(cx):
            cx = [cx]*nz
            cy = [cy]*nz
        if bpm_mask is not None:
            if bpm_mask.ndim == 2:
                bpm_mask = [bpm_mask]*n_frames
                bpm_mask = np.array(bpm_mask)
        for i in Progressbar(range(n_frames), desc="processing frames"):
            if bpm_mask is not None:
                bpm_mask_tmp = bpm_mask[i]
            else:
                bpm_mask_tmp = None
            res = frame_fix_badpix_isolated(array[i], bpm_mask=bpm_mask_tmp,
                                            sigma_clip=sigma_clip,
                                            num_neig=num_neig, size=size,
                                            protect_mask=protect_mask,
                                            verbose=False, cxy=(cx[i], cy[i]),
                                            ignore_nan=ignore_nan,
                                            full_output=True)
            array_out[i] = res[0]
            final_bpm[i] = res[1]
        count_bp = np.sum(final_bpm)
    else:
        if bpm_mask is None:
            ori_nan_mask = np.where(np.isnan(np.nanmean(array, axis=0)))
            ind = clip_array(np.nanmean(array, axis=0), sigma_clip, sigma_clip,
                             neighbor=neigh, num_neighbor=num_neig, mad=mad)
            final_bpm = np.zeros_like(array[0], dtype=bool)
            final_bpm[ind] = 1
            if ignore_nan:
                final_bpm[ori_nan_mask] = 0
            if protect_mask:
                cir = disk((cy, cx), protect_mask, shape=final_bpm.shape)
                final_bpm[cir] = 0
            final_bpm = final_bpm.astype('bool')
        else:
            if bpm_mask.ndim == 3:
                final_bpm = np.median(bpm_mask, axis=0)
            else:
                final_bpm = bpm_mask.copy()

        for i in Progressbar(range(n_frames), desc="processing frames"):
            frame = array_out[i]
            smoothed = median_filter(frame, size, mode='mirror')
            frame[np.where(final_bpm)] = smoothed[np.where(final_bpm)]
            if verbose:
                count_bp += np.sum(final_bpm)

    if verbose:
        msg = "Done replacing {:.0f} bad pixels using the median of neighbors"
        print(msg.format(count_bp))
        if not frame_by_frame:
            msg = "(i.e. {:.0f} static bad pixels per channel))"
            print(msg.format(count_bp/n_frames))
        timing(start)

    if full_output:
        return array_out, final_bpm
    else:
        return array_out


def cube_fix_badpix_annuli(array, fwhm, cy=None, cx=None, sig=5.,
                           protect_mask=0, r_in_std=50, r_out_std=None,
                           verbose=True, half_res_y=False, min_thr=None,
                           max_thr=None, full_output=False):
    """
    Function to correct the bad pixels annulus per annulus (centered on the
    provided location of the star), in an input frame or cube.
    This function is faster than bp_clump_removal; hence to be prefered in all
    cases where there is only one bright source with circularly symmetric PSF.
    The bad pixel values are replaced by: ann_median + ann_stddev*random_gauss;
    where ann_median is the median of the annulus, ann_stddev is the standard
    deviation in the annulus, and random_gauss is a random factor picked from a
    gaussian distribution centered on 0 and with variance 1.

    Parameters
    ----------
    array : 3D or 2D array
        Input 3d cube or 2d image.
    fwhm: float or 1D array
        Vector containing the full width half maximum of the PSF in pixels, for
        each channel (cube_like); or single value (frame_like)
    cy, cx : None, float or 1D array, optional
        If None: will use the barycentre of the image found by
        photutils.centroid_com()
        If floats: coordinates of the center, assumed to be the same in all
        frames if the input is a cube.
        If 1D arrays: they must be the same length as the 0th dimension of the
        input cube.
    sig: Float scalar, optional
        Number of stddev above or below the median of the pixels in the same
        annulus, to consider a pixel as bad.
    protect_mask : int or float, optional
        If larger than 0, radius of a circular aperture (at the center of the
        frames) in which no bad pixels will be identified. This can be useful
        to protect the star and vicinity.
    r_in_std: float or None, optional
        Inner radius (in pixels) of the annulus used for the calculation of the
        standard deviation of the background noise - used as a min threshold
        when identifying bad pixels. Default: 50.
    r_out_std: float or None, optional
        Outer radius in pixels of the annulus used for the calculation of the
        standard deviation of the background noise - used as a min threshold
        when identifying bad pixels. If set to None, the default will be to
        consider the largest annulus starting at r_in_std which fits within the
        frame.
    verbose: bool, {False, True}, optional
        Whether to print out the number of bad pixels in each frame.
    half_res_y: bool, {True,False}, optional
        Whether the input data have only half the angular resolution vertically
        compared to horizontally (e.g. SINFONI data).
        The algorithm will then correct the bad pixels every other row.
    min_thr, max_thr: {None,float}, optional
        Any pixel whose value is lower (resp. larger) than this threshold will
        be automatically considered bad and hence sigma_filtered. If None, it
        is not used.
    full_output: bool, {False,True}, optional
        Whether to return as well the cube of bad pixel maps and the cube of
        defined annuli.

    Returns:
    --------
    obj_tmp: 2d or 3d array
        The bad pixel corrected frame/cube.
    If full_output is set to True, it returns as well:
    bpix_map: 2d or 3d array
        The bad pixel map or the cube of bpix maps
    ann_frame_cumul: 2 or 3d array
        The cube of defined annuli
    """

    obj_tmp = array.copy()
    ndims = obj_tmp.ndim
    assert ndims == 2 or ndims == 3, "Object is not two or three dimensional.\n"

    # thresholds
    if min_thr is None:
        min_thr = np.amin(obj_tmp)-1
    if max_thr is None:
        max_thr = np.amax(obj_tmp)-1

    def bp_removal_2d(obj_tmp, cy, cx, fwhm, sig, protect_mask, r_in_std,
                      r_out_std, verbose):

        n_x = obj_tmp.shape[1]
        n_y = obj_tmp.shape[0]

        # Squash the frame if twice less resolved vertically than horizontally
        if half_res_y:
            if n_y % 2 != 0:
                msg = 'The input frames do not have of an even number of rows. '
                msg2 = 'Hence, you should not use option half_res_y = True'
                raise ValueError(msg+msg2)
            n_y = int(n_y/2)
            cy = int(cy/2)
            frame = obj_tmp.copy()
            obj_tmp = np.zeros([n_y, n_x])
            for yy in range(n_y):
                obj_tmp[yy] = frame[2*yy]

        # 1/ Stddev of background
        if r_in_std or r_out_std:
            r_in_std = min(r_in_std*fwhm, cx-2, cy-2, n_x-cx-2, n_y-cy-2)
            if r_out_std:
                r_out_std *= fwhm
            else:
                r_out_std = min(n_y-(cy+r_in_std), cy-r_in_std,
                                n_x-(cx+r_in_std), cx-r_in_std)
            ceny, cenx = frame_center(obj_tmp)
            width = max(2, r_out_std-r_in_std)
            obj_tmp_crop = get_annulus_segments(obj_tmp, r_in_std, width,
                                                mode="val")
        else:
            obj_tmp_crop = obj_tmp
        _, _, stddev = sigma_clipped_stats(obj_tmp_crop, sigma=2.5)

        # 2/ Define each annulus, its median and stddev

        ymax = max(cy, n_y-cy)
        xmax = max(cx, n_x-cx)
        if half_res_y:
            ymax *= 2
        rmax = np.sqrt(ymax**2+xmax**2)
        # the annuli definition is optimized for Airy rings
        ann_width = max(1.5, 0.5*fwhm)  # 0.61*fwhm
        nrad = int(rmax/ann_width)+1
        d_bord_max = max(n_y-cy, cy, n_x-cx, cx)
        if half_res_y:
            d_bord_max = max(2*(n_y-cy), 2*cy, n_x-cx, cx)

        big_ell_frame = np.zeros_like(obj_tmp)
        sma_ell_frame = np.zeros_like(obj_tmp)
        ann_frame_cumul = np.zeros_like(obj_tmp)
        n_neig = np.zeros(nrad, dtype=np.int16)
        med_neig = np.zeros(nrad)
        std_neig = np.zeros(nrad)
        neighbours = np.zeros([nrad, n_y*n_x])

        for rr in range(nrad):
            if rr > int(d_bord_max/ann_width):
                # just to merge farthest annuli with very few elements
                rr_big = nrad
                rr_sma = int(d_bord_max/ann_width)
            else:
                rr_big = rr
                rr_sma = rr
            if half_res_y:
                big_ell_idx = ellipse(r=cy, c=cx,
                                      r_radius=((rr_big+1)*ann_width)/2,
                                      c_radius=(rr_big+1)*ann_width,
                                      shape=(n_y, n_x))
                if rr != 0:
                    small_ell_idx = ellipse(r=cy, c=cx,
                                            r_radius=(rr_sma*ann_width)/2,
                                            c_radius=rr_sma*ann_width,
                                            shape=(n_y, n_x))
            else:
                big_ell_idx = disk((cy, cx), radius=(rr_big+1)*ann_width,
                                   shape=(n_y, n_x))
                if rr != 0:
                    small_ell_idx = disk((cy, cx), radius=rr_sma*ann_width,
                                         shape=(n_y, n_x))
            big_ell_frame[big_ell_idx] = 1
            if rr != 0:
                sma_ell_frame[small_ell_idx] = 1
            ann_frame = big_ell_frame - sma_ell_frame
            n_neig[rr] = ann_frame[np.where(ann_frame)].shape[0]
            neighbours[rr, :n_neig[rr]] = obj_tmp[np.where(ann_frame)]
            ann_frame_cumul[np.where(ann_frame)] = rr

            # We delete iteratively max and min outliers in each annulus,
            # so that the annuli median and stddev are not corrupted by bpixs
            neigh = neighbours[rr, :n_neig[rr]]
            n_rm = 0
            n_pix_init = neigh.shape[0]
            while neigh.shape[0] >= np.amin(n_neig[rr]) and n_rm < n_pix_init/5:
                min_neigh = np.amin(neigh)
                if reject_outliers(neigh, min_neigh, m=5, stddev=stddev):
                    min_idx = np.argmin(neigh)
                    neigh = np.delete(neigh, min_idx)
                    n_rm += 1
                else:
                    max_neigh = np.amax(neigh)
                    if reject_outliers(neigh, max_neigh, m=5, stddev=stddev):
                        max_idx = np.argmax(neigh)
                        neigh = np.delete(neigh, max_idx)
                        n_rm += 1
                    else:
                        break
            n_neig[rr] = neigh.shape[0]
            neighbours[rr, :n_neig[rr]] = neigh
            neighbours[rr, n_neig[rr]:] = 0
            med_neig[rr] = np.median(neigh)
            std_neig[rr] = np.std(neigh)

        # 3/ Create a tuple-array with coordinates of a circle of radius 1.8*fwhm
        # centered on the provided coordinates of the star
        if protect_mask:
            if half_res_y:
                circl_new = ellipse(cy, cx, r_radius=protect_mask/2.,
                                    c_radius=protect_mask, shape=(n_y, n_x))
            else:
                circl_new = disk((cy, cx), radius=protect_mask,
                                 shape=(n_y, n_x))
        else:
            circl_new = []

        # 4/ Loop on all pixels to check bpix
        obj_tmp_corr = obj_tmp.copy()
        obj_tmp_corr, bpix_map = correct_ann_outliers(obj_tmp, ann_width, sig,
                                                      med_neig, std_neig, cy,
                                                      cx, min_thr, max_thr,
                                                      stddev, half_res_y)

        # 5/ Count bpix and uncorrect if within the circle
        nbpix_tot = np.sum(bpix_map)
        nbpix_tbc = nbpix_tot - np.sum(bpix_map[circl_new])
        bpix_map[circl_new] = 0
        obj_tmp_corr[circl_new] = obj_tmp[circl_new]
        if verbose:
            print(nbpix_tot, ' bpix in total, and ', nbpix_tbc, ' corrected.')

        # Unsquash all the frames
        if half_res_y:
            frame = obj_tmp_corr.copy()
            frame_bpix = bpix_map.copy()
            n_y = 2*n_y
            obj_tmp_corr = np.zeros([n_y, n_x])
            bpix_map = np.zeros([n_y, n_x])
            ann_frame = ann_frame_cumul.copy()
            ann_frame_cumul = np.zeros([n_y, n_x])
            for yy in range(n_y):
                obj_tmp_corr[yy] = frame[int(yy/2)]
                bpix_map[yy] = frame_bpix[int(yy/2)]
                ann_frame_cumul[yy] = ann_frame[int(yy/2)]

        return obj_tmp_corr, bpix_map, ann_frame_cumul

    if cy is None or cx is None:
        cy, cx = frame_center(obj_tmp)
    if ndims == 2:
        obj_tmp, bpix_map, ann_frame_cumul = bp_removal_2d(obj_tmp, cy, cx,
                                                           fwhm, sig,
                                                           protect_mask,
                                                           r_in_std, r_out_std,
                                                           verbose)
    if ndims == 3:
        n_z = obj_tmp.shape[0]
        bpix_map = np.zeros_like(obj_tmp)
        ann_frame_cumul = np.zeros_like(obj_tmp)
        if isinstance(fwhm, (int, float)):
            fwhm = [fwhm]*n_z
        if isinstance(cy, (float, int)) and isinstance(cx, (float, int)):
            cy = [cy]*n_z
            cx = [cx]*n_z
        for i in range(n_z):
            if verbose:
                print('************Frame # ', i, ' *************')
                print('centroid assumed at coords:', cx[i], cy[i])
            res_i = bp_removal_2d(obj_tmp[i], cy[i], cx[i], fwhm[i], sig,
                                  protect_mask, r_in_std, r_out_std, verbose)
            obj_tmp[i], bpix_map[i], ann_frame_cumul[i] = res_i

    if full_output:
        return obj_tmp, bpix_map, ann_frame_cumul
    else:
        return obj_tmp


def cube_fix_badpix_clump(array, bpm_mask=None, cy=None, cx=None, fwhm=4.,
                          sig=4., protect_mask=0, verbose=True,
                          half_res_y=False, min_thr=None, max_nit=15, mad=True,
                          full_output=False):
    """
    Function to identify and correct clumps of bad pixels. Very fast when a bad
    pixel map is provided. If a bad pixel map is not provided, the bad pixel
    clumps will be searched iteratively and replaced by the median of good
    neighbouring pixel values, when enough of them are available. The size of
    the box is set by the closest odd integer larger than fwhm (to avoid
    accidentally replacing point sources).



    Parameters
    ----------
    array : 3D or 2D array
        Input 3d cube or 2d image.
    bpix_map: 3D or 2D array, opt
        Input bad pixel array. Should have same dimenstions as array. If not
        provided, the algorithm will attempt to identify bad pixel clumps
        automatically.
    cy,cx : float or 1D array. opt
        Vector with approximate y and x coordinates of the star for each channel
        (cube_like), or single 2-elements vector (frame_like). These will be
        used if bpix_map is None and protect_mask set to True. If left to None,
        default values will correspond to the central pixel coordinates.
    fwhm: float or 1D array, opt
        Vector containing the full width half maximum of the PSF in pixels, for
        each channel (cube_like); or single value (frame_like). Shouod be
        provided if bpix map is None.
    sig: float, optional
        Value representing the number of "sigmas" above or below the "median" of
        the neighbouring pixel, to consider a pixel as bad. See details on
        parameter "m" of function reject_outlier.
    protect_mask : int or float, optional
        If larger than 0, radius of a circular aperture (at the center of the
        frames) in which no bad pixels will be identified. This can be useful
        to protect the star and vicinity.
    verbose: bool, {False,True}, optional
        Whether to print the number of bad pixels and number of iterations
        required for each frame.
    half_res_y: bool, {True,False}, optional
        Whether the input data has only half the angular resolution vertically
        compared to horizontally (e.g. the case of SINFONI data); in other words
        there are always 2 rows of pixels with exactly the same values.
        The algorithm will just consider every other row (hence making it
        twice faster), then apply the bad pixel correction on all rows.
    min_thr: float, tuple or None, opt
        If a float is provided, corresponds to a minimum absolute threshold
        below which pixels are not considered bad (can be used to avoid the
        identification of bad pixels within noise).
        If a tuple of 2 values, corresponds to the range of values within which
        not to consider a pixel as bad. (e.g. (-0.1, 10.)).
    max_nit: float, optional
        Maximum number of iterations on a frame to correct bpix. Typically, it
        should be set to less than ny/2 or nx/2. This is a mean of precaution in
        case the algorithm gets stuck with 2 neighbouring pixels considered bpix
        alternately on two consecutively iterations hence leading to an infinite
        loop (very very rare case).
    mad : {False, True}, bool optional
        If True, the median absolute deviation will be used instead of the
        standard deviation.
    full_output: bool, {False,True}, optional
        Whether to return as well the cube of bad pixel maps and the cube of
        defined annuli.

    Returns:
    --------
    obj_tmp: 2d or 3d array
        the bad pixel corrected frame/cube.
    If full_output is set to True, it returns as well:
    bpix_map: 2d or 3d array
        the bad pixel map or the cube of bpix maps
    """

    obj_tmp = array.copy()
    ndims = obj_tmp.ndim
    assert ndims == 2 or ndims == 3, "Object is not two or three dimensional.\n"

    if bpm_mask is not None:
        if bpm_mask.shape[-2:] != array.shape[-2:]:
            raise TypeError("Bad pixel map has wrong y/x dimensions.")

    def bp_removal_2d(obj_tmp, cy, cx, fwhm, sig, protect_mask, min_thr,
                      half_res_y, mad, verbose):
        n_x = obj_tmp.shape[1]
        n_y = obj_tmp.shape[0]

        if half_res_y:
            if n_y % 2 != 0:
                msg = 'The input frames do not have of an even number of rows. '
                msg2 = 'Hence, you should not use option half_res_y = True'
                raise ValueError(msg+msg2)
            n_y = int(n_y/2)
            frame = obj_tmp.copy()
            obj_tmp = np.zeros([n_y, n_x])
            for yy in range(n_y):
                obj_tmp[yy] = frame[2*yy]

        fwhm_round = int(round(fwhm))
        # This should reduce the chance to accidently correct a bright planet:
        if fwhm_round % 2 == 0:
            neighbor_box = max(3, fwhm_round+1)
        else:
            neighbor_box = max(3, fwhm_round)
        nneig = sum(np.arange(3, neighbor_box+2, 2))

        # 1/ Create a tuple-array with coordinates of a circle of radius 1.8*fwhm
        # centered on the approximate coordinates of the star
        if protect_mask:
            if half_res_y:
                circl_new = ellipse(int(cy/2), cx, r_radius=0.5*protect_mask,
                                    c_radius=protect_mask, shape=(n_y, n_x))
            else:
                circl_new = disk((cy, cx), radius=protect_mask,
                                  shape=(n_y, n_x))
        else:
            circl_new = []

        # 3/ Create a bad pixel map, by detecting them with clip_array
        bp = clip_array(obj_tmp, sig, sig, out_good=False, neighbor=True,
                        num_neighbor=neighbor_box, mad=mad,
                        half_res_y=half_res_y)
        bpix_map = np.zeros_like(obj_tmp)
        bpix_map[bp] = 1
        if min_thr is not None:
            if np.isscalar(min_thr):
                min_thr = (-min_thr, min_thr)
            elif not isinstance(min_thr, tuple):
                msg = "if provided, min_thr should be float or tuple"
                raise ValueError(msg)
            else:
                if len(min_thr) != 2:
                    msg = "if min_thr is a tuple, it should have 2 elements"
                    raise ValueError(msg)
            cond1 = obj_tmp > min_thr[0]
            cond2 = obj_tmp < min_thr[1]
            bpix_map[np.where(cond1 & cond2)] = 0
        nbpix_tot = np.sum(bpix_map)
        bpix_map[circl_new] = 0
        nbpix_tbc = np.sum(bpix_map)
        bpix_map_cumul = np.zeros_like(bpix_map)
        bpix_map_cumul[:] = bpix_map[:]
        nit = 0

        # 4/ Loop over the bpix correction with sigma_filter, until 0 bpix left
        while nbpix_tbc > 0 and nit < max_nit:
            nit = nit+1
            if verbose:
                print("Iteration {}: {} bpix in total, {} to be "
                      "corrected".format(nit, nbpix_tot, nbpix_tbc))
            obj_tmp = sigma_filter(obj_tmp, bpix_map, neighbor_box=neighbor_box,
                                   min_neighbors=nneig, half_res_y=half_res_y,
                                   verbose=verbose)
            bp = clip_array(obj_tmp, sig, sig, out_good=False, neighbor=True,
                            num_neighbor=neighbor_box, mad=mad,
                            half_res_y=half_res_y)
            bpix_map = np.zeros_like(obj_tmp)
            bpix_map[bp] = 1
            if min_thr is not None:
                cond1 = obj_tmp > min_thr[0]
                cond2 = obj_tmp < min_thr[1]
                bpix_map[np.where(cond1 & cond2)] = 0
            nbpix_tot = np.sum(bpix_map)
            bpix_map[circl_new] = 0
            nbpix_tbc = np.sum(bpix_map)
            bpix_map_cumul = bpix_map_cumul+bpix_map

        if verbose:
            print('All bad pixels are corrected.')

        if half_res_y:
            frame = obj_tmp.copy()
            frame_bpix = bpix_map_cumul.copy()
            n_y = 2*n_y
            obj_tmp = np.zeros([n_y, n_x])
            bpix_map_cumul = np.zeros([n_y, n_x])
            for yy in range(n_y):
                obj_tmp[yy] = frame[int(yy/2)]
                bpix_map_cumul[yy] = frame_bpix[int(yy/2)]

        return obj_tmp, bpix_map_cumul

    if ndims == 2:
        if bpm_mask is None:
            if (cy is None or cx is None) and protect_mask:
                cy, cx = frame_center(array)
            obj_tmp, bpix_map_cumul = bp_removal_2d(obj_tmp, cy, cx, fwhm, sig,
                                                    protect_mask, min_thr,
                                                    half_res_y, mad, verbose)
        else:
            fwhm_round = int(round(fwhm))
            fwhm_round = fwhm_round+1-(fwhm_round % 2)  # make it odd
            neighbor_box = max(3, fwhm_round)  # to not replace a companion
            nneig = sum(np.arange(3, neighbor_box+2, 2))
            obj_tmp = sigma_filter(obj_tmp, bpm_mask, neighbor_box, nneig,
                                   half_res_y, verbose)
            bpix_map_cumul = bpm_mask

    if ndims == 3:
        n_z = obj_tmp.shape[0]
        if bpm_mask is None:
            if cy is None or cx is None:
                cy, cx = frame_center(array)
                cy = [cy]*n_z
                cx = [cx]*n_z
            elif isinstance(cy, (float, int)) and isinstance(cx, (float, int)):
                cy = [cy]*n_z
                cx = [cx]*n_z
            if isinstance(fwhm, (float, int)):
                fwhm = [fwhm]*n_z
            bpix_map_cumul = np.zeros_like(obj_tmp)
            for i in range(n_z):
                if verbose:
                    print('************Frame # ', i, ' *************')
                obj_tmp[i], bpix_map_cumul[i] = bp_removal_2d(obj_tmp[i], cy[i],
                                                              cx[i], fwhm[i],
                                                              sig, protect_mask,
                                                              min_thr,
                                                              half_res_y, mad,
                                                              verbose)
        else:
            if isinstance(fwhm, (float, int)):
                fwhm_round = int(round(fwhm))
            else:
                fwhm_round = int(np.median(fwhm))
            fwhm_round = fwhm_round+1-(fwhm_round % 2)  # make it odd
            neighbor_box = max(3, fwhm_round)  # to not replace a companion
            nneig = sum(np.arange(3, neighbor_box+2, 2))
            for i in range(n_z):
                if verbose:
                    print('************Frame # ', i, ' *************')
                if bpm_mask.ndim == 3:
                    bpm = bpm_mask[i]
                else:
                    bpm = bpm_mask
                obj_tmp[i] = sigma_filter(obj_tmp[i], bpm, neighbor_box,
                                          nneig, half_res_y, verbose)
            bpix_map_cumul = bpm_mask

    # make it a binary map
    bpix_map_cumul[np.where(bpix_map_cumul>1)] = 1

    if full_output:
        return obj_tmp, bpix_map_cumul
    else:
        return obj_tmp


def cube_fix_badpix_ifs(array, lbdas, fluxes=None, mask=None, cy=None, cx=None,
                        clumps=True, sigma_clip=3, num_neig=5, size=5,
                        protect_mask=0, mad=False, fwhm=4, min_thr=None,
                        max_nit=15, imlib='vip-fft', interpolation='lanczos4',
                        ignore_nan=True, verbose=True, full_output=False):
    """
    Function to identify and correct bad pixels in an IFS cube, leveraging on
    the radial expansion of the PSF with wavelength.
    Bad pixel identification is done with either the `cube_fix_badpix_isolated`
    or the `cube_fix_badpix_clump` function in PSF subtracted frames (through
{'fraction_non_alphanumeric': 0.07670757547746697, 'fraction_numerical': 0.00759396464682376, 'mean_word_length': 1.7661080953504837, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 128, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '225233', 'n_tokens_mistral': 21795, 'n_tokens_neox': 20763, 'n_words': 6908}