content: # -*- coding: utf-8 -*-
  """
  Combines all the images for a given (TARGET, FILTER, HWP) combination to
  produce a single, average image.
  
  Estimates the sky background level of the on-target position at the time of the
  on-target observation using a bracketing pair of off-target observations through
  the same HWP polaroid rotation value. Subtracts this background level from
  each on-target image to produce background free images. Applies an airmass
  correction to each image, and combines these final image to produce a background
  free, airmass corrected, average image.
  """
  
  # Core imports
  import os
  import sys
  import copy
  import warnings
  
  # Import scipy/numpy packages
  import numpy as np
  from scipy import ndimage
  
  # Import astropy packages
  from astropy.table import Table
  import astropy.units as u
  from astropy.convolution import Gaussian2DKernel
  from astropy.modeling import models, fitting
  from astropy.stats import gaussian_fwhm_to_sigma, sigma_clipped_stats
  from photutils import (make_source_mask,
      MedianBackground, SigmaClip, Background2D)
  
  # Import plotting utilities
  from matplotlib import pyplot as plt
  
  # Add the AstroImage class
  import astroimage as ai
  
  # Add the header handler to the BaseImage class
  from Mimir_header_handler import Mimir_header_handler
  ai.reduced.ReducedScience.set_header_handler(Mimir_header_handler)
  ai.set_instrument('mimir')
  
  #==============================================================================
  # *********************** CUSTOM USER CODE ************************************
  # this is where the user specifies where the raw data is stored
  # and some of the subdirectory structure to find the actual .FITS images
  #==============================================================================
  
  # This is a list of targets for which to process each subgroup (observational
  # group... never spanning multiple nights, etc...) instead of combining into a
  # single "metagroup" for all observations of that target. The default behavior
  # is to go ahead and combine everything into a single, large "metagroup". The
  # calibration data should probably not be processed as a metagroup though.
  processSubGroupList = []
  processSubGroupList = [t.upper() for t in processSubGroupList]
  
  # Define the location of the PPOL reduced data to be read and worked on
  PPOL_data = 'C:\\Users\\Jordan\\FITS_data\\Mimir_data\\PPOL_Reduced\\201611\\'
  S3_dir    = os.path.join(PPOL_data, 'S3_Astrometry')
  
  # This is the location where all pyPol data will be saved
  pyPol_data = 'C:\\Users\\Jordan\\FITS_data\\Mimir_data\\pyPol_Reduced\\201611'
  
  # This is the location of the previously generated masks (step 4)
  maskDir = os.path.join(pyPol_data, 'Masks')
  
  # Setup new directory for polarimetry data
  polarimetryDir = os.path.join(pyPol_data, 'Polarimetry')
  if (not os.path.isdir(polarimetryDir)):
      os.mkdir(polarimetryDir, 0o755)
  
  HWPDir = os.path.join(polarimetryDir, 'HWPImgs')
  if (not os.path.isdir(HWPDir)):
      os.mkdir(HWPDir, 0o755)
  
  bkgPlotDir = os.path.join(HWPDir, 'bkgPlots')
  if (not os.path.isdir(bkgPlotDir)):
      os.mkdir(bkgPlotDir, 0o755)
  
  # # Setup PRISM detector properties
  # read_noise = 13.0 # electrons
  # effective_gain = 3.3 # electrons/ADU
  
  #########
  ### Establish the atmospheric extinction (magnitudes/airmass)
  #########
  # Following table from Hu (2011)
  # Data from Gaomeigu Observational Station
  # Passband | K'(lambda) [mag/airmass] | K'' [mag/(color*airmass)]
  # U	         0.560 +/- 0.023            0.061 +/- 0.004
  # B          0.336 +/- 0.021            0.012 +/- 0.003
  # V          0.198 +/- 0.024           -0.015 +/- 0.004
  # R          0.142 +/- 0.021           -0.067 +/- 0.005
  # I          0.093 +/- 0.020            0.023 +/- 0.006
  
  
  # Following table from Schmude (1994)
  # Data from Texas A & M University Observatory
  # Passband | K(lambda) [mag/airmass] | dispersion on K(lambda)
  # U	         0.60 +/- 0.05             0.120
  # B          0.40 +/- 0.06             0.165
  # V          0.26 +/- 0.03             0.084
  # R          0.19 +/- 0.03             0.068
  # I          0.16 +/- 0.02             0.055
  
  # TODO: Ask Dan about atmospheric extinction from airmass at NIR
  kappa = dict(zip(['U',  'B',  'V',  'R',  'I',  'J',  'H',  'K'  ],
                   [0.60, 0.40, 0.26, 0.19, 0.16, 0.05, 0.01, 0.005]))
  
  
  # Read in the indexFile data and select the filenames
  print('\nReading file index from disk')
  indexFile = os.path.join(pyPol_data, 'reducedFileIndex.csv')
  fileIndex = Table.read(indexFile, format='ascii.csv')
  
  # Determine which parts of the fileIndex pertain to HEX dither science images
  useFiles    = np.logical_and(
      fileIndex['USE'] == 1,
      fileIndex['DITHER_TYPE'] == 'ABBA'
  )
  useFileRows = np.where(useFiles)
  
  # Cull the file index to only include files selected for use
  fileIndex = fileIndex[useFileRows]
  
  # Define an approximate pixel scale
  pixScale  = 0.5789*(u.arcsec/u.pixel)
  
  # TODO: implement a FWHM seeing cut... not yet working because PSF getter seems
  # to be malfunctioning in step 2
  #
  #
  # # Loop through each unique GROUP_ID and test for bad seeing conditions.
  # groupByID = fileIndex.group_by(['GROUP_ID'])
  # for subGroup in groupByID.groups:
  #     # Grab the FWHM values for this subGroup
  #     thisFWHMs = subGroup['FWHM']*u.pixel
  #
  #     # Grab the median and standard deviation of the seeing for this subgroup
  #     medianSeeing = np.median(thisFWHMs)
  #     stdSeeing    = np.std(thisFWHMs)
  #
  #     # Find bad FWHM values
  #     badFWHMs = np.logical_not(np.isfinite(subGroup['FWHM']))
  #     badFWHMs = np.logical_or(
  #         badFWHMs,
  #         thisFWHMs <= 0
  #     )
  #     badFWHM = np.logical_and(
  #         badFWHM,
  #         thisFWHMs > 2.0*u.arcsec
  #     )
  #     import pdb; pdb.set_trace()
  
  # Group the fileIndex by...
  # 1. Target
  # 2. Waveband
  fileIndexByTarget = fileIndex.group_by(['TARGET', 'FILTER'])
  
  # Loop through each group
  for group in fileIndexByTarget.groups:
      # Grab the current group information
      thisTarget     = str(np.unique(group['TARGET'].data)[0])
      thisFilter     = str(np.unique(group['FILTER'].data)[0])
  
      # # Skip the Merope nebula for now... not of primary scientific importance
      # if thisTarget == 'MEROPE': continue
  
      # Update the user on processing status
      print('\nProcessing images for')
      print('Target : {0}'.format(thisTarget))
      print('Filter : {0}'.format(thisFilter))
  
      # Grab the atmospheric extinction coefficient for this wavelength
      thisKappa = kappa[thisFilter]
  
      # Further divide this group by its constituent HWP values
      indexByPolAng = group.group_by(['IPPA'])
  
      # Loop over each of the HWP values, as these are independent from
      # eachother and should be treated entirely separately from eachother.
      for IPPAgroup in indexByPolAng.groups:
          # Grab the current HWP information
          thisIPPA = np.unique(IPPAgroup['IPPA'].data)[0]
  
          # Update the user on processing status
          print('\tIPPA : {0}'.format(thisIPPA))
  
          # For ABBA dithers, we need to compute the background levels on a
          # sub-group basis. If this target has not been selected for subGroup
          # averaging, then simply append the background subtracted images to a
          # cumulative list of images to align and average.
  
          # Initalize an image list to store all the images for this
          # (target, filter, pol-ang) combination
          imgList = []
  
          indexByGroupID = IPPAgroup.group_by(['GROUP_ID'])
          for subGroup in indexByGroupID.groups:
              # Grab the numae of this subGroup
              thisSubGroup = str(np.unique(subGroup['OBJECT'])[0])
  
              # if (thisSubGroup != 'NGC2023_R1') and (thisSubGroup != 'NGC2023_R2'): continue
  
              # Construct the output file name and test if it alread exsits.
              if thisTarget in processSubGroupList:
                  outFile = '_'.join([thisTarget, thisSubGroup, str(thisIPPA)])
                  outFile = os.path.join(HWPDir, outFile) + '.fits'
              elif thisTarget not in processSubGroupList:
                  outFile = '_'.join([thisTarget, thisFilter, str(thisIPPA)])
                  outFile = os.path.join(HWPDir, outFile) + '.fits'
  
              # Test if this file has already been constructed and either skip
              # this subgroup or break out of the subgroup loop.
              if os.path.isfile(outFile):
                  print('\t\tFile {0} already exists...'.format(os.path.basename(outFile)))
                  if thisTarget in processSubGroupList:
                      continue
                  elif thisTarget not in processSubGroupList:
                      break
  
              # Update the user on the current execution status
              print('\t\tProcessing images for subgroup {0}'.format(thisSubGroup))
  
              # Initalize lists to store the A and B images.
              AimgList  = []
              BimgList  = []
  
              # Initalize a list to store the off-target sky background levels
              BbkgList  = []
  
              # Initilaze lists to store the times of observation
              AdatetimeList = []
              BdatetimeList = []
  
              # Read in all the images for this subgroup
              progressString = '\t\tNumber of Images : {0}'
              for iFile, filename in enumerate(subGroup['FILENAME']):
                  # Update the user on processing status
                  print(progressString.format(iFile+1), end='\r')
  
                  # Read in a temporary compy of this image
                  PPOL_file = os.path.join(S3_dir, filename)
                  tmpImg = ai.reduced.ReducedScience.read(PPOL_file)
  
                  # Crop the edges of this image
                  ny, nx = tmpImg.shape
                  binningArray = np.array(tmpImg.binning)
  
                  # Compute the amount to crop to get a 1000 x 1000 image
                  cy, cx = (ny - 1000, nx - 1000)
  
                  # Compute the crop boundaries and apply them
                  lf = np.int(np.round(0.5*cx))
                  rt = lf + 1000
                  bt = np.int(np.round(0.5*cy))
                  tp = bt + 1000
  
                  tmpImg = tmpImg[bt:tp, lf:rt]
  
                  # Grab the on-off target value for this image
                  thisAB = subGroup['AB'][iFile]
  
                  # Place the image in a list and store required background values
                  if thisAB == 'B':
                      # Place B images in the BimgList
                      BimgList.append(tmpImg)
  
                      # Place the median value of this off-target image in list
                      mask = make_source_mask(
                          tmpImg.data, snr=2, npixels=5, dilate_size=11
                          )
                      mean, median, std = sigma_clipped_stats(
                          tmpImg.data, sigma=3.0, mask=mask
                      )
                      BbkgList.append(median)
  
                      # Place the time of this image in a list of time values
                      BdatetimeList.append(tmpImg.julianDate)
  
                  if thisAB == 'A':
                      # Read in any associated masks and store them.
                      maskFile = os.path.join(maskDir, os.path.basename(filename))
  
                      # If there is a mask for this file, then apply it!
                      if os.path.isfile(maskFile):
                          # Read in the mask file
                          tmpMask = ai.reduced.ReducedScience.read(maskFile)
  
                          # Crop the mask to match the shape of the original image
                          tmpMask = tmpMask[cy:ny-cy, cx:nx-cx]
  
                          # Grab the data to be masked
                          tmpData = tmpImg.data
  
                          # Mask the data and put it back into the tmpImg
                          maskInds = np.where(tmpMask.data)
                          tmpData[maskInds] = np.NaN
                          tmpImg.data = tmpData
  
                      # Place B images in the BimgList
                      AimgList.append(tmpImg)
  
                      # Place the time of this image in a list of time values
                      AdatetimeList.append(tmpImg.julianDate)
  
              # Create a new line for shell output
              print('')
  
              # Construct an image stack of the off-target images
              BimageStack = ai.utilitywrappers.ImageStack(BimgList)
  
              # Build a supersky image from these off-target images
              superskyImage = BimageStack.produce_supersky()
  
              import pdb; pdb.set_trace()
  
              # Locate regions outside of a 5% deviation
              tmpSuperskyData  = superskyImage.data
              maskedPix  = np.abs(tmpSuperskyData - 1.0)  > 0.05
  
              # Get rid of the small stuff and expand the big stuff
              maskedPix = ndimage.binary_opening(maskedPix, iterations=2)
              maskedPix = ndimage.binary_closing(maskedPix, iterations=2)
              maskedPix = ndimage.binary_dilation(maskedPix, iterations=4)
  
              # TODO: Make the box_size and filter_size sensitive to binning.
              binningArray = np.array(superskyImage.binning)
              box_size = tuple((100/binningArray).astype(int))
              filter_size = tuple((10/binningArray).astype(int))
  
              # Setup the sigma clipping and median background estimators
              sigma_clip = SigmaClip(sigma=3., iters=10)
              bkg_estimator = MedianBackground()
  
              # Compute a smoothed background image
              bkgData = Background2D(superskyImage.data,
                  box_size=box_size, filter_size=filter_size, mask=maskedPix,
                  sigma_clip=sigma_clip, bkg_estimator=bkg_estimator)
  
              # Construct a smoothed supersky image object
              smoothedSuperskyImage = ai.reduced.ReducedScience(
                  bkgData.background/bkgData.background_median,
                  uncertainty = bkgData.background_rms,
                  properties={'unit':u.dimensionless_unscaled}
              )
  
              # Interpolate background values to A times
              AbkgList = np.interp(
                  AdatetimeList,
                  BdatetimeList,
                  BbkgList,
                  left=-1e6,
                  right=-1e6
              )
  
              # Cut out any extrapolated data (and corresponding images)
              goodInds = np.where(AbkgList > -1e5)
              AimgList = np.array(AimgList)[goodInds]
              AdatetimeList = np.array(AdatetimeList)[goodInds]
              AbkgList = AbkgList[goodInds]
  
              AsubtractedList = []
              # Loop through the on-target images and subtract background values
              for Aimg, Abkg in zip(AimgList, AbkgList):
                  # Subtract the interpolated background values from the A images
                  tmpImg = Aimg - smoothedSuperskyImage*(Abkg*Aimg.unit)
  
                  # Apply an airmass correction
                  tmpImg = tmpImg.correct_airmass(thisKappa)
  
                  # Append the subtracted and masked image to the list.
                  AsubtractedList.append(tmpImg)
  
              # Now that the images have been fully processed, pause to generate
              # a plot to store in the "background plots" folder. These plots
              # constitute a good sanity check on background subtraction.
              plt.plot(BdatetimeList, BbkgList, '-ob')
              plt.scatter(AdatetimeList, AbkgList, marker='o', facecolor='r')
              plt.xlabel('Julian Date')
              plt.ylabel('Background Value [ADU]')
              figName = '_'.join([thisTarget, thisSubGroup, str(thisIPPA)])
              figName = os.path.join(bkgPlotDir, figName) + '.png'
              plt.savefig(figName, dpi=300)
              plt.close('all')
  
              # Here is where I need to decide if each subgroup image should be
              # computed or if I should just continue with the loop.
              if thisTarget.upper() in processSubGroupList:
                  # Construct an image combiner for the A images
                  AimgStack = ai.utilitywrappers.ImageStack(AsubtractedList)
  
                  # Align the images
                  AimgStack.align_images_with_wcs(
                      subPixel=False,
                      padding=np.NaN
                      )
  
                  # Combine the images
                  AoutImg = imgStack.combine_images()
  
                  # Save the image
                  AoutImg.write(outFile, dtype=np.float64)
              else:
                  # Extend the imgList variable with background corrected images
                  imgList.extend(AsubtractedList)
  
  
          if len(imgList) > 0:
              # At the exit of the loop, process ALL the files from ALL the groups
              # Construct an image combiner for the A images
              imgStack = ai.utilitywrappers.ImageStack(imgList)
  
              # Align the images
              imgStack.align_images_with_wcs(
                  subPixel=False,
                  padding=np.NaN
                  )
  
              # Combine the images
              outImg = imgStack.combine_images()
  
              # Save the image
              outImg.write(outFile, dtype=np.float64)
  
  print('\nDone computing average images!')
  