text: ''' Code to modify the existing Bad Pixel Masks by adding new masked regions.
  Written for Python 2.7 (as this is the supported eups version so far). If want
  to migrate to Python 3.+ test the multiprocessing call (the one with partial)
  '''
  
  import os
  import sys
  import time
  import logging
  import argparse
  import copy
  import numpy as np
  from functools import partial
  import multiprocessing as mp
  try:
      import matplotlib.pyplot as plt
  except:
      pass
  import fitsio
  
  # Setup logging
  logging.basicConfig(
      format='%(asctime)s - %(levelname)s - %(message)s',
      level=logging.INFO,
  )
  
  # BPM definitions
  # From imsupport::mask_bits.h
  bpmdef = dict()
  bpmdef.update({'BPMDEF_FLAT_MIN' : 1})
  #Pixels that are hot in the flats.
  bpmdef.update({'BPMDEF_FLAT_MAX' : 2})
  #Pixels that are in the BPM for the flats.
  bpmdef.update({'BPMDEF_FLAT_MASK' : 4})
  # Pixels that are hot in the biases.
  bpmdef.update({'BPMDEF_BIAS_HOT' : 8})
  # Pixels that are warm in the biases.
  bpmdef.update({'BPMDEF_BIAS_WARM' : 16})
  # Pixels that are in the BPM for the biases.
  bpmdef.update({'BPMDEF_BIAS_MASK' : 32})
  # Pixels that are downstream of a hot pixel in the bias.
  bpmdef.update({'BPMDEF_BIAS_COL' : 64})
  # Pixels on the glowing edges of the CCD.
  bpmdef.update({'BPMDEF_EDGE' : 128})
  # Correctable pixels (usually downstream of hot pixels).
  bpmdef.update({'BPMDEF_CORR' : 256})
  # Imperfect calibration or excess noise, such as tape bumps.
  # Ignore for highest-precision work.
  bpmdef.update({'BPMDEF_SUSPECT' : 512})
  # Columns with charge redistribution in sky exposures.
  bpmdef.update({'BPMDEF_FUNKY_COL' : 1024})
  # Outliers in stacked sky exposures.
  bpmdef.update({'BPMDEF_WACKY_PIX' : 2048})
  # Pixel on non-functional amplifier.
  bpmdef.update({'BPMDEF_BADAMP' : 4096})
  # suspect due to edge proximity.
  bpmdef.update({'BPMDEF_NEAREDGE' : 8192})
  # suspect due to known tape bump.
  bpmdef.update({'BPMDEF_TAPEBUMP' : 16384})
  # End of BPMDEF
  #
  # Change data type to unsigned integers, to match the dtype of the BPMs
  #
  for k in bpmdef:
      bpmdef[k] = np.uint(bpmdef[k])
  
  
  def open_fits(fnm, ):
      ''' Open the FITS, read the data and store it on a list
      '''
      tmp = fitsio.FITS(fnm)
      data, header = [], []
      for extension in tmp:
          ext_tmp = np.copy(extension.read())
          data.append(ext_tmp)
          header_tmp = copy.deepcopy(extension.read_header())
          header.append(header_tmp)
      tmp.close()
      if (len(data) == 0):
          logging.error('No extensions were found on {0}'.format(fnm))
          exit(1)
      elif (len(data) >= 1):
          return data, header
  
  def open_txt(fnm, Nskip):
      ''' Open the text file containing the 2 coordinates per box (bottom-left
      and top-right). Skip the header lines, if any
      The coordinates must be integer, and all rows must have the same number
      of columns
      '''
      # If I don't specify the names key on dtype, then the output will be an
      # array with one more dimension, to harbor the columns. Below array has
      # shape (N,) but if no names are inputs, the shape would be (N, 4)
      m_kw = {
          'names' : ('x1', 'y1', 'x2', 'y2', 'flag'),
          'formats' : ('i4', 'i4', 'i4', 'i4', '|S50'),
      }
      m = np.loadtxt(fnm, dtype=m_kw, comments='#',
                     converters=None, skiprows=Nskip, usecols=None,
                     unpack=False, ndmin=0)
      return m
  
  def get_dict_key(kdict, xval):
      ''' Returns the key associated to the value, for the input dictionary.
      It the key is not in the dictionary, returns None
      '''
      res = [nm for nm, val in kdict.iteritems() if val == xval]
      if (res == []):
          return None
      else:
          return res
  
  def flatten_list(list_2levels):
      ''' Function to flatten a list of lists, generating an output with
      all elements ia a single level. No duplicate drop neither sort are
      performed
      '''
      f = lambda x: [item for sublist in x for item in sublist]
      res = f(list_2levels)
      return res
  
  def bit_count(int_type):
      ''' Function to count the amount of bits composing a number, that is the
      number of base 2 components on which it can be separated, and then
      reconstructed by simply sum them. Idea from Wiki Python.
      Brian Kernighan's way for counting bits. Thismethod was in fact
      discovered by Wegner and then Lehmer in the 60s
      This method counts bit-wise. Each iteration is not simply a step of 1.
      Example: iter1: (2066, 2065), iter2: (2064, 2063), iter3: (2048, 2047)
      Inputs
      - int_type: integer
      Output
      - integer with the number of base-2 numbers needed for the decomposition
      '''
      counter = 0
      while int_type:
          int_type &= int_type - 1
          counter += 1
      return counter
  
  def bit_decompose(int_x):
      ''' Function to decompose a number in base-2 numbers. This is performed by
      two binary operators. Idea from Stackoverflow.
          x << y
      Returns x with the bits shifted to the left by y places (and new bits on
      the right-hand-side are zeros). This is the same as multiplying x by 2**y.
          x & y
      Does a "bitwise and". Each bit of the output is 1 if the corresponding bit
      of x AND of y is 1, otherwise it's 0.
      Inputs
      - int_x: integer
      Returns
      - list of base-2 values from which adding them, the input integer can be
      recovered
      '''
      base2 = []
      i = 1
      while (i <= int_x):
          if (i & int_x):
              base2.append(i)
          i <<= 1
      return base2
  
  def bit_superpose(box_i,
                    bpm=None,
                    aux_bit1=None,
                    aux_bit2_decomp=None,
                    non_masked_val=None,
                    mask_unmasked=None):
      ''' Function to stack bits. Intended to be parallelized inside joint_mask()
      This function works in one box at a time, one flag at a time. The only
      variable that iterates is the box_i_msk
      Returns
      - mask_box: a mask of the same dimensions as the input BPM array. The only
      unmasked values are the one inside the current box
      '''
      logging.info('Working in box: {0}'.format(box_i))
      # Mask all numerical values, then unmask only the pixels on which to work
      # This works: mask_box = np.ma.masked_where(bpm is not np.nan, bpm)
      mask_box = np.ma.array(np.copy(bpm), mask=True)
      if (mask_box.all() is not np.ma.masked):
          logging.warning('The masking of all pixels failed')
      # Unmask the box region to operate on it.
      # NOTE: np.ma.nomask assign zeros instead of recover original data
      # NOTE: np.ma.getdata returns ndarray
      mask_box[box_i['y1']:box_i['y2'] ,
               box_i['x1']:box_i['x2']] = mask_box[box_i['y1']:box_i['y2'] ,
                                                   box_i['x1']:box_i['x2']].data
      # Inside the unmasked section, look for each of the bits and its
      # correspondence in the pixels.
      # Add non-masked values to the equation
      if (box_i.size == 1): 
          box_i_flag = str(box_i['flag'])
      else:
          box_i_flag = box_i['flag']
      for idx, unib in enumerate(list(aux_bit1) + [non_masked_val]):
          # Copy the mask of the box, and operate on the unique bits
          mask_bitx = np.ma.masked_where(mask_box != unib, mask_box, copy=True)
          # Control flow, if there are not non-masked values
          if (np.ma.count(mask_bitx) == 0):
              continue
          if (unib != 0):
              # Get the composing bits for the value, and check if the number
              # we want to assign is already on the components of the pixel
              # bad pixel mask
              unique_bit_comp = aux_bit2_decomp[idx]
              # Control flow, if value is already in the bit decomposition
              # Check for non-masked values also
              if (bpmdef[box_i_flag] in unique_bit_comp):
                  continue
              #
              # Here an error was raised and solved with str()
              # monitor this section for further
              #
              # Add the flag to the entire unmasked region
              mask_bitx += bpmdef[box_i_flag]
              # The values from mask_bitx must go to mask_box. Then the
              # values from mask_box must go to the BPM array
              mask_box[~mask_bitx.mask] = mask_bitx[~mask_bitx.mask]
          elif (unib == non_masked_val) and (mask_unmasked):
              # Here we take care of the non-masked pixels
              mask_bitx += bpmdef[box_i_flag]
              # The values from mask_bitx must go to mask_box. Then the
              # values from mask_box must go to the BPM array
              mask_box[~mask_bitx.mask] = mask_bitx[~mask_bitx.mask]
      return mask_box
  
  def joint_mask(bpm,
                 header_orig,
                 box_msk,
                 fname_bpm_orig,
                 non_masked_val=0,
                 mask_unmasked=True,
                 Nproc=None,
                 out_dtype='uint16',
                 file_check=False,
                 out_fname=None,
                 checksum=False,):
      ''' Applies the box-defined mask to the original BPM, avoiding replace
      bits values. The flag need to have an entry in BPMDEF.
      The bit comparison is carried out in parallel, one process per box.
      To be used: ['BPMDEF_SUSPECT', 'BPMDEF_NEAREDGE']
      Inputs
      - bpm: BPM mask array
      - box_msk: numpy structured array corner-coordinates (left-bottom,
      right-top) for the different boxes
      - non_masked_val: value for the non-masked pixels
      - mask_unmasked: whether to mask unmasked pixels. If False then only add
      the flag value to the already masked pixels
      - Nproc: number of processes to be run in parallel for the bit comparison,
      one per box
      - out_dtype: dtype to be used in the BPMs
      - file_check: wether to check for existing file
      '''
      # Check the bits
      #
      # Unique bits in the mask not considering the zeros
      # Transform to tuple for protection
      aux_bit1 = np.unique(bpm[np.where(bpm != non_masked_val)].ravel())
      aux_bit1 = tuple(aux_bit1)
      if (not len(aux_bit1)):
          logging.warning('No masking was detected (values != 0)')
          return False
          # continue
      # Decomposing the above, in base-2 bits
      aux_bit2_decomp = tuple(map(bit_decompose, aux_bit1))
      aux_bit2_count = [len(x) for x in aux_bit2_decomp]
      aux_bit2 = flatten_list(aux_bit2_decomp)
      aux_bit2 = np.unique(aux_bit2)
      # Check if these bits are in the dictionary of BPMDEF
      aux_xm_key = [get_dict_key(bpmdef, y) for y in aux_bit2]
      # If the list contains None, then some of the values weren't found
      if aux_xm_key.count(None):
          logging.warning('Some values in the mask are not defined in BPMDEF')
          aux_xm_key = [x for x in aux_xm_key if x is not None]
      # Flatten the list, because it is composed by lists inside the larger list
      aux_xm_key = flatten_list(aux_xm_key)
      logging.info('Found keys from BPMDEF: {0}'.format(aux_xm_key))
      #
      # Superpose masks, do not duplicate the flag, if already in the bits
      # composing the pixel mask value
      #
      bpm_twin = np.copy(bpm)
      # ====Run in parallel====
      # For Python 3+ the map can be applied in a easier way
      kw_super = {
          'bpm' : bpm,
          'aux_bit1' : aux_bit1,
          'aux_bit2_decomp' : aux_bit2_decomp,
          'non_masked_val' : non_masked_val,
          'mask_unmasked' : mask_unmasked,
      }
      if (box_msk.size == 0):
          logging.warning('No boxes to process. Exiting')
          exit(1)
      elif (box_msk.size == 1):
          t0 = time.time()
          mask_i_box = [bit_superpose(box_msk, **kw_super),]
          t1 =  time.time()
          txt_time = 'Time in comparing bits: {0:.2f} min'.format((t1 - t0) / 60.)
          logging.info(txt_time)
      elif (box_msk.size > 1):
          t0 = time.time()
          if (Nproc is None):
              Nproc = len(box_msk)
          P1 = mp.Pool(processes=Nproc)
          # Using the partial() method define an auxiliary function with the fixed
          # values. Then, in the map the iterating value will be provided
          part_bit_superpose = partial(bit_superpose, **kw_super)
          # The map blocks until the result is ready. If want to use apply_async()
          # or other method, processes could be freed when ends, without wait for
          # the others
          mask_i_box = P1.map(part_bit_superpose, box_msk)
          t1 =  time.time()
          txt_time = 'Time in comparing bits: {0:.2f} min'.format((t1 - t0) / 60.)
          logging.info(txt_time)
      # =======================
      # Replace values on the BPM array
      for i_box in mask_i_box:
          bpm_twin[~i_box.mask] = i_box[~i_box.mask]
      # Unmask the array
      bpm_final = np.ma.getdata(bpm_twin)
      #
      # Save into FITS
      # Use the original keys for CCDNUM and DES_EXT
      #
      if (out_fname is None):
          out_fname = 'bpm_modif_{0}'.format(fname_bpm_orig)
      #
      # Check if file already exists
      if (os.path.exists(out_fname) and file_check):
          logging.warning('File {0} already exists'.format(out_fname))
          while True:
              ask_txt = '\n\n===========================\n'
              ask_txt += 'File {0} already exists\n'.format(out_fname)
              ask_txt += 'Do you want to: \n\t-Overwrite? [Y]'
              ask_txt += '\n\t-Do not modify and exit? [N]'
              ask_txt += '\n\t-Add result as a new FITS extension? [A]'
              ask_txt += '\nAnswer [Y/N/A] (default=Y): '
              choice = raw_input(ask_txt).lower()
              choice = str.strip(choice)
              if (choice == 'y' or (choice == '')):
                  os.remove(out_fname)
                  break
              elif (choice == 'n'):
                  logging.info('Ok, exiting.')
                  exit(0)
                  break
              elif (choice == 'a'):
                  break
              else:
                  pass
      elif (os.path.exists(out_fname) and (not file_check)):
          logging.info('Overwritting {0}'.format(out_fname))
          os.remove(out_fname)
      #
      # Initialize object
      fits = fitsio.FITS(out_fname, 'rw')
      # Write array and header
      fits.write(bpm_final, dtype=out_dtype)#, header=header_orig)
      # Modify header
      comm1 = 'Wedge mask added by F Paz-Chinchon to {0}'.format(fname_bpm_orig)
      hlist = [
          {'name' : 'CCDNUM',
           'value' : header_orig['CCDNUM'],
           'comment' : 'CCD Number'
          },
          {'name' : 'DES_EXT',
           'value' : header_orig['DES_EXT'],
           'comment' : 'Extension type',
          },
          {'name' : 'DESMKBPM', 'value' : time.ctime(),
           'comment' : 'bad pixel map created'},
          {'name' : 'COMMENT',
           'value' : comm1},
      ]
      fits[0].write_keys(hlist)
      # Write checksum
      if checksum:
          fits[0].write_checksum()
      # Close FITS
      fits.close()
      logging.info('Modified BPM saved as {0}'.format(out_fname))
      #
      # Diagnosis plot
      #
      # Stacked mask
      if False:
          plt.close('all')
          fig = plt.figure(figsize=(5, 8))
          ax0 = fig.add_subplot(1, 1, 1)
          cmap = plt.cm.tab20
          a = np.ma.masked_where(bpm_twin == 0, bpm_twin)
          cmap.set_bad(color='white')
          im0 = ax0.imshow(a, cmap='tab20', origin='lower', vmin=256)
          plt.colorbar(im0)
          plt.show()
      #
      # Plot mask, showing in colormap the number of bits per pixel
      if False:
          aux_bpm = np.copy(bpm)
          for ind, item in enumerate(aux_bit1):
              aux_bpm[np.where(aux_bpm == item)] = aux_bit2_count[ind]
          plt.close('all')
          fig = plt.figure(figsize=(4, 8))
          ax = fig.add_subplot(1, 1, 1)
          im = ax.imshow(aux_bpm, origin='lower', cmap='viridis_r')
          plt.colorbar(im)
          plt.show()
          del aux_bpm
      return True
  
  
  if __name__ == '__main__':
      desc_txt = 'Code to add a mask to a BPM FITS file, based on coordinates'
      desc_txt += ' (given in an input table), and the BPMDEF flag to be assumed'
      desc_txt += ' by the masked pixels. It works per-CCD, one file at a time.'
      desc_txt += ' The input table must contain rows as (x1, y1, x2, y2,'
      desc_txt += ' bpmdef_flag). NOTE: it add the bits to the mask'
      abc = argparse.ArgumentParser(description=desc_txt)
      h0 = '<NOT YET IMPLEMENTED> Image to be loaded, to check the masking'
      abc.add_argument('--image', help=h0)
      tmp_bpm = 'D_n20170815t0824_c03_r3370p01_bpm.fits'
      h1 = 'BPM FITS file. Default: {0}'.format(tmp_bpm)
      abc.add_argument('--bpm', help=h1, default=tmp_bpm)
      tmp_add = 'vert_flag_op3.txt'
      h2 = 'Space-separated file with 2 vertices (bottom-left and top-right) per'
      h2 += ' box and its BPMDEF flag, to be be included in the BPM masking.'
      h2 += ' Format should be: x_ini y_ini x_end y_end bpmdef_flag'
      h2 += ' Default: {0}'.format(tmp_add)
      abc.add_argument('--tab', help=h2, default=tmp_add)
      h3 = 'How many lines skip at the top of the space-separated file defining'
      h3 += ' the boxes. Default: 1'
      abc.add_argument('--skip', help=h3, default=1, type=int)
      h4 = 'How many processes to run in parallel for the mask stacking. Default'
      h4 += ' is the number of boxes'
      abc.add_argument('--n1', help=h4, type=int)
      h5 = 'Value for non-masked pixels. Default: 0'
      abc.add_argument('--val', help=h5, type=int, default=0)
      h6 = 'Data type for the output BPM file. Default: uint16 (unsigned integer)'
      abc.add_argument('--dtype', help=h6, type=str, default='uint16')
      h7 = 'Flag to check for existing file, avoiding overwrite (prompt is'
      h7 += ' displayed). Default is to overwrite'
      abc.add_argument('--check', help=h7, action='store_true')
      h8 = 'Output filename. Default is bpm_modif_{filename of input BPM}.fits'
      abc.add_argument('--out', help=h8)
      h9 = 'Flag to write checksums in FITS header. Default not to write them'
      abc.add_argument('--csum', help=h9, action='store_true')
      # Variables
      abc = abc.parse_args()
      #
      # aux_img, header_img = open_fits(abc.image)
      aux_bpm, header_bpm = open_fits(abc.bpm)
      aux_tab = open_txt(abc.tab, abc.skip)
      # We expect one extension.
      if (len(aux_bpm) != 1):
          logging.error('This code is implemented to update one-extension only')
          exit(1)
      # Call different flags, for different boxes. One flag at a time
      aux_mask = joint_mask(aux_bpm[0], header_bpm[0], aux_tab, abc.bpm,
                            Nproc=abc.n1, non_masked_val=abc.val,
                            out_dtype=abc.dtype, file_check=abc.check,
                            out_fname=abc.out, checksum=abc.csum)
  