text: #! /usr/bin/env python
  
  """
  Module with S/N calculation functions.
  We strongly recommend users to read Mawet et al. (2014) before using routines 
  of this module: https://ui.adsabs.harvard.edu/abs/2014ApJ...792...97M/abstract
  """
  
  __author__ = 'Carlos Alberto Gomez Gonzalez, O. Absil @ ULg, V. Christiaens'
  __all__ = ['snr',
             'snrmap',
             'significance',
             'frame_report']
  
  import numpy as np
  import photutils
  from scipy.stats import norm, t
  from hciplot import plot_frames
  from skimage.draw import disk, circle_perimeter
  from matplotlib import pyplot as plt
  from astropy.convolution import convolve, Tophat2DKernel
  from astropy.stats import median_absolute_deviation as mad
  from multiprocessing import cpu_count
  from ..conf.utils_conf import pool_map, iterable, sep
  from ..conf import time_ini, timing, check_array
  from ..var import get_annulus_segments, frame_center, dist
  
  
  def snrmap(array, fwhm, approximated=False, plot=False, known_sources=None,
             nproc=None, array2=None, use2alone=False, 
             exclude_negative_lobes=False, verbose=True, **kwargs):
      """Parallel implementation of the S/N map generation function. Applies the
      S/N function (small samples penalty) at each pixel.
      
      The S/N is computed as in Mawet et al. (2014) for each radial separation.    
      https://ui.adsabs.harvard.edu/abs/2014ApJ...792...97M/abstract
      
      *** DISCLAIMER ***
      Signal-to-noise ratio is not significance! For a conversion from snr to 
      n-sigma (i.e. the equivalent confidence level of a Gaussian n-sigma), use 
      the significance() function.    
      
      
      Parameters
      ----------
      array : numpy ndarray
          Input frame (2d array).
      fwhm : float
          Size in pixels of the FWHM.
      approximated : bool, optional
          If True, an approximated S/N map is generated.
      plot : bool, optional
          If True plots the S/N map. False by default.
      known_sources : None, tuple or tuple of tuples, optional
          To take into account existing sources. It should be a tuple of float/int
          or a tuple of tuples (of float/int) with the coordinate(s) of the known
          sources.
      nproc : int or None
          Number of processes for parallel computing.
      array2 : numpy ndarray, optional
          Additional image (e.g. processed image with negative derotation angles) 
          enabling to have more noise samples. Should have the 
          same dimensions as array.
      use2alone: bool, optional
          Whether to use array2 alone to estimate the noise (might be useful to 
          estimate the snr of extended disk features).
      verbose: bool, optional
          Whether to print timing or not.
      **kwargs : dictionary, optional
          Arguments to be passed to ``plot_frames`` to customize the plot (and to
          save it to disk).
  
      Returns
      -------
      snrmap : 2d numpy ndarray
          Frame with the same size as the input frame with each pixel.
      """
      if verbose:
          start_time = time_ini()
  
      check_array(array, dim=2, msg='array')
      sizey, sizex = array.shape
      snrmap_array = np.zeros_like(array)
      width = min(sizey, sizex) / 2 - 1.5 * fwhm
      mask = get_annulus_segments(array, (fwhm / 2) + 2, width, mode="mask")[0]
      mask = np.ma.make_mask(mask)
      # by making a bool mask *after* applying the mask to the array, we also mask
      # out zero values from the array. This logic cannot be simplified by using
      # mode="ind"!
      yy, xx = np.where(mask)
      coords = zip(xx, yy)
  
      if nproc is None:
          nproc = cpu_count() // 2        # Hyper-threading doubles the # of cores
  
      if known_sources is None:
  
          # proxy to S/N calculation
          if approximated:
              cy, cx = frame_center(array)
              tophat_kernel = Tophat2DKernel(fwhm / 2)
              array = convolve(array, tophat_kernel)
              width = min(sizey, sizex) / 2 - 1.5 * fwhm
              mask = get_annulus_segments(array, (fwhm / 2) + 1, width - 1,
                                          mode="mask")[0]
              mask = np.ma.make_mask(mask)
              yy, xx = np.where(mask)
              coords = [(int(x), int(y)) for (x, y) in zip(xx, yy)]
              res = pool_map(nproc, _snr_approx, array, iterable(coords), fwhm,
                             cy, cx)
              res = np.array(res)
              yy = res[:, 0]
              xx = res[:, 1]
              snr_value = res[:, 2]
              snrmap_array[yy.astype(int), xx.astype(int)] = snr_value
  
          # computing s/n map with Mawet+14 definition
          else:
              res = pool_map(nproc, snr, array, iterable(coords), fwhm, True,
                             array2, use2alone, exclude_negative_lobes)
              res = np.array(res)
              yy = res[:, 0]
              xx = res[:, 1]
              snr_value = res[:, -1]
              snrmap_array[yy.astype('int'), xx.astype('int')] = snr_value
  
      # masking known sources
      else:
          if not isinstance(known_sources, tuple):
              raise TypeError("`known_sources` must be a tuple or tuple of "
                              "tuples")
          else:
              source_mask = np.zeros_like(array)
              if isinstance(known_sources[0], tuple):
                  for coor in known_sources:
                      source_mask[coor[::-1]] = 1
              elif isinstance(known_sources[0], int):
                  source_mask[known_sources[1], known_sources[0]] = 1
              else:
                  raise TypeError("`known_sources` seems to have wrong type. It "
                                  "must be a tuple of ints or tuple of tuples "
                                  "(of ints)")
  
          # checking the mask with the sources
          if source_mask[source_mask == 1].shape[0] > 50:
              msg = 'Input source mask is too crowded (check its validity)'
              raise RuntimeError(msg)
  
          soury, sourx = np.where(source_mask == 1)
          sources = []
          coor_ann = []
          arr_masked_sources = array.copy()
          centery, centerx = frame_center(array)
          for y, x in zip(soury, sourx):
              radd = dist(centery, centerx, int(y), int(x))
              if int(radd) < centery - np.ceil(fwhm):
                  sources.append((y, x))
  
          for source in sources:
              y, x = source
              radd = dist(centery, centerx, int(y), int(x))
              anny, annx = get_annulus_segments(array, int(radd-fwhm),
                                                int(np.round(3 * fwhm)))[0]
  
              ciry, cirx = disk((y, x), int(np.ceil(fwhm)))
              # masking the sources positions (using the MAD of pixels in annulus)
              arr_masked_sources[ciry, cirx] = mad(array[anny, annx])
  
              # S/Ns of annulus without the sources
              coor_ann = [(x, y) for (x, y) in zip(annx, anny) if (x, y) not in
                          zip(cirx, ciry)]
              res = pool_map(nproc, snr, arr_masked_sources, iterable(coor_ann),
                             fwhm, True, array2, use2alone, 
                             exclude_negative_lobes)
              res = np.array(res)
              yy_res = res[:, 0]
              xx_res = res[:, 1]
              snr_value = res[:, 4]
              snrmap_array[yy_res.astype('int'), xx_res.astype('int')] = snr_value
              coor_ann += coor_ann
  
          # S/Ns of the rest of the frame without the annulus
          coor_rest = [(x, y) for (x, y) in zip(xx, yy) if (x, y) not in coor_ann]
          res = pool_map(nproc, snr, array, iterable(coor_rest), fwhm, True,
                         array2, use2alone, exclude_negative_lobes)
          res = np.array(res)
          yy_res = res[:, 0]
          xx_res = res[:, 1]
          snr_value = res[:, 4]
          snrmap_array[yy_res.astype('int'), xx_res.astype('int')] = snr_value
  
      if plot:
          plot_frames(snrmap_array, colorbar=True, title='S/N map', **kwargs)
  
      if verbose:
          print("S/N map created using {} processes".format(nproc))
          timing(start_time)
      return snrmap_array
  
  
  def _snr_approx(array, source_xy, fwhm, centery, centerx):
      """
      array - frame convolved with top hat kernel
      """
      sourcex, sourcey = source_xy
      rad = dist(centery, centerx, sourcey, sourcex)
      ind_aper = disk((sourcey, sourcex), fwhm/2.)
      # noise : STDDEV in convolved array of 1px wide annulus (while
      # masking the flux aperture) * correction of # of resolution elements
      ind_ann = circle_perimeter(int(centery), int(centerx), int(rad))
      array2 = array.copy()
      array2[ind_aper] = mad(array[ind_ann])  # mask
      n2 = (2 * np.pi * rad) / fwhm - 1
      noise = array2[ind_ann].std(ddof=1) * np.sqrt(1+(1/n2))
      # signal : central px minus the mean of the pxs (masked) in 1px annulus
      signal = array[sourcey, sourcex] - array2[ind_ann].mean()
      snr_value = signal / noise
      return sourcey, sourcex, snr_value
  
  
  def snr(array, source_xy, fwhm, full_output=False, array2=None, use2alone=False,
          exclude_negative_lobes=False, plot=False, verbose=False):
      """
      Calculate the S/N (signal to noise ratio) of a test resolution element
      in a residual frame (e.g. post-processed with LOCI, PCA, etc). Implements
      the approach described in Mawet et al. 2014 on small sample statistics,
      where a student t-test (eq. 9) can be used to determine S/N (and contrast)
      in high contrast imaging. 3 extra possibilities compared to Mawet et al. 
      2014 (https://ui.adsabs.harvard.edu/abs/2014ApJ...792...97M/abstract):
  
          * possibility to provide a second array (e.g. obtained with opposite
          derotation angles) to have more apertures for noise estimation
          * possibility to exclude negative ADI lobes directly adjacent to the
          tested xy location, to not bias the noise estimate
          * possibility to use only the second array for the noise estimation
          (useful for images containing a lot of disk/extended signals).
          
      *** DISCLAIMER ***
      Signal-to-noise ratio is not significance! For a conversion from snr to 
      n-sigma (i.e. the equivalent confidence level of a Gaussian n-sigma), use 
      the significance() function.    
      
      Parameters
      ----------
      array : numpy ndarray, 2d
          Post-processed frame where we want to measure S/N.
      source_xy : tuple of floats
          X and Y coordinates of the planet or test speckle.
      fwhm : float
          Size in pixels of the FWHM.
      full_output : bool, optional
          If True returns back the S/N value, the y, x input coordinates, noise
          and flux.
      array2 : None or numpy ndarray, 2d, optional
          Additional image (e.g. processed image with negative derotation angles) 
          enabling to have more apertures for noise estimation at each radial
          separation. Should have the same dimensions as array.
      use2alone : bool, opt
          Whether to use array2 alone to estimate the noise (can be useful to 
          estimate the S/N of extended disk features)
      exclude_negative_lobes : bool, opt
          Whether to include the adjacent aperture lobes to the tested location 
          or not. Can be set to True if the image shows significant neg lobes.
      plot : bool, optional
          Plots the frame and the apertures considered for clarity. 
      verbose: bool, optional
          Chooses whether to print some output or not.
  
      Returns
      -------
      [if full_output=True:]
      sourcey : numpy ndarray
          [full_output=True] Input coordinates (``source_xy``) in Y.
      sourcex : numpy ndarray
          [full_output=True] Input coordinates (``source_xy``) in X.
      f_source : float
          [full_output=True] Flux in test elemnt.
      fluxes : numpy ndarray
          [full_output=True] Background apertures fluxes.
      [always:]    
      snr_vale : float
          Value of the S/N for the given test resolution element.
      """
      check_array(array, dim=2, msg='array')
      if not isinstance(source_xy, tuple):
          raise TypeError("`source_xy` must be a tuple of floats")
      if array2 is not None:
          if not array2.shape == array.shape:
              raise TypeError('`array2` has not the same shape as input array')
  
      sourcex, sourcey = source_xy
      centery, centerx = frame_center(array)
      sep = dist(centery, centerx, float(sourcey), float(sourcex))
  
      if not sep > (fwhm/2)+1:
          raise RuntimeError('`source_xy` is too close to the frame center')
  
      sens = 'clock'  # counterclock
  
      angle = np.arcsin(fwhm/2./sep)*2
      number_apertures = int(np.floor(2*np.pi/angle))
      yy = np.zeros((number_apertures))
      xx = np.zeros((number_apertures))
      cosangle = np.cos(angle)
      sinangle = np.sin(angle)
      xx[0] = sourcex - centerx
      yy[0] = sourcey - centery
      for i in range(number_apertures-1):
          if sens == 'clock':
              xx[i+1] = cosangle*xx[i] + sinangle*yy[i] 
              yy[i+1] = cosangle*yy[i] - sinangle*xx[i] 
          elif sens == 'counterclock':
              xx[i+1] = cosangle*xx[i] - sinangle*yy[i] 
              yy[i+1] = cosangle*yy[i] + sinangle*xx[i]
  
      xx += centerx
      yy += centery
      rad = fwhm/2.
      if exclude_negative_lobes:
          xx = np.concatenate(([xx[0]], xx[2:-1]))
          yy = np.concatenate(([yy[0]], yy[2:-1]))
  
      apertures = photutils.CircularAperture((xx, yy), r=rad)  # Coordinates (X,Y)
      fluxes = photutils.aperture_photometry(array, apertures, method='exact')
      fluxes = np.array(fluxes['aperture_sum'])
  
      if array2 is not None:
          fluxes2 = photutils.aperture_photometry(array2, apertures,
                                                  method='exact')
          fluxes2 = np.array(fluxes2['aperture_sum'])
          if use2alone:
              fluxes = np.concatenate(([fluxes[0]], fluxes2[:]))
          else:
              fluxes = np.concatenate((fluxes, fluxes2))
  
      f_source = fluxes[0].copy()
      fluxes = fluxes[1:]
      n2 = fluxes.shape[0]
      backgr_apertures_std = fluxes.std(ddof=1)
      snr_vale = (f_source - fluxes.mean())/(backgr_apertures_std *
                                             np.sqrt(1+(1/n2)))
  
      if verbose:
          msg1 = 'S/N for the given pixel = {:.3f}'
          msg2 = 'Integrated flux in FWHM test aperture = {:.3f}'
          msg3 = 'Mean of background apertures integrated fluxes = {:.3f}'
          msg4 = 'Std-dev of background apertures integrated fluxes = {:.3f}'
          print(msg1.format(snr_vale))
          print(msg2.format(f_source))
          print(msg3.format(fluxes.mean()))
          print(msg4.format(backgr_apertures_std))
  
      if plot:
          _, ax = plt.subplots(figsize=(6, 6))
          ax.imshow(array, origin='lower', interpolation='nearest', alpha=0.5, 
                    cmap='gray')
          for i in range(xx.shape[0]):
              # Circle takes coordinates as (X,Y)
              aper = plt.Circle((xx[i], yy[i]), radius=fwhm/2., color='r',
                                fill=False, alpha=0.8)
              ax.add_patch(aper)
              cent = plt.Circle((xx[i], yy[i]), radius=0.8, color='r', fill=True,
                                alpha=0.5)
              ax.add_patch(cent)
              aper_source = plt.Circle((sourcex, sourcey), radius=0.7,
                                       color='b', fill=True, alpha=0.5)
              ax.add_patch(aper_source)
          ax.grid(False)
          plt.show()
  
      if full_output:
          return sourcey, sourcex, f_source, fluxes, snr_vale
      else:
          return snr_vale
  
  
  
  def significance(snr, rad, fwhm, student_to_gauss=True):
      """ Converts a S/N ratio (measured as in Mawet et al. 2014) into the 
      equivalent gaussian significance, i.e. the n-sigma with the same confidence 
      level as the S/N at the given separation.
       
       
      Parameters
      ----------
      snr : float or numpy array
          SNR value(s)
      rad : float or numpy array
          Radial separation(s) from the star in pixels. If an array, it should be
          the same shape as snr and provide the radial separation corresponding
          to each snr measurement.
      fwhm : float
          Full Width Half Maximum of the PSF.
      student_to_gauss : bool, optional
          Whether the conversion is from Student SNR to Gaussian significance. If 
          False, will assume the opposite: Gaussian significance to Student SNR.
      
      Returns
      -------
      sigma : float
          Gaussian significance in terms of n-sigma
      
      """
      
      if student_to_gauss:
          sigma = norm.ppf(t.cdf(snr,(rad/fwhm)*2*np.pi-2))
          if t.cdf(snr,(rad/fwhm)*2*np.pi-2)==1.0:
              print("Warning high S/N! cdf>0.9999999999999999 is rounded to 1")
              print(r"Returning 8.2 instead of inf. Quote detection>8.2 $\sigma$")
              return 8.2
      else:
          sigma = t.ppf(norm.cdf(snr), (rad/fwhm)*2*np.pi-2)
          
      return sigma      
              
  
  def frame_report(array, fwhm, source_xy=None, verbose=True):
      """ Gets information from given frame: Integrated flux in aperture, S/N of
      central pixel (either ``source_xy`` or max value), mean S/N in aperture.
  
      Parameters
      ----------
      array : numpy ndarray
          2d array or input frame.
      fwhm : float
          Size of the FWHM in pixels.
      source_xy : tuple of floats or list (of tuples of floats)
          X and Y coordinates of the center(s) of the source(s).
      verbose : bool, optional
          If True prints to stdout the frame info.
  
      Returns
      -------
      source_xy : tuple of floats or list (of tuples of floats)
          X and Y coordinates of the center(s) of the source(s).
      obj_flux : list of floats
          Integrated flux in aperture.
      snr_centpx : list of floats
          S/N of the ``source_xy`` pixels.
      meansnr_pixels : list of floats
          Mean S/N of pixels in 1xFWHM apertures centered on ``source_xy``.
      """
      if array.ndim != 2:
          raise TypeError('Array is not 2d.')
  
      obj_flux = []
      meansnr_pixels = []
      snr_centpx = []
  
      if source_xy is not None:
          if isinstance(source_xy, (list, tuple)):
              if not isinstance(source_xy[0], tuple):
                  source_xy = [source_xy]
          else:
              raise TypeError("`source_xy` must be a tuple of floats or tuple "
                              "of tuples")
  
          for xy in source_xy:
              x, y = xy
              if verbose:
                  print(sep)
                  print('Coords of chosen px (X,Y) = {:.1f}, {:.1f}'.format(x, y))
  
              # we get integrated flux on aperture with diameter=1FWHM
              aper = photutils.CircularAperture((x, y), r=fwhm / 2.)
              obj_flux_i = photutils.aperture_photometry(array, aper,
                                                         method='exact')
              obj_flux_i = obj_flux_i['aperture_sum'][0]
  
              # we get the mean and stddev of SNRs on aperture
              yy, xx = disk((y, x), fwhm / 2)
              snr_pixels_i = [snr(array, (x_, y_), fwhm, plot=False,
                                  verbose=False) for y_, x_ in zip(yy, xx)]
              meansnr_i = np.mean(snr_pixels_i)
              stdsnr_i = np.std(snr_pixels_i,ddof=1)
              pxsnr_i = snr(array, (x, y), fwhm, plot=False, verbose=False)
  
              obj_flux.append(obj_flux_i)
              meansnr_pixels.append(meansnr_i)
              snr_centpx.append(pxsnr_i)
  
              if verbose:
                  msg0 = 'Flux in a centered 1xFWHM circular aperture = {:.3f}'
                  print(msg0.format(obj_flux_i))
                  print('Central pixel S/N = {:.3f}'.format(pxsnr_i))
                  print(sep)
                  print('Inside a centered 1xFWHM circular aperture:')
                  msg1 = 'Mean S/N (shifting the aperture center) = {:.3f}'
                  print(msg1.format(meansnr_i))
                  msg2 = 'Max S/N (shifting the aperture center) = {:.3f}'
                  print(msg2.format(np.max(snr_pixels_i)))
                  msg3 = 'stddev S/N (shifting the aperture center) = {:.3f}'
                  print(msg3.format(stdsnr_i))
                  print('')
  
      else:
          y, x = np.where(array == array.max())
          y, x = y[0], x[0]  # assuming there is only one max, taking 1st if clump
          source_xy = (x, y)
          if verbose:
              print(sep)
              print('Coords of Max px (X,Y) = {:.1f}, {:.1f}'.format(x, y))
  
          # we get integrated flux on aperture with diameter=1FWHM
          aper = photutils.CircularAperture((x, y), r=fwhm / 2.)
          obj_flux_i = photutils.aperture_photometry(array, aper, method='exact')
          obj_flux_i = obj_flux_i['aperture_sum'][0]
  
          # we get the mean and stddev of SNRs on aperture
          yy, xx = disk((y, x), fwhm / 2.)
          snr_pixels_i = [snr(array, (x_, y_), fwhm, plot=False,
                              verbose=False) for y_, x_ in zip(yy, xx)]
          meansnr_pixels = np.mean(snr_pixels_i)
          stdsnr_i = np.std(snr_pixels_i,ddof=1)
          pxsnr_i = snr(array, (x, y), fwhm, plot=False, verbose=False)
  
          obj_flux.append(obj_flux_i)
          snr_centpx.append(pxsnr_i)
  
          if verbose:
              msg0 = 'Flux in a centered 1xFWHM circular aperture = {:.3f}'
              print(msg0.format(obj_flux_i))
              print('Central pixel S/N = {:.3f}'.format(pxsnr_i))
              print(sep)
              print('Inside a centered 1xFWHM circular aperture:')
              msg1 = 'Mean S/N (shifting the aperture center) = {:.3f}'
              print(msg1.format(meansnr_pixels))
              msg2 = 'Max S/N (shifting the aperture center) = {:.3f}'
              print(msg2.format(np.max(snr_pixels_i)))
              msg3 = 'stddev S/N (shifting the aperture center) = {:.3f}'
              print(msg3.format(stdsnr_i))
              print(sep)
  
      return source_xy, obj_flux, snr_centpx, meansnr_pixels
  