#!/usr/bin/env python
#
# 2007 Jul 8  - Andrew W Stephens - alpha version
# 2007 Jul 9  - AWS - beta version
# 2007 Jul 10 - AWS - move most operations to cleanquad function
# 2007 Jul 11 - AWS - use stddev to decide whether to keep orig quad
# 2007 Jul 14 - AWS - generalized code to allow different pattern sizes
# 2007 Jul 18 - AWS - fix bug generating index arrays
# 2007 Jul 20 - AWS - add quadrant bias-level normalization
# 2007 Jul 23 - AWS - add force option
# 2007 Aug 06 - AWS - f/6 spectroscopy mode: use top & bottom for pattern
# 2007 Aug 22 - AWS - add -a flag to use all pixels
# 2008 Jan 11 - AWS - check for available image extensions
# 2008 Feb 05 - AWS - don't close input file until the end (req'd if next>2)
# 2008 Oct 02 - AWS - don't write pattern unless given the -p flag
# 2009 May 03 - AWS - use conformant default output file name
# 2009 May 13 - AWS - verify FITS header (old images have unquoted release date)
# 2009 May 22 - AWS - output full-frame pattern
# 2009 May 22 - AWS - improve quadrant bias normalization
# 2009 May 23 - AWS - add optional sky frame
# 2009 May 26 - AWS - add user-supplied bias offset correction
# 2009 Jul 04 - AWS - do not try to bias correct spectral flats
# 2009 Oct 24 - AWS - add basic row filtering
# 2009 Nov 06 - AWS - ignore bad pixels flagged in DQ extension
# 2009 Nov 08 - AWS - use mode for quadrant bias level normalization
# 2009 Nov 12 - AWS - use sigma-clipped stddev to judge quality of bias normalization
# 2009 Nov 17 - AWS - fit a Gaussian to the sky pixel distribution for bias norm.
# 2010 Feb 02 - AWS - sky subtract before quadrant normalization
# 2010 Feb 18 - AWS - add sigma-clipping to row filtering
# 2010 Apr 09 - AWS - only check for gcal status if OBSTYPE = FLAT
# 2010 Apr 10 - AWS - accept list input
# 2010 Apr 13 - AWS - minor tweak of the spectroscopic regions
# 2010 Jul 11 - AWS - allow images sizes which are multiples of the pattern size
# 2010 Oct 08 - AWS - option to read in bad pixel mask (e.g. object mask from nisky)
# 2010 Oct 10 - AWS - change the -g flag to take arguments
# 2010 Oct 11 - AWS - pad GNIRS images (2 row at the top)
# 2010 Oct 12 - AWS - GNIRS row filtering using an 8-pixel wide kernel
# 2010 Dec 21 - AWS - add grid filter
# 2010 Dec 28 - AWS - select GNIRS pattern region based on camera & slit
# 2011 Feb 03 - AWS - use extension 2 for nsprepared GNIRS data
# 2011 Feb 05 - AWS - add input glob expansion
# 2011 May 05 - AWS - output 32-bit files
# 2011 Jun 17 - AWS - catch modified GNIRS XD decker name
# 2012 Mar 29 - AWS - use the mode instead of median for pattern determination
# 2012 May 17 - AWS - do not modify input flag values
# 2013 Jun 16 - AWS - allow processing of non-standard FITS (with a warning)
# 2013 Jun 22 - AWS - add option to ignore DQ plane
# 2013 Jun 22 - AWS - use whole array if SKYIMAGE header keyword is present

# To Do:
# GNIRS: Mask out padding when a DQ or pixel mask is available
# Detect and mask out objects before doing any calculations
# check if run before
# properly handle images < 1024 x 1024
# Specification of image section to use to calculate pattern
# Specification of image section affected by pattern
# Look at stddev of each row to identify which have pattern noise

#-----------------------------------------------------------------------

import datetime
import getopt
import glob
import matplotlib.pyplot as pyplot
import numpy
import os
from astropy.io import fits as pyfits
from scipy.optimize import leastsq
import string
import sys

version = '2013 Jun 22'

#-----------------------------------------------------------------------

def usage():
    print ''
    print 'NAME'
    print '       cleanir.py - filter pattern noise out of NIRI and GNIRS frames\n'
    print 'SYNOPSIS'
    print '       cleanir.py [options] infile/list\n'
    print 'DESCRIPTION'
    print '       This script assumes that the NIRI/GNIRS pattern noise in a quadrant'
    print '       can be represented by a fixed pattern which is repeated over the'
    print '       entire quadrant.  The default size for this pattern is 16 pixels'
    print '       wide and 4 pixels high (which may be changed via the -x and -y'
    print '       flags).  The pattern is determined for each quadrant by taking the'
    print '       mode of the pixel value distribution at each position in the'
    print '       pattern.  Once the pattern has been determined for a quadrant'
    print '       it is replicated to cover the entire quadrant and subtracted,'
    print '       and the mean of the pattern is added back to preserve flux.'
    print '       The standard deviation of all the pixels in the quadrant is'
    print '       compared to that before the pattern subtraction, and if no'
    print '       reduction was achieved the subtraction is undone.  The pattern'
    print '       subtraction may be forced via the -f flag.  This process is'
    print '       repeated for all four quadrants and the cleaned frame is written'
    print '       to c<infile> (or the file specified with the -o flag).  The'
    print '       pattern derived for each quadrant may be saved with the -p flag.'
    print ''
    print '       Pattern noise is often accompanied by an offset in the bias'
    print '       values between the four quadrants.  One may want to use the'
    print '       -q flag to try to remove this offset.  This attempts to match'
    print '       the iteratively determined median value of each quadrant.'
    print '       This method works best with sky subtraction (i.e. with the -s'
    print '       flag), and does not work well if there are large extended objects'
    print '       in the frame.  By default the median is determined from the'
    print '       entire frame, although the -c flag will only use a central'
    print '       portion of the image.  Note that the derived quadrant offsets'
    print '       will be applied to the output pattern file.'
    print ''
    print '       Removing the pattern from spectroscopy is more difficult because'
    print '       of many vertical sky lines.  By default NIRI f/6 spectroscopy with'
    print '       the 2-pixel or blue slits (which do not fill the detector), uses the'
    print '       empty regions at the bottom (1-272) and top (720-1024) of the'
    print '       array for measuring the pattern.  This is not possible for other'
    print '       modes of spectroscopy where the spectrum fills the detector.'
    print '       For these modes it is best to do sky subtraction before pattern'
    print '       removal.  The quickest method is to pass a sky frame (or an offset'
    print '       frame) via the -s flag.  The manual method is to generate and'
    print '       subtract the sky, determine and save the pattern via the -p flag,'
    print '       then subtract the pattern from the original image.  One may use'
    print '       the -a flag to force using all of the pixels for the pattern'
    print '       determination.  If the SKYIMAGE FITS header keyword is present'
    print '       it is assumed that the sky has already been subtracted and all'
    print '       pixels will be used for the pattern determination.'
    print ''
    print '       Note that you may use glob expansion in infile, however, the'
    print '       entire string must then be quoted or any pattern matching'
    print '       characters (*,?) must be escaped with a backslash.'
    print ''
    print 'OPTIONS'
    print '       -a : use all pixels for pattern determination'
    print '       -b <badpixelmask> : specify a bad pixel mask (overrides DQ plane)'
    print '       -c <frac> : use central <fraction> of image for bias adjustment [1]'
    print '       -d <dir> : specify an input data directory'
    print '       -f : force cleaning of all quads even if stddev does not decrease'
    print '       -g # : graph results (0=none, 1=pattern, 2=offsets, 3=both)'
    print '       -m : use median for quadrant normalization (instead of fitting a Gaussian)'
    print '       --nodq : ignore the DQ plane'
    print '       -o <file> : write output to <file> (instead of c<infile>)'
    print '       -p <file> : write full-frame pattern to <file>'
    print '       -q : adjust quadrant offsets'
    print '       -r : row filtering (useful for GNIRS XD spectra)'
    print '       -s <sky> : sky frame to help in pattern recognition'
    print '       -t : apply test grid filter before normalizing quadrants'
    print '       -v : verbose debugging output'
    print '       -x <size> : set pattern x size in pix [16]'
    print '       -y <size> : set pattern y size in pix [4]\n'
    print 'VERSION'
    print '       ', version
    print ''
    raise SystemExit

#-----------------------------------------------------------------------

def main():
    global allpixels, applygridfilter
    global bad, badmask, bias1, bias2, bias3, bias4, biasadjust
    global cfrac, datadir, force, graph
    global output, patternfile, pxsize, pysize, quadmedian, rowfilter
    global savepattern, skyfile, subtractsky, subtractrowmedian
    global usedq, verbose

    try:
        opts, args = getopt.getopt(sys.argv[1:], 'ab:c:d:fg:hmo:p:qrs:tx:y:v', ['nodq','q1=','q2=','q3=','q4='])
    except getopt.GetoptError, err:
        print str(err)
        usage()
        sys.exit(2)

    if (len(args)) != 1:
        usage()

    nargs = len(sys.argv[1:])
    nopts = len(opts)

    allpixels = False
    applygridfilter = False
    bad = -9.e6              # value assigned to bad pixels
    badmask = 'DQ'
    bias1 = 0.0
    bias2 = 0.0
    bias3 = 0.0
    bias4 = 0.0
    biasadjust = False
    cfrac = 1.0              # use whole image
    datadir = ''
    force = False
    graph = 0
    output = 'default'
    patternfile = ''
    pxsize = 16
    pysize = 4
    quadmedian = False
    rowfilter = False
    savepattern = False
    skyfile = ''
    subtractsky = False
    subtractrowmedian = False
    usedq = True
    verbose = False

    for o, a in opts:
        if o in ('-a'):        # force using all pixels for pattern determination
            allpixels = True
        elif o in ('-b'):
            badmask = a
        elif o in ('-c'):      # use central fraction for bias normalization
            cfrac = float(a)
        elif o in ('-d'):      # input data directory
            datadir = a
        elif o in ('-f'):      # force pattern subtraction in every quadrant
            force = True
        elif o in ('-g'):      # graph results
            graph = int(a)
        elif o in ('-o'):      # specify cleaned output file
            output = a
        elif o in ('-m'):
            quadmedian = True
        elif o in ('--nodq'):
            usedq = False
        elif o in ('-p'):      # write pattern file
            patternfile = a
            savepattern = True
        elif o in ('-q'):      # try to adjust quadrant bias values
            biasadjust = True
        elif o in ('--q1'):    # bias offset for quadrant 1
            bias1 = float(a)
        elif o in ('--q2'):
            bias2 = float(a)
        elif o in ('--q3'):
            bias3 = float(a)
        elif o in ('--q4'):
            bias4 = float(a)
        elif o in ('-r'):      # row filtering
            rowfilter = True
        elif o in ('-s'):      # sky frame
            skyfile = a
            subtractsky = True
        elif o in ('-t'):      # test grid filter
            applygridfilter = True
        elif o in ('-x'):      # specify pattern x-dimension
            pxsize = int(a)
        elif o in ('-y'):      # specify pattern y-dimension
            pysize = int(a)
        elif o in ('-v'):      # verbose debugging output
            verbose = True
        else:
            assert False, "unhandled option"

    inputfile = args[0]

    files = glob.glob(inputfile)
    if verbose:
        print '...input = ', inputfile
        print '...files = ', files

    print ''

    for f in files:
        if IsFits(f):
            cleanir(f)
        else: # file list
            print 'Expanding ' + f + '...\n'
            inlist = open(f,'r')
            for line in inlist:
                cleanir(line.strip())
            inlist.close()

#-----------------------------------------------------------------------

def IsFits(infile):
    global datadir

    # If the file exists and has a .fits extension assume that it is FITS:
    if os.path.exists(datadir + infile):
        if infile.endswith('.fits'):
            fits = True
        else:
            fits = False
    elif os.path.exists(infile): # Check for lists in the CWD
        if infile.endswith('.fits'):
            fits = True
        else:
            fits = False
    else: # assume it is a FITS image missing the .fits extension
        fits = True
    return fits

#-----------------------------------------------------------------------

def IterStat(vector, lowsigma=3, highsigma=3):
    global verbose
    median = numpy.median(vector)
    stddev = numpy.std(vector)
    minval = median - lowsigma  * stddev
    maxval = median + highsigma * stddev
    num = numpy.size(vector)
    dn = 1000
    while (dn > 1 and stddev > 0 ):
        tmpvec = vector[(vector>minval) & (vector<maxval)]
        median = numpy.median(tmpvec)
        stddev = numpy.std(tmpvec)
        dn = num - numpy.size(tmpvec)
        num = numpy.size(tmpvec)
        if verbose:
            print '   ...median=',median,' stddev=',stddev,' min=',minval,' max=',maxval,' N=',num,' dN=',dn
        minval = median - lowsigma  * stddev
        maxval = median + highsigma * stddev

    return (median, stddev)

#-----------------------------------------------------------------------

def Mode(x):
    if x.size < 1:
        print 'ERROR: cannot determine the mode of an empty array'
        sys.exit(2)
    y = numpy.sort(x)
    starts = numpy.concatenate(([1], numpy.diff(y).astype(bool).astype(int)))
    starts_sum = starts.cumsum()
    counts = numpy.bincount(starts_sum)
    arg_mode_freq = counts.argmax()
    counts_sum = counts.cumsum()
    mode = y[counts_sum[arg_mode_freq-1]]
    if counts[arg_mode_freq] >= 20:
        return mode
    else:
        print 'WARNING: low mode frequency =', counts[arg_mode_freq], ' switching to median'
        return numpy.median(x)

#-----------------------------------------------------------------------

def CleanQuad(quad,patternin):
    # quad = quadrant to be pattern-subtracted
    # patternin = region to use for pattern determination
    global qxsize, qysize                                 # quadrant size
    global pxsize, pysize                                 # pattern size
    global inputmedian, inputstddev

    if verbose:
        print '   ...mean of input quadrant =', numpy.mean(quad)
        print '   ...median of input quadrant =', numpy.median(quad)

    # create arrays of indices which correspond to the pattern tiled over
    # the region of the input quadrant to be used for pattern determination
    inpx = len(patternin[0])
    inpy = len(patternin)
    if verbose:
        print '   ...size of pattern determination region =',inpx,'x',inpy
    indx = numpy.tile(numpy.arange(0,inpx,pxsize), inpy/pysize)
    indy = numpy.arange(0,inpy,pysize).repeat(inpx/pxsize)
    if verbose:
        print '   ...indx:', indx
        print '   ...indy:', indy

    # create blank pattern array:
    pattern = numpy.zeros(pysize*pxsize).reshape(pysize,pxsize)
    origstddev = numpy.std(quad)
    print '   ...standard deviation of input quadrant =%9.3f' % origstddev

    # find median pattern across quadrant:
    if (graph > 0):
        binwidth = 1
        binmin = inputmedian - 5. * inputstddev
        binmax = inputmedian + 5. * inputstddev
        #binmax = numpy.max(quad) # show all the data
        bins = numpy.arange( binmin, binmax, binwidth )
        bincenters = bins[1:bins.size] - binwidth/2.
        iplot = 0

    for iy in range(0, pysize):
        for ix in range(0, pxsize):
            tmpdata = patternin[indy+iy,indx+ix]
            #pattern[iy,ix] = numpy.median(tmpdata[tmpdata!=bad]) # filter out bad pix
            pattern[iy,ix] = Mode(tmpdata[tmpdata!=bad])

            if (graph==1 or graph==3):
                iplot += 1
                hist,bins = numpy.histogram(tmpdata, bins=bins)
                plot = pyplot.subplot(pysize,pxsize,iplot)
                pyplot.plot(bincenters, hist, linestyle='', marker='.')
                pyplot.axvline(x=pattern[iy,ix], ls='--', color='green')
                if ix != 0:
                    plot.set_yticklabels([])

    if verbose:
        print '...pattern:', pattern

    if (graph==1 or graph==3):
        print ('...graphing results...')
        pyplot.subplots_adjust(left=0.05, bottom=0.05, right=0.95, top=0.95, wspace=0., hspace=0.2)
        pyplot.show()

    # tile pattern over quadrant:
    quadpattern = numpy.tile(pattern, (qysize/pysize, qxsize/pxsize))
    quadpattern -= numpy.mean(pattern)      # set the mean value to zero
    #print '   ...mean of pattern = ', numpy.mean(quadpattern)

    cleanedquad = quad - quadpattern                      # subtract pattern
    cleanstddev = numpy.std(cleanedquad)      # calculate standard deviation

    print '   ...standard deviation of cleaned quadrant = %.3f' % cleanstddev

    if (force):
        print '   ...forcing pattern subtraction'
    else:
        # has subtracting the pattern reduced the standard deviation?
        if ( origstddev - cleanstddev > 0.01 ):
            print '   ...improvement!'
        else:
            print '   ...no significant improvement so using the original quadrant'
            cleanedquad = quad      # reset quadrant pixels to original values
            quadpattern = quadpattern * 0         # set pattern to zeros

    return cleanedquad, quadpattern

#-----------------------------------------------------------------------

def CleanRow(row, sample, value):
    # row = row to be pattern-subtracted
    # sample = sample used to measure pattern
    # value = desired final median value of row

    indx = numpy.arange(0,len(sample),8)
    pattern = numpy.zeros(8)
    for ix in range(0, 8):
        tmpdata = sample[indx+ix]
        tmpdata = tmpdata[tmpdata!=bad] # filter out bad pix
        # pattern[ix] = numpy.median(tmpdata)
        median,stddev = IterStat(tmpdata)
        pattern[ix] = median

    if verbose:
        print '...pattern:', pattern

    # repeat the pattern over the row and subtract:
    rowpattern = numpy.tile(pattern, len(row)/8)
    cleanedrow = row - rowpattern + value

    #median,stddev = IterStat(cleanedrow, lowsigma=3, highsigma=1)
    #cleanedrow = cleanedrow + (value - median)

    return cleanedrow

#-----------------------------------------------------------------------

def ApplyRowFilter(quad, patternin):
    # quad = quadrant to be pattern-subtracted
    # patternin = region to use for pattern determination
    global qxsize, qysize                                 # quadrant size

    median,stddev = IterStat(patternin)           # iterative median
    print '...median of input sample quadrant =', median, '+/-', stddev

    for iy in range(0,qysize):  # this is not correct, but will work for GNIRS
        if verbose:
            print '...row =', iy
        quad[iy] = CleanRow(quad[iy], patternin[iy], median)
        #quad[iy] = CleanRow(quad[iy], patternin[iy], inputmedian)

    return quad

#-----------------------------------------------------------------------

def gaussian(t,p):    # p[0] = mu   p[1] = sigma    p[2] = peak
    return(p[2] * numpy.exp( -(t - p[0])**2 / (2 * p[1]**2) ))

def residuals(p,data,t):
    err = data - gaussian(t,p)
    return err

def NormQuadMedian(quad):
    global bins, bincenters, inputmean, inputmedian, inputstddev
    global lsigma, hsigma, bad
    hist,bins = numpy.histogram(quad, bins=bins)
    if verbose:
        print '...calculating median using low-sigma =',lsigma,' and high-sigma =',hsigma
    mincts = inputmedian - lsigma * inputstddev
    maxcts = inputmedian + hsigma * inputstddev
    if verbose:
        print '...input median=',inputmedian,' min=',mincts,' max=',maxcts
    flatquad = quad[quad != bad]  # flatten array and filter out bad pix
    npix = numpy.size(flatquad)
    dn = 100
    while (npix > 10000 and dn > 10):
        tmpquad = flatquad[(flatquad>mincts) & (flatquad<maxcts)]
        median = numpy.median(tmpquad)
        stddev = numpy.std(tmpquad)
        mincts = median - lsigma * stddev
        maxcts = median + hsigma * stddev
        dn = npix - numpy.size(tmpquad)
        npix = numpy.size(tmpquad)
        if verbose:
            print '...median=',median,' stddev=',stddev,' min=',mincts,' max=',maxcts,' npix=',npix,' dn=',dn
    offset = inputmedian - median
    print '   ...offset = %.3f' % offset
    return hist, median, offset

#-----------------------------------------------------------------------

def NormQuadGauss(quad):
    global bins, bincenters, inputmean, inputmedian, inputstddev
    hist,bins = numpy.histogram(quad, bins=bins)
    mode = bins[ hist.argmax() ]
    peak = hist.max()
    fitsigma = 1.0      # this should probably be a command-line parameter
    mincts = mode - fitsigma * inputstddev
    maxcts = mode + fitsigma * inputstddev
    t = bincenters[ (bincenters>mincts) & (bincenters<maxcts) ]
    data =    hist[ (bincenters>mincts) & (bincenters<maxcts) ]
    p0 = [mode, inputstddev, peak]
    print '   ...initial parameter guesses = %.3f %.3f %.0f' % (p0[0],p0[1],p0[2])
    pbest = leastsq(residuals, p0, args=(data,t), full_output=1)
    p = pbest[0]
    print '   ...best fit parameters =       %.3f %.3f %.0f' % (p[0],p[1],p[2])
    offset = inputmean - p[0]
    print '   ...offset = %.3f' % offset
    xfit = numpy.linspace(mincts, maxcts, 100)
    yfit = gaussian(xfit, p)
    return hist, p[0], offset, xfit, yfit

#-----------------------------------------------------------------------

def GridFilter(img):
    global qxsize, qysize                                 # quadrant size
    gsize = 64
    indx = numpy.tile(numpy.arange(0,qxsize-gsize+1,gsize), qysize/gsize)
    indy = numpy.arange(0,qysize-gsize+1,gsize).repeat(qxsize/gsize)
    tmpimg = numpy.zeros((gsize, gsize))
    for iy in range(0, gsize):
        for ix in range(0, gsize):
            tmpdata = img[indy+iy,indx+ix]
            tmpimg[iy,ix] = numpy.median(tmpdata)
    return tmpimg

#-----------------------------------------------------------------------

def cleanir(inputfile):
    global allpixels
    global badmask, biasadjust, bias1, bias2, bias3, bias4, bins, bincenters
    global cfrac, datadir, force
    global inputmean, inputmedian, inputstddev, lsigma, hsigma
    global output, pattern, patternfile, pxsize, pysize, qxsize, qysize, quadmedian, rowfilter
    global skyfile, subtractsky
    global verbose

    print 'CLEANIR v.', version

    havedq = False           # we have DQ information
    lampsonflat = False      # this is a lamps-on flat

    if verbose:
        print '...inputfile =', inputfile
        print '...allpixels =', allpixels
        print '...badmask =', badmask
        print '...biasadjust =', biasadjust
        print '...bias1 =', bias1
        print '...bias2 =', bias2
        print '...bias3 =', bias3
        print '...bias4 =', bias4
        print '...cfrac =', cfrac
        print '...datadir =', datadir
        print '...force =', force
        print '...quadmedian =', quadmedian
        print '...output =', output
        print '...patternfile =', patternfile
        print '...row filter =', rowfilter
        print '...skyfile =', skyfile
        print '...subtractsky =', subtractsky
        print '...pxsize =', pxsize
        print '...pysize =', pysize

    if not inputfile.endswith('.fits'):
        inputfile = inputfile + '.fits'

    if (output == 'default'):
        outputfile = 'c' + os.path.basename(inputfile)
    else:
        outputfile = output
        if ( not outputfile.endswith('.fits') ):
            outputfile = outputfile + '.fits'

    if (datadir != ''):
        inputfile = datadir + '/' + inputfile

    print 'Removing pattern noise from', inputfile

    if (savepattern):
        if ( not patternfile.endswith('.fits') ):
            patternfile = patternfile + '.fits'

    if subtractsky:
        if ( not skyfile.endswith('.fits') ):
            skyfile = skyfile + '.fits'

    if not os.path.exists(inputfile):      # check whether input file exists
        print 'ERROR: ', inputfile, 'does not exist'
        sys.exit(2)

    if os.path.exists(outputfile):        # check whether output file exists
        print '...removing old', outputfile
        os.remove(outputfile)

    if savepattern:
        if os.path.exists(patternfile):  # check whether pattern file exists
            print '...removing old', patternfile
            os.remove(patternfile)

    if subtractsky:
        if not os.path.exists(skyfile):      # check whether sky file exists
            print skyfile, 'does not exist'
            sys.exit(2)

    if cfrac < 0.1:
        print 'ERROR: central fraction must be >= 0.1'
        sys.exit(2)

    if cfrac > 1:
        print 'ERROR: central fraction must be <= 1'
        sys.exit(2)

    if verbose:
        print '...reading', inputfile
    hdulist = pyfits.open(inputfile)
    if verbose:
        print '...hdulist:', hdulist.info()

    # This isn't always reliable, so just ignore non-standard FITS for now...
    #if verbose:
    #    print '...verifying FITS header...'
    #    hdulist.verify('fix')
    #else:
    #    hdulist.verify('silentfix')

    next = len(hdulist)
    if verbose:
        print '...number of extensions = ', next
    if ( next == 1 ):
        sci = 0
    elif ( next < 5 ):
        sci = 1
    else:
        sci = 2
    print '...assuming the science data are in extension', sci

    image = numpy.array(hdulist[sci].data)
    if verbose:
        print '...SCI: ', image
        print image.dtype.name

    try:
        naxis1,naxis2 = hdulist[sci].header['naxis1'], hdulist[sci].header['naxis2']
    except:
        print 'ERROR: cannot get the dimensions of extension ', sci
        pyfits.info(inputfile)
        sys.exit(2)
    print '...image dimensions = ', naxis1, 'x', naxis2

    try:
        instrument = hdulist[0].header['INSTRUME']
        if verbose:
            print '...instrument =', instrument
    except:
        print 'WARNING: cannot determine instrument'
        instrument = 'INDEF'
        allpixels = True

    try:
        nscut = hdulist[0].header['NSCUT']
        nscut = True
    except:
        nscut = False
    if verbose:
        print '...nscut =', nscut

    if instrument == 'GNIRS':
        print '...padding the top of GNIRS image...'
        pad = numpy.zeros((2,naxis1), dtype=numpy.float32) # create 2D array of padding
        image = numpy.append(image,pad,axis=0)  # append the padding array to the end
        if verbose:
            print '...new image:', image
        naxis2 = naxis2 + 2
        print '...image dimensions = ', naxis1, 'x', naxis2

    print '...pattern size =', pxsize, 'x', pysize

    qxsize = naxis1 / 2                                       # quadrant x size
    qysize = naxis2 / 2                                       # quadrant y size

    if qxsize%pxsize != 0 or qysize%pysize != 0:
        print 'ERROR: quadrant size is not a multiple of the pattern size'
        sys.exit(2)

    if pxsize > qxsize or pysize > qysize:
        print 'ERROR: pattern size is larger than the quadrant size!'
        sys.exit(2)

    try:
        skyimage = hdulist[0].header['SKYIMAGE']
        skysubtracted = True
    except:
        skysubtracted = False
    print '...sky subtracted =', skysubtracted

    #-----------------------------------------------------------------------

    if usedq and badmask == 'DQ':
        if verbose:
            print '...reading data quality extension...'
        try:
            dq = numpy.array(hdulist['DQ'].data)
            havedq = True
            if verbose:
                print '...DQ: ', dq
            if ( numpy.size(dq[dq>0]) > numpy.size(dq)/2 ):
                print 'WARNING:', numpy.size(dq[dq>0]), 'pixels are flagged as bad in the DQ plane!'
        except:
            print '...no DQ data found'
            # dq = numpy.zeros(naxis2*naxis1,int)
            # dq.resize(naxis2,naxis1)

    elif os.path.exists(badmask): # bad pixel mask specified on the command line

        if verbose:
            print '...reading bad pixel mask', badmask

        if badmask.endswith('.pl'):
            if verbose:
                print '...converting pl to fits'
            # fitsfile = inputfile.replace('.pl', '.fits')
            tmpbadmask = 'cleanir-badpixelmask.fits'

            if os.path.exists(tmpbadmask):
                os.remove(tmpbadmask)

            from pyraf import iraf

            iraf.imcopy(badmask, tmpbadmask)
            badmask = tmpbadmask

        badmaskhdu = pyfits.open(badmask)

        if verbose:
            badmaskhdu.info()

        dq = numpy.array(badmaskhdu[0].data)
        havedq = True

        if badmask.endswith('.pl'):
            os.remove(tmpbadmask)

        if verbose:
            print '...DQ: ', dq

        badmaskhdu.close()

    elif badmask == 'DQ' and not usedq:
        print '! WARNING: ignoring DQ plane!'
        havedq = False

    else:
        print '! WARNING: ', badmask, 'does not exist!'
        havedq = False

    #-----------------------------------------------------------------------

    if biasadjust:
        try:
            obstype = hdulist[0].header['OBSTYPE']
        except:
            print 'WARNING: cannot determine obstype'
            obstype = 'INDEF'
        if verbose:
            print '...obstype =', obstype

        if (obstype == 'FLAT'):
            try:
                gcalshutter = hdulist[0].header['GCALSHUT']
            except:
                print 'WARNING: cannot determine GCAL shutter status'
                if verbose:
                    print '...gcal shutter =', gcalshutter
            if (gcalshutter == 'OPEN'):
                print '! WARNING: This is a lamps-on flat, so turning off quadrant normalization.'
                lampsonflat = True

    # Bias level adjustment should probably only be done on flat-fielded data.

    #-----------------------------------------------------------------------

    if subtractsky:
        print '...reading sky', skyfile
        sky = pyfits.open(skyfile)
        print '...verifying sky...'
        if verbose:
            sky.verify('fix')
        else:
            sky.verify('silentfix')
        skyimage = numpy.array(sky[sci].data)

        if instrument == 'GNIRS':
            print '...padding the top of the GNIRS sky...'
            skyimage = numpy.append(skyimage,pad,axis=0)  # append the padding array to the end

        skysubtracted = True # This variable used when determining the area to use for pattern determination.

    # NEED ERROR CHECKING HERE! (extensions, image size, filter, instrument, etc.)

    #-----------------------------------------------------------------------

    if subtractrowmedian:
        print '...subtracting the median of each rows...'
        imagemean = numpy.mean(image)
        for iy in range(0, naxis2):
            tmprow = image[iy,:]
            if verbose:
                print '...row ', iy
            median,stddev = IterStat(tmprow)  # iterative median
            image[iy,:] -= median
        image += ( imagemean - image.mean() )  # reset image to the original mean value

        # image[iy,:] -= numpy.median(image[iy,:])  # simple row-filtering over the whole image
        # Median filter each quadrant:
        # image[iy,0:naxis1/2]      -= numpy.median(image[iy,0:naxis1/2])
        # image[iy,naxis1/2:naxis1] -= numpy.median(image[iy,naxis1/2:naxis1])

    #-----------------------------------------------------------------------
    # Set regions to be used for pattern determination:
    # +-------+
    # | 1 | 2 |
    # +---+---+
    # | 3 | 4 |
    # +---+---+

    if instrument == 'NIRI':
        camera = 'INDEF'
        decker = 'INDEF'
        try:
            fpmask = hdulist[0].header['FPMASK']
        except:
            print 'WARNING: cannot find FPMASK header keyword'
            print '   Assuming that this is imaging...'
            fpmask = 'f6-cam_G5208'

    elif instrument == 'GNIRS':
        fpmask = 'INDEF'
        try:
            camera = hdulist[0].header['CAMERA']
        except:
            print 'WARNING: cannot find CAMERA header keyword'
            camera = 'INDEF'

        try:
            decker = hdulist[0].header['DECKER']
        except:
            print 'WARNING: cannot find DECKER header keyword'
            decker = 'INDEF'

        try:
            lnrs = hdulist[0].header['LNRS']
        except:
            print 'WARNING: cannot find LNRS header keyword'
            lnrs = 'INDEF'

    else:
        fpmask = 'INDEF'
        camera = 'INDEF'
        decker = 'INDEF'
        allpixels = True

    if verbose:
        print '...fpmask = ', fpmask
        print '...camera = ', camera
        print '...decker = ', decker

    if allpixels:
        print '...using whole image for pattern determination'
        q1x1,q1x2, q1y1,q1y2 = 0,qxsize,      qysize,naxis2  # quad 1
        q2x1,q2x2, q2y1,q2y2 = qxsize,naxis1, qysize,naxis2  # quad 2
        q3x1,q3x2, q3y1,q3y2 = 0,qxsize,           0,qysize  # quad 3
        q4x1,q4x2, q4y1,q4y2 = qxsize,naxis1,      0,qysize  # quad 4
        lsigma = 3.0
        hsigma = 1.0      # set a very small upper threshold to reject stars

    elif instrument == 'NIRI' and not skysubtracted and \
             (fpmask == 'f6-2pixBl_G5214' or \
              fpmask == 'f6-4pixBl_G5215' or \
              fpmask == 'f6-6pixBl_G5216' or \
              fpmask == 'f6-2pix_G5211'):
        print '...using region above and below slit (y<=272 and y>=728) for pattern determination'
        q1x1,q1x2, q1y1,q1y2 = 0,qxsize,      728,naxis2
        q2x1,q2x2, q2y1,q2y2 = qxsize,naxis1, 728,naxis2
        q3x1,q3x2, q3y1,q3y2 = 0,qxsize,      0,272
        q4x1,q4x2, q4y1,q4y2 = qxsize,naxis1, 0,272
        lsigma = 3.0
        hsigma = 3.0

    elif instrument == 'NIRI' and not skysubtracted and \
             (fpmask == 'f6-4pix_G5212'  or \
              fpmask == 'f6-6pix_G5213'  or \
              fpmask == 'f32-6pix_G5229' or \
              fpmask == 'f32-9pix_G5230'):
        print '...using whole image for pattern determination'
        print 'WARNING: Sky lines may be altered by pattern removal!'
        q1x1,q1x2, q1y1,q1y2 = 0,qxsize,      qysize,naxis2
        q2x1,q2x2, q2y1,q2y2 = qxsize,naxis1, qysize,naxis2
        q3x1,q3x2, q3y1,q3y2 = 0,qxsize,           0,qysize
        q4x1,q4x2, q4y1,q4y2 = qxsize,naxis1,      0,qysize
        lsigma = 3.0
        hsigma = 3.0

    elif instrument == 'GNIRS' and not skysubtracted and \
             'Short' in camera and not 'XD' in decker:  # GNIRS short-camera long-slit
        print '...using x<=160 and x>=864 for pattern determination'
        q1x1,q1x2, q1y1,q1y2 = 0,160,         qysize,naxis2
        q2x1,q2x2, q2y1,q2y2 = 864,naxis1,    qysize,naxis2
        q3x1,q3x2, q3y1,q3y2 = 0,160,         0,qysize
        q4x1,q4x2, q4y1,q4y2 = 864,naxis1,    0,qysize
        lsigma = 3.0
        hsigma = 3.0

    else:
        print '...using whole image for pattern determination'
        q1x1,q1x2, q1y1,q1y2 = 0,qxsize,      qysize,naxis2  # quad 1
        q2x1,q2x2, q2y1,q2y2 = qxsize,naxis1, qysize,naxis2  # quad 2
        q3x1,q3x2, q3y1,q3y2 = 0,qxsize,           0,qysize  # quad 3
        q4x1,q4x2, q4y1,q4y2 = qxsize,naxis1,      0,qysize  # quad 4
        lsigma = 3.0
        hsigma = 1.0      # set a very small upper threshold to reject stars

    patternin  = image.copy()
    patternin1 = patternin[q1y1:q1y2, q1x1:q1x2]
    patternin2 = patternin[q2y1:q2y2, q2x1:q2x2]
    patternin3 = patternin[q3y1:q3y2, q3x1:q3x2]
    patternin4 = patternin[q4y1:q4y2, q4x1:q4x2]

    #-------------------------------------------------------------------
    # Subtract sky frame
    if subtractsky:
        print '...subtracting sky...'
        patternin1 -= skyimage[q1y1:q1y2, q1x1:q1x2]
        patternin2 -= skyimage[q2y1:q2y2, q2x1:q2x2]
        patternin3 -= skyimage[q3y1:q3y2, q3x1:q3x2]
        patternin4 -= skyimage[q4y1:q4y2, q4x1:q4x2]

    #-------------------------------------------------------------------
    # Flag pixels with bad DQ
    if havedq:
        print '...flagging bad pixels...'
        dq1 = dq[q1y1:q1y2, q1x1:q1x2]
        dq2 = dq[q2y1:q2y2, q2x1:q2x2]
        dq3 = dq[q3y1:q3y2, q3x1:q3x2]
        dq4 = dq[q4y1:q4y2, q4x1:q4x2]
        patternin1[dq1==1] = bad
        patternin2[dq2==1] = bad
        patternin3[dq3==1] = bad
        patternin4[dq4==1] = bad

    #-------------------------------------------------------------------
    # Calculate means and medians for reference:

    inputmean = numpy.mean(image)
    print '...mean of input image = %.3f' % inputmean

    if (biasadjust or graph>0):
        #inputmedian, inputstddev = IterStat(image)  # sigma-clipped
        allpatternin = numpy.concatenate(seq=(patternin1,patternin2,patternin3,patternin4))
        allpatternin = allpatternin[allpatternin!=bad] # filter out bad values
        inputmedian, inputstddev = IterStat(allpatternin)
        print '...sigma-clipped median = %.3f' % inputmedian
        print '...sigma-clipped stddev = %.3f' % inputstddev

    #-------------------------------------------------------------------
    # calculate and subtract pattern:

    quads = image.copy()
    quad1 = quads[qysize:naxis2,      0:qxsize]
    quad2 = quads[qysize:naxis2, qxsize:naxis1]
    quad3 = quads[    0:qysize,       0:qxsize]
    quad4 = quads[    0:qysize,  qxsize:naxis1]

    print '...upper left quadrant:'
    clean1, pattern1 = CleanQuad(quad1,patternin1)
    print '...upper right quadrant:'
    clean2, pattern2 = CleanQuad(quad2,patternin2)
    print '...lower left quadrant:'
    clean3, pattern3 = CleanQuad(quad3,patternin3)
    print '...lower right quadrant:'
    clean4, pattern4 = CleanQuad(quad4,patternin4)

    if verbose:
        print '...reassembling new image...'
    newimage = image.copy()
    newimage[qysize:naxis2,      0:qxsize] = clean1
    newimage[qysize:naxis2, qxsize:naxis1] = clean2
    newimage[     0:qysize,      0:qxsize] = clean3
    newimage[     0:qysize, qxsize:naxis1] = clean4

    if verbose:
        print '...updating header...'
    timestamp = datetime.datetime.now().strftime("%Y.%m.%d %H:%M:%S")
    if verbose:
        print '...time stamp =', timestamp
    hdulist[0].header.add_history('Cleaned with cleanir.py ' + timestamp)

    #-----------------------------------------------------------------------
    # Use the cleaned image from here on out

    patternin  = newimage.copy()

    if subtractsky:
        print '...subtracting sky...'
        patternin -= skyimage

    patternin1 = patternin[q1y1:q1y2, q1x1:q1x2]
    patternin2 = patternin[q2y1:q2y2, q2x1:q2x2]
    patternin3 = patternin[q3y1:q3y2, q3x1:q3x2]
    patternin4 = patternin[q4y1:q4y2, q4x1:q4x2]

    #-----------------------------------------------------------------------
    # GNIRS 8-pixel row filtering

    # Go through each row of each quadrant and generate an 8-pixel wide kernel,
    # subtract it, and then add back the previously measured quadrant mean.

    if rowfilter:
        print '...filtering rows...'

        print '...upper left quadrant:'
        clean1 = ApplyRowFilter(newimage[qysize:naxis2,      0:qxsize], patternin1)
        print '...upper right quadrant:'
        clean2 = ApplyRowFilter(newimage[qysize:naxis2, qxsize:naxis1], patternin2)
        print '...lower left quadrant:'
        clean3 = ApplyRowFilter(newimage[     0:qysize,      0:qxsize], patternin3)
        print '...lower right quadrant:'
        clean4 = ApplyRowFilter(newimage[     0:qysize, qxsize:naxis1], patternin4)

        if verbose:
            print '...reassembling image...'

        newimage[qysize:naxis2,      0:qxsize] = clean1
        newimage[qysize:naxis2, qxsize:naxis1] = clean2
        newimage[     0:qysize,      0:qxsize] = clean3
        newimage[     0:qysize, qxsize:naxis1] = clean4

        # Use the cleaned image from here on out
        patternin  = newimage.copy()
        if subtractsky:
            print '...subtracting sky...'
            patternin -= skyimage
        patternin1 = patternin[q1y1:q1y2, q1x1:q1x2]
        patternin2 = patternin[q2y1:q2y2, q2x1:q2x2]
        patternin3 = patternin[q3y1:q3y2, q3x1:q3x2]
        patternin4 = patternin[q4y1:q4y2, q4x1:q4x2]

    #-------------------------------------------------------------------
    # Normalize each quadrant:

    if (biasadjust and not lampsonflat):
        print '...normalizing the bias level of each quadrant...'
        # And apply the measured offset to the pattern output

        if havedq:                                # Flag pixels with bad DQ
            print '...flagging bad pixels...'
            dq1 = dq[qysize:(1+cfrac)*qysize, (1-cfrac)*qxsize:qxsize]
            dq2 = dq[qysize:(1+cfrac)*qysize, qxsize:(1+cfrac)*qxsize]
            dq3 = dq[(1-cfrac)*qysize:qysize, (1-cfrac)*qxsize:qxsize]
            dq4 = dq[(1-cfrac)*qysize:qysize, qxsize:(1+cfrac)*qxsize]
            patternin1[dq1==1] = bad
            patternin2[dq2==1] = bad
            patternin3[dq3==1] = bad
            patternin4[dq4==1] = bad

        binmin = inputmedian - 5. * inputstddev
        binmax = inputmedian + 5. * inputstddev
        binwidth = 1.0
        if (binmax - binmin) / binwidth < 50: # if too few bins the least-squares minimization will fail
            binwidth = (binmax - binmin) / 50.

        if instrument == 'GNIRS': # GNIRS pixels are integers when LNRS=1, so use integer bins.
            if lnrs == 1:         # This definitely makes the graph look better but it
                binwidth = 1      # doesn't seem to help the quality of the normalization.

        if verbose:
            print '...median =', inputmedian,'  stddev =', inputstddev
        bins = numpy.arange( binmin, binmax, binwidth )
        bincenters = bins[1:bins.size] - binwidth/2.
        print '...binning into', bins.size, 'bins from', binmin, 'to', binmax, '(bin width =', binwidth,')'

        if applygridfilter:
            print '...applying grid filter to each quadrant...'
            patternin1 = GridFilter(patternin1)
            patternin2 = GridFilter(patternin2)
            patternin3 = GridFilter(patternin3)
            patternin4 = GridFilter(patternin4)

        if quadmedian:
            fit = False
            print '...Using median for quadrant normalization.'
            print '...upper left quadrant:'
            hist1,center1,offset1 = NormQuadMedian(patternin1)
            print '...upper right quadrant:'
            hist2,center2,offset2 = NormQuadMedian(patternin2)
            print '...lower left quadrant:'
            hist3,center3,offset3 = NormQuadMedian(patternin3)
            print '...lower right quadrant:'
            hist4,center4,offset4 = NormQuadMedian(patternin4)
        else:
            fit = True
            print '...upper left quadrant:'
            hist1,center1,offset1,xfit1,yfit1 = NormQuadGauss(patternin1)
            print '...upper right quadrant:'
            hist2,center2,offset2,xfit2,yfit2 = NormQuadGauss(patternin2)
            print '...lower left quadrant:'
            hist3,center3,offset3,xfit3,yfit3 = NormQuadGauss(patternin3)
            print '...lower right quadrant:'
            hist4,center4,offset4,xfit4,yfit4 = NormQuadGauss(patternin4)

        newimage[qysize:naxis2,      0:qxsize] += offset1
        newimage[qysize:naxis2, qxsize:naxis1] += offset2
        newimage[0:qysize,           0:qxsize] += offset3
        newimage[0:qysize,      qxsize:naxis1] += offset4
        pattern1 -= offset1
        pattern2 -= offset2
        pattern3 -= offset3
        pattern4 -= offset4

        print '...checking quality of bias normalization...'
        newmedian, newstddev = IterStat(newimage)

        if ( inputstddev - newstddev > 0.001 ):
            print '   ...sigma-clipped stddev has decreased from %.3f to %.3f' % (inputstddev, newstddev)
            offset = inputmean - numpy.mean(newimage)
            print '...adjusting whole image by %.3f to match input image...' % offset
            newimage += offset
        else:
            print '   ...sigma-clipped stddev has not significantly improved: %.3f -> %.3f' % (inputstddev, newstddev)
            print '   ...undoing quadrant bias offsets...'
            outimage = newimage
            image[qysize:naxis2,      0:qxsize] -= offset1
            image[qysize:naxis2, qxsize:naxis1] -= offset2
            image[0:qysize,           0:qxsize] -= offset3
            image[0:qysize,      qxsize:naxis1] -= offset4
            pattern1 += offset1
            pattern2 += offset2
            pattern3 += offset3
            pattern4 += offset4

        #-------------------------------------------------------------------

        if (graph>1): # 2x2 grid of pixel distributions, fits & estimated sky values
            print ('...graphing pixel distributions in each quadrant...')

            xlimits = numpy.array([binmin, binmax])

            plot = pyplot.subplot(2,2,1)
            pyplot.plot(bincenters, hist1, linestyle='', marker='.')
            if fit:
                pyplot.plot(xfit1, yfit1, linestyle='-', color='red', linewidth=2)
            pyplot.xlim(xlimits)
            pyplot.axvline(x=center1, ls='--', color='green')
            pyplot.text(0.05, 0.85, 'mean = %.2f'  % center1, horizontalalignment='left', transform=plot.transAxes)
            pyplot.text(0.95, 0.85, 'delta = %.2f' % offset1, horizontalalignment='right', transform=plot.transAxes)
            pyplot.title('Quadrant 1')

            plot = pyplot.subplot(2,2,2)
            pyplot.plot(bincenters, hist2, linestyle='', marker='.')
            if fit:
                pyplot.plot(xfit2, yfit2, linestyle='-', color='red', linewidth=2)
            pyplot.xlim(xlimits)
            pyplot.axvline(x=center2, ls='--', color='green')
            pyplot.text(0.05, 0.85, 'mean = %.2f' % center2, horizontalalignment='left', transform=plot.transAxes)
            pyplot.text(0.95, 0.85, 'delta = %.2f' % offset2, horizontalalignment='right', transform=plot.transAxes)
            pyplot.title('Quadrant 2')

            plot = pyplot.subplot(2,2,3)
            pyplot.plot(bincenters, hist3, linestyle='', marker='.')
            if fit:
                pyplot.plot(xfit3, yfit3, linestyle='-', color='red', linewidth=2)
            pyplot.xlim(xlimits)
            pyplot.axvline(x=center3, ls='--', color='green')
            pyplot.text(0.05, 0.85, 'mean = %.2f' % center3, horizontalalignment='left', transform=plot.transAxes)
            pyplot.text(0.95, 0.85, 'delta = %.2f' % offset3, horizontalalignment='right', transform=plot.transAxes)
            pyplot.title('Quadrant 3')

            plot = pyplot.subplot(2,2,4)
            pyplot.plot(bincenters, hist4, linestyle='', marker='.')
            if fit:
                pyplot.plot(xfit4, yfit4, linestyle='-', color='red', linewidth=2)
            pyplot.xlim(xlimits)
            pyplot.axvline(x=center4, ls='--', color='green')
            pyplot.text(0.05, 0.85, 'mean = %.2f' % center4, horizontalalignment='left', transform=plot.transAxes)
            pyplot.text(0.95, 0.85, 'delta = %.2f' % offset4, horizontalalignment='right', transform=plot.transAxes)
            pyplot.title('Quadrant 4')

            pyplot.subplots_adjust(left=0.05, bottom=0.05, right=0.95,
                                   top=0.95, wspace=0.2, hspace=0.2)

            # top label = inputfile
            pyplot.show()

    #-------------------------------------------------------------------
    # Apply manual bias correction if supplied:
    if (bias1 != 0.0 or bias2 != 0.0 or bias3 != 0.0 or bias4 != 0.0):
        print '...applying user-supplied bias offset...'
        newimage[qysize:naxis2,      0:qxsize] += bias1
        newimage[qysize:naxis2, qxsize:naxis1] += bias2
        newimage[0:qysize,           0:qxsize] += bias3
        newimage[0:qysize,      qxsize:naxis1] += bias4
        pattern1 -= bias1
        pattern2 -= bias2
        pattern3 -= bias3
        pattern4 -= bias4

    if verbose:
        print '...mean of input image = %.3f' % inputmean
        print '...mean of output image = %.3f' % numpy.mean(newimage)
        print '...median of output image = %.3f' % numpy.median(newimage)

    #-------------------------------------------------------------------
    # Write cleaned output image

    if instrument == 'GNIRS':
        print '...removing GNIRS padding...'
        # remove 2-pixel padding on top of image:
        # syntax: delete(array, [rows to delete], axis=0)
        newimage = numpy.delete(newimage, [naxis2-1,naxis2-2], axis=0)

    print '...writing', outputfile
    hdulist[sci].data = newimage
    hdulist.writeto(outputfile, output_verify='warn')

    #-------------------------------------------------------------------
    # Write pattern image

    if (savepattern):
        print '...assembling and writing pattern image...'
        # create blank pattern array:
        fullpattern = numpy.zeros(naxis2*naxis1).reshape(naxis2,naxis1)
        # assemble the quadrants into a full pattern image:
        fullpattern[qysize:naxis2,      0:qxsize] = pattern1
        fullpattern[qysize:naxis2, qxsize:naxis1] = pattern2
        fullpattern[     0:qysize,      0:qxsize] = pattern3
        fullpattern[     0:qysize, qxsize:naxis1] = pattern4
        # normalize to zero:
        fullpattern -= fullpattern.mean()
        print '...writing', patternfile
        hdu = pyfits.PrimaryHDU(fullpattern)
        hdu.writeto(patternfile)

    #-------------------------------------------------------------------
    # Close file

    hdulist.close()
    print ' '

#-----------------------------------------------------------------------

if __name__ == '__main__':
    main()

#-----------------------------------------------------------------------


# -*- coding: utf-8 -*-
"""
Created on Sat Jul 18 21:29:33 2020

Script for showing the effectiveness of the WebRTC VAD speech-removal pipeline
with a variety of files

@author: benpg
"""

from webrtc_vad import VAD
from webrtc_utils import plt_speech, speech_stats
import os
from os.path import join
import numpy as np
import matplotlib.pyplot as plt
import librosa
from VAD_pipeline import VAD_pipeline

def webrtc_tester(audiodir, labelspath, aggressiveness, norm_LU=False, plots_on=False):
    '''
    Parameters
    ----------
    audiodir : str
        Path to folder containing audio files to analyse.
    labelspath : str
        Path to folder containing speech label txt files with identical names to the .wav files.
    aggressiveness : int
        Aggressiveness with which noise is removed from file (0-3).
    norm_LU : int or False, optional
        Takes a value that must be a negative integer as the target LU at which to loudness normalise the file (dB). The default is False, i.e. no normalisation.
    plots_on : bool or int, optional
        Decides if/how many plots of true/predicted speech to plot for the files given. False: no plots (default); True: plot all files; integer n: plot the first n files

    Returns
    -------
    true_speech: list of floats
        List with each element the number of seconds of correctly identified speech in each file.
    speech_total_length: list of floats
        List with each element the total number of seconds of speech in each file, given by the speech labels
    true_noise: list of floats
        List with each element the number of seconds of correctly identified noise in each file.
    noise_total_length: list of floats
        List with each element the total number of seconds of noise in each file.

    '''
    pltcount = 0
    true_speech, speech_total_length, true_noise, noise_total_length = [],[],[],[]
    if plots_on==True: plots_on = np.inf     # plot all files if no limit set
    if norm_LU:
        normdir = join(audiodir,'norm')
        if not os.path.exists(normdir): os.mkdir(normdir)
        for file in os.listdir(audiodir):
            if file.endswith('.wav'):
                VAD_pipeline(audiodir, file, normdir, agg=aggressiveness, sr=8000, LU=norm_LU, norm_only=True)
        audiodir = normdir

    for file in os.listdir(audiodir):
        if file.endswith('.wav'):
            audio, sr = librosa.load(join(audiodir,file), sr=2000)
            speech_pred_labels = VAD(join(audiodir,file), aggressiveness, chunks=False)
            filename = file.rstrip('.wav')
            try:
                speech_labels = np.loadtxt(join(labelspath,filename)+'.txt', usecols=(0,1))
            except:
                # If label file not found, print warning and process with no speech labels
                print('\nFile: %s\nLabel file not found, assuming no speech.\n' % file)
                speech_labels = np.array([[0,0]])
            file_length = audio.shape[0]/sr

            # reshape to 2D for files with only one pair of labels
            if speech_labels.shape == (2,): speech_labels = np.array([speech_labels])

            # Plot the true and predicted speech
            if pltcount < plots_on:
                plt_speech(speech_pred_labels, speech_labels, file_length, filename, audio, sr)
                pltcount += 1

            # Get and print test information
            print('\nFile:', file)
            stats = speech_stats(speech_pred_labels, speech_labels, file_length)
            # Keep running total of length of speech and length predicted correctly
            true_speech.append(stats[0])
            speech_total_length.append(stats[1])
            true_noise.append(stats[2])
            noise_total_length.append(stats[3])
    return(true_speech, speech_total_length, true_noise, noise_total_length)

plt.close('all')

aggressiveness = 0
plots_on = True  # Plot true/predicted speech. False for none, True for all, or an int number
norm_LU = -35     # False or a number

path = r'D:\Big Files\Humbug OneDrive\Experiments\IHI_small_sample'
labelspath = join(path, 'speech_labels')

audiodir = path

true_speech, speech_total_length, true_noise, noise_total_length = \
    webrtc_tester(audiodir, labelspath, aggressiveness, norm_LU, plots_on)

tpr_overall = np.sum(true_speech)/np.sum(speech_total_length)
tnr_overall = np.sum(true_noise)/np.sum(noise_total_length)
print('\nOVERALL:\n%.4f percent of speech detected (%.2f out of %.2f seconds)' \
      % (100*tpr_overall, np.sum(true_speech), np.sum(speech_total_length)))
print('%.4f percent of noise preserved (%.2f out of %.2f seconds)' \
  % (100*tnr_overall, np.sum(true_noise), np.sum(noise_total_length)))


# This is meant to be run in python 3.4 or higher

# This module is a basic, raw writer of the DAQ data to disk.

from __future__ import print_function

import numpy as n
import time, calendar

#@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
#@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

LZ_EPOCH_STRING = '20150101T0000'
#LZ_EPOCH_DATETIME_DT = datetime.strptime(LZ_EPOCH_STRING, '%Y%m%dT%H%M')
LZ_EPOCH_DATETIME = n.uint64(calendar.timegm(time.strptime(LZ_EPOCH_STRING, '%Y%m%dT%H%M')) * 1e9)

NS_PER_S_UINT64 = n.uint64(1e9)

UINT64_VERSION_OF_32 = n.uint64(32)

#@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
#@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
#@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

class processStream:
    def __init__(self, digi_time_at_start_ns, outfile_base='data/lzrd_sue_raw', buffer_size_bytes=1048576, buffer_write_timeout_s = 60, max_file_size_bytes = 1073741824, verbose=0):
        """

        """
        # store verbose level. 0 means no output. 1 mean minimum output. And so forth
        self.verbose = verbose
        # start time of the acquisition. Technically, it's the time this object got instantiated
        # minus whatever offset needed
        self.abs_starttime_ns = n.uint64(time.time()*1e9) - LZ_EPOCH_DATETIME - n.uint64(digi_time_at_start_ns)

        # create a base of a filename to write to
        gmtime = time.gmtime((self.abs_starttime_ns+LZ_EPOCH_DATETIME)/1e9)
        timestamp = time.strftime('%Y%m%dT%H%M%S', gmtime)
        self.outfile_base = "%s_%s" %(outfile_base, timestamp)

        self.outfile_handle = None
        self.max_file_size_bytes = max_file_size_bytes
        self.current_file_ind = 0

        # initialize a buffer for storing the events before writing them.
        self.write_buffer = bytearray()
        ## A simple byte object that we continously append to and delete may be too slow. So, try a fixed bytearray
        #self.write_buffer = bytearray(int(buffer_size_bytes*2))
        #self.write_buffer_size = 0
        # turns out that, with what I'm trying right now, the fixed bytearray isn't any faster
        # keep track of when we last wrote the data to a buffer
        self.last_buffer_write = time.time()
        # file write variables
        self.buffer_size_bytes = buffer_size_bytes
        self.buffer_write_timeout_s = buffer_write_timeout_s

        # get the info string for this data set
        self.acq_string = self.construct_acq_string()

    def __del__(self):
        """
        Event destructor. It writes the remaining frames to file.
        ____________________________________________________________________
        Inputs:

        Optional:

        ____________________________________________________________
        Returns:

        """
        if self.verbose > 1:
            print("processStream instance destructor started")

        self.write_final()
        if self.verbose > 1:
            print("processStream instance destructor executed after final file write")

    def run_subprocess(self, lock, conn):
        """
        function executed from a separate thread
        ____________________________________________________________________
        Inputs:
            lock
                The lock object needed for safe threading

            conn
                The pipe object used for communication

        Optional:

        ____________________________________________________________
        Returns:

        """
        while True:
            # the conn.recv call will cause this thread to stop and wait for input
            break_condition, frame_bytes, deadtime = conn.recv()
            if break_condition:
                # exit loop and end function if told to do
                break
            lockstatus = lock.acquire()
            try:
                # let the parent process know that the processor is working
                conn.send_bytes(b'1')
                self.write_ready(frame_bytes, deadtime=deadtime)
            finally:
                lock.release()
        # initiate final write
        self.write_final()

    def new_file(self):
        """
        Open a new file
        ____________________________________________________________________
        Inputs:

        Optional:

        ____________________________________________________________
        Returns:

        """
        filename = "%s_%05d.dat" %(self.outfile_base, self.current_file_ind)
        self.outfile_handle = open(filename, 'bw')
        self.current_file_ind += 1  # increment file index
        # write the endian check string
        self.outfile_handle.write(n.array([0x01020304]).astype('<u4').tostring())
        # write out settings
        self.outfile_handle.write(n.array([len(self.acq_string)]).astype('<u4').tostring())
        self.outfile_handle.write(self.acq_string)
        if self.verbose > 3:
            print("Opened file %s of initial size %d" %(filename, len(setting_string) + 8))

    def construct_acq_string(self):
        """
        Constructs an event header string
        ____________________________________________________________________
        Inputs:

        Optional:

        ____________________________________________________________
        Returns:
            A binary setting string and the size of the settings string.
        """
        # Store some info.
        # event timestamp
        timestamp_highbits = n.uint32(self.abs_starttime_ns >> UINT64_VERSION_OF_32)    # timestamp_highbits
        timestamp_lowbits = n.uint32(self.abs_starttime_ns)
        acq_header_string = b"acqStartTime_ns_high:"
        acq_header_string +=timestamp_highbits.tostring()
        acq_header_string += b"acqStartTime_ns_low:"
        acq_header_string += timestamp_lowbits.tostring()
        # attach the length of the entry
        acq_header_string = n.array([len(acq_header_string)]).astype('<u4').tostring() + acq_header_string
        return acq_header_string

    def write_ready(self, frame_bytes, deadtime=0):
        """
        Add a new frame to the write buffer.
        ____________________________________________________________________
        Inputs:
            frame_bytes
                The byte string frame

        Optional:
            deadtime = 0
                The dead time to prepend before each frame.
        ____________________________________________________________
        Returns:

        """
        # open file for writing if None open
        if self.outfile_handle is None: self.new_file()
        # store the dead time along with the size of the current frame
        deadtime = n.uint64(deadtime)
        deadtime_highbits = n.uint32(deadtime >> UINT64_VERSION_OF_32)
        deadtime_lowbits = n.uint32(deadtime)
        deadtime_string = b"deadtime_ns_high:"
        deadtime_string +=deadtime_highbits.tostring()
        deadtime_string += b"deadtime_ns_low:"
        deadtime_string += deadtime_lowbits.tostring()
        self.write_buffer.extend(n.array([len(deadtime_string) + len(frame_bytes)]).astype('<u4').tostring())
        self.write_buffer.extend(deadtime_string)
        # append data
        self.write_buffer.extend(frame_bytes)
        # write the event if needed
        write_buffer_len = len(self.write_buffer)
        if self.verbose > 5: print("writing to buffer of total size %d bytes" %(write_buffer_len))
        current_time = time.time()  # current time for timeout purposes
        if (write_buffer_len > self.buffer_size_bytes) or \
            ((current_time - self.last_buffer_write) > self.buffer_write_timeout_s):
            # If the buffer size limit or timeout are reached, the buffer
            # will be written.
            print("writing file")
            self.outfile_handle.write(self.write_buffer)
            if self.verbose > 3:
                print("Dumping buffer of size %d to file after %.2f seconds" %(write_buffer_len, (current_time - self.last_buffer_write)))
            # clear the buffer
            self.write_buffer.clear()
            self.last_buffer_write = current_time
        # check if file should be closed
        if self.outfile_handle.tell() >= self.max_file_size_bytes:
            # reached maximum file size
            self.outfile_handle.close() # close file
            self.outfile_handle = None  # clear variable

    def write_final(self):
        """

        ____________________________________________________________________
        Inputs:

        Optional:

        ____________________________________________________________
        Returns:

        """
        if self.verbose > 2:
            print("Entered write_final with %d PODs left to write" %(len(self.waveforms)))
        write_buffer_len = len(self.write_buffer)
        if write_buffer_len > 0:
            # open file for writing if None open
            if self.outfile_handle is None: self.new_file()
            # write the event if needed
            write_buffer_len = len(self.write_buffer)
            if self.verbose > 5: print("writing to buffer of total size %d bytes" %(write_buffer_len))
            # write the buffer
            self.outfile_handle.write(self.write_buffer)
            # clear the buffer
            self.write_buffer.clear()
            current_time = time.time()  # current time
            self.last_buffer_write = current_time
            # close the file
            self.outfile_handle.close() # close file
            self.outfile_handle = None  # clear variable


import cv2
import numpy as np
from os import listdir
from os.path import isfile, join
#import learning_sample

''' 샘플링한 얼굴 사진이 들어있는 경로 '''
data_path = 'faces/'

''' 샘플 파일만 리스트로 가져오기 '''
sample_files = [f for f in listdir(data_path) if isfile(join(data_path,f))]
Training_Data, Labels = [], []

''' Training_Data에 샘플 사진 넣기'''
for i, files in enumerate(sample_files):
    image_path = data_path + sample_files[i]
    images = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if images is None:
        continue
    ''' uint8 형식으로 배열 저장 '''
    Training_Data.append(np.asarray(images, dtype=np.uint8))
    Labels.append(np.asarray(i, dtype=np.int32))

if len(Labels) == 0:
    print("There is no data to train.")
    exit()
''' LBPHFaceRecongizer 학습 '''
recognition_lbph = cv2.face.LBPHFaceRecognizer_create()
recognition_lbph.train(np.asarray(Training_Data), np.asarray(Labels))
print("Model Training Complete.")

''' 모자이크 기능 시작 '''
Conf_threshold = 0.4
NMS_threshold = 0.4
COLORS = [(0, 255, 0), (0, 0, 255), (255, 0, 0),
          (255, 255, 0), (255, 0, 255), (0, 255, 255)]

''' YOLO 학습된 객체 불러오기 '''
class_name = []
with open('classes-custom.txt', 'r') as f:
    class_name = [cname.strip() for cname in f.readlines()]

net = cv2.dnn.readNet('yolov4-tiny-custom.weights', 'yolov4-tiny-custom.cfg')
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA_FP16)

detection_yolo = cv2.dnn_DetectionModel(net)
detection_yolo.setInputParams(size=(416, 416), scale=1 / 255, swapRB=True)

def face_detector(img, size = 0.5):
    '''
    YOLO로 검출된 얼굴을 LBPH에 입력으로 쓰고 본인인지 확인한다.
    본인이 아닌 걸로 결과가 나오면 모자이크 처리한다.
    '''
    classes, scores, boxes = detection_yolo.detect(img, Conf_threshold, NMS_threshold)

    if len(boxes) == 0:
        cv2.imshow('mosaic_cam', img)
        return

    for (classid, score, box) in zip(classes, scores, boxes):
        color = COLORS[int(classid) % len(COLORS)]
        cv2.rectangle(img, box, color, 1)

        roi = img[box[1]:box[1]+box[3],box[0]:box[0]+box[2]]

        try:
            roi = cv2.resize(roi, (200, 200))
        except:
            break

        ''' 검출된 사진을 흑백으로 변환 '''
        roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)

        ''' 위에서 학습한 모델로 예측 시도 '''
        result = recognition_lbph.predict(roi)

        ''' result[1]은 신뢰도이고 0에 가까울수록 자신과 같다는 뜻이다. '''
        # 500?
        if result[1] < 500:
            confidence = int(100*(1-(result[1])/300))
            display_string = str(confidence)+'%'

        cv2.putText(img,display_string,(box[0],box[1]), cv2.FONT_HERSHEY_COMPLEX,1,(250,120,255),2)

        ''' 80%이하면 다른 사람으로 간주'''
        if confidence >= 77:
            cv2.putText(img, "It's you!", (box[0], box[1]+box[3]+19), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 255, 0), 2)

        else:
            roi = img[box[1]:box[1]+box[3], box[0]:box[0]+box[1]]

            try:
                roi = cv2.resize(roi,(box[2]//15,box[3]//15))
                roi = cv2.resize(roi,(box[2],box[3]), interpolation=cv2.INTER_AREA)
            except:
                break

            img[box[1]:box[1]+box[3],box[0]:box[0]+box[2]] = roi
            cv2.putText(img, "who?", (box[0], box[1]+box[3]+19), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)

    cv2.imshow('mosaic_cam',img)

''' 카메라 열기 '''
cap = cv2.VideoCapture(0)

while True:
    ''' 카메라로 부터 사진 한장 읽기 '''
    ret, frame = cap.read()
    if ret ==False:
        break
    ''' 얼굴 검출 시도 '''
    face_detector(frame)
    if cv2.waitKey(1) & 0xFF==ord('q'):
        break

cap.release()
cv2.destroyAllWindows()


# Licensed with the 3-clause BSD license.  See LICENSE for details.
from itertools import cycle
from logging import Logger

import numpy as np

import sqlalchemy as sa
from sqlalchemy.orm import Session
from sqlalchemy.orm.exc import NoResultFound

from astropy.coordinates import Angle
import astropy.units as u
from astropy.time import Time
from sbpy.data import Ephem

from . import util, schema, ephem
from .schema import Obj, Eph, Found, Obs, GenericObs
from .util import RADec, Point
from .logging import ProgressTriangle
from .exceptions import (
    BadObjectID,
    NoEphemerisError,
    SourceNotFoundError,
    AddFoundFailure,
    MultipleSourceError
)

class SBDB:
    """Database object for SBSearch.

    Parameters
    ----------
    url_or_session : string or sqlalchemy Session
        The sqlalchemy-formatted database URL or a sqlalchmey session
        to use.

    *args
        `sqlalchemy.create_engine` arguments.

    """

    DB_NAMES = ['obj', 'eph', 'obs', 'generic_obs', 'found']

    def __init__(self, url_or_session, *args):
        if isinstance(url_or_session, Session):
            self.session = url_or_session
            self.engine = self.session.get_bind()
            self.sessionmaker = None
        else:
            self.engine = sa.create_engine(url_or_session, *args)
            self.sessionmaker = sa.orm.sessionmaker(bind=self.engine)
            self.session = self.sessionmaker()

    def __del__(self):
        self.close()

    def close(self):
        self.session.close()

    @classmethod
    def create_test_db(cls):
        """Create a test database.

        Requires database named "sbsearch_test".

        Scheme for handling temporary databases from:
        http://alextechrants.blogspot.com/2014/01/unit-testing-sqlalchemy-apps-part-2.html

        """
        from .test.skytiles import sky_tiles, N_tiles

        url = "postgresql:///sbsearch_test"
        db = SBDB(url)

        # remove previous temp work
        metadata = sa.MetaData(bind=db.engine)
        metadata.reflect()
        for table in metadata.tables.keys():
            if table == 'spatial_ref_sys':
                continue
            db.session.execute('DROP TABLE IF EXISTS {} CASCADE'
                               .format(table))

        # create test database
        db.verify_database(Logger('test'))
        db.add_object('C/1995 O1')
        db.add_object('2P')

        # 30 s exposures
        exptime = 30 / 86400
        for i in range(N_tiles**2):
            obs = GenericObs(
                obsid=i,
                jd_start=2458119.5 + exptime * i,
                jd_stop=2458119.5 + exptime * (i + 1),
                fov=sky_tiles[i],
                filter='r',
                exposure=exptime,
                seeing=1.5,
                airmass=1.3,
                maglimit=25)
            db.session.add(obs)

        db.add_ephemeris(2, '500', 2458119.5, 2458121.5, step='1d',
                         source='jpl', cache=True)
        db.add_found_by_id(2, [1, 2, 3], '500', cache=True)
        return db

    def verify_database(self, logger, names=[]):
        """Verify SBSearch tables.

        Parameters
        ----------
        logger : `~logging.Logger`
            Log messages to this logger.

        names : list, optional
            Additional database names to consider.

        """

        metadata = sa.MetaData()
        metadata.reflect(self.engine)

        missing = False
        for name in self.DB_NAMES + names:
            if name not in metadata.tables.keys():
                missing = True
                logger.error('{} is missing from database'.format(name))

        if missing:
            schema.create(self.engine)
            logger.info('Created database tables.')

    def add_ephemeris(self, objid, location, start, stop, step=None,
                      source='jpl', cache=False):
        """Add ephemeris data to databse.

        Parameters
        ----------
        objid : int
            Database object ID.

        location : string
            Observer location.

        start, stop : float, string, `~astropy.time.Time`
            Observation date range (inclusive).

        step : string or `~astropy.units.Quantity`, optional
            Ephemeris step size or ``None`` to use an adaptable size
            based on heliocentric distance.

        source : string
            Source to use: 'mpc' or 'jpl'.

        cache : bool, optional
            Use cached ephemerides; primarily for testing.

        Returns
        -------
        count : int
            Number of inserted rows.

        """

        desg = self.resolve_object(objid)[1]
        today = Time.now().iso[:10]
        count = 0

        t = util.epochs_to_time((start, stop))
        epochs = {
            'start': t[0].iso,
            'stop': t[1].iso,
            'step': step
        }
        eph = ephem.generate(desg, location, epochs, source=source,
                             cache=cache)

        jd = util.epochs_to_jd(eph['Date'].value)
        coords = RADec(eph['RA'], eph['Dec'])
        vmag = util.vmag_from_eph(eph)

        for i in range(len(eph)):
            # save ephemeris segment, which is used for
            # searching, yes they overlap
            p0 = max(i - 1, 0)
            p1 = min(len(eph) - 1, i + 1)

            segment = 'SRID=40001;LINESTRING({} {}, {} {})'.format(
                eph['RA'][p0].to('deg').value,
                eph['Dec'][p0].to('deg').value,
                eph['RA'][p1].to('deg').value,
                eph['Dec'][p1].to('deg').value)

            # ephemeris table data
            data = Eph(
                objid=objid,
                jd=jd[i],
                rh=eph['r'][i].value,
                delta=eph['Delta'][i].value,
                ra=eph['RA'][i].to('rad').value,
                dec=eph['Dec'][i].to('rad').value,
                dra=eph['dRA cos(Dec)'][i].to('arcsec/hr').value,
                ddec=eph['ddec'][i].to('arcsec/hr').value,
                unc_a=eph['SMAA_3sigma'][i].to('rad').value,
                unc_b=eph['SMIA_3sigma'][i].to('rad').value,
                unc_theta=eph['Theta_3sigma'][i].to('rad').value,
                vmag=vmag[i],
                segment=segment,
                retrieved=today
            )
            self.session.add(data)
            count += 1

        self.session.commit()

        return count

    def add_found(self, objid, observations, update=False, cache=False):
        """Add found objects to found database.

        Parameters
        ----------
        objid : int
            Found object ID.

        observations : array of Obs
            Observations with found objects, all from a single data source.

        update : bool, optional
            ``True`` to overwrite existing entries.

        cache : bool, optional
            Use cached ephemerides; primarily for testing.

        Returns
        -------
        foundids : list
            All found IDs.

        newids : list
            Just the newly found IDs.

        """

        # verify that all observations are from a single data source
        sources = list(set([type(obs) for obs in observations]))
        if len(sources) > 1:
            source_list = ', '.join(
                [source.__data_source_name__ for source in sources])
            raise MultipleSourceError(
                'Must only add one source at a time, but found {}: {}'
                .format(len(sources), source_list))
        source = sources[0]

        jd = np.array([(obs.jd_start + obs.jd_stop) / 2
                       for obs in observations])

        jd_sorted, unsort_jd = np.unique(jd, return_inverse=True)
        desg = self.resolve_object(objid)[1]
        eph = ephem.generate(desg, source.__obscode__, jd_sorted, source='jpl',
                             cache=cache)
        orb = ephem.generate_orbit(desg, jd_sorted, cache=cache)

        # restore requested order
        eph = eph[unsort_jd]
        orb = orb[unsort_jd]

        vmag = util.vmag_from_eph(eph)
        sangle = Angle(eph['sunTargetPA'] - 180 * u.deg)
        sangle = sangle.wrap_at(360 * u.deg).deg
        vangle = Angle(eph['velocityPA'] - 180 * u.deg)
        vangle = vangle.wrap_at(360 * u.deg).deg
        Tp = Time(orb['Tp'], format='jd', scale='tt').utc.jd
        tmtp = jd - Tp

        found = []
        new = []
        for i, obs in enumerate(observations):
            try:
                f = (self.session.query(Found)
                     .filter(Found.objid == objid)
                     .filter(Found.obsid == obs.obsid)
                     .one())
            except NoResultFound:
                f = None

            if f is None:
                # create a new row
                f = Found(objid=objid, obsid=obs.obsid)
                new.append(True)
            else:
                new.append(False)
                if not update:
                    # not new, not updating, nothing to do
                    found.append(f)
                    continue

            f.jd = jd[i]
            f.ra = eph['ra'][i].to('deg').value
            f.dec = eph['dec'][i].to('deg').value
            f.dra = eph['dRA cos(Dec)'][i].to('arcsec/hr').value
            f.ddec = eph['ddec'][i].to('arcsec/hr').value
            f.unc_a = eph['SMAA_3sigma'][i].to('arcsec').value
            f.unc_b = eph['SMIA_3sigma'][i].to('arcsec').value
            f.unc_theta = eph['Theta_3sigma'][i].to('deg').value
            f.vmag = vmag[i]
            f.rh = eph['r'][i].to('au').value
            f.rdot = eph['r_rate'][i].to('km/s').value
            f.delta = eph['Delta'][i].to('au').value
            f.phase = eph['alpha'][i].to('deg').value
            f.selong = eph['elong'][i].to('deg').value
            f.sangle = sangle[i]
            f.vangle = vangle[i]
            f.trueanomaly = orb['nu'][i].to('deg').value
            f.tmtp = tmtp[i]

            if new[-1]:
                self.session.add(f)

            found.append(f)

        self.session.commit()

        foundids = [f.foundid for f in found]
        newids = list(np.array(foundids)[new])
        return foundids, newids

    def add_found_by_id(self, objid, obsids, **kwargs):
        """Add found objects to found database using observation ID.

        Parameters
        ----------
        objid : int
            Found object ID.

        obsids : array-like
            Observation IDs with found object.  Must be from a single source.

        **kwargs
            Any ``add_found`` keyword arguments.

        Returns
        -------
        foundids : list
            All found IDs.

        newids : list
            Just the newly found IDs.

        """

        observations = self.get_observations_by_id(obsids).all()
        return self.add_found(objid, observations, **kwargs)

    def add_object(self, desg):
        """Add new object to object database.

        Parameters
        ----------
        desg : string
            Object designation.

        Returns
        -------
        objid : int
            The new object ID.

        """

        if not isinstance(desg, str):
            raise ValueError('desg must be a string')

        obj = Obj(desg=desg)
        self.session.add(obj)
        self.session.commit()
        return obj.objid

    def add_observations(self, observations, update=False, logger=None):
        """Add observations to database.

        If observations already exist for a given observation ID, the
        old data are updated.

        Parameters
        ----------
        observations : list of Obs
            Observations to insert.

        update : bool, optional
            Update database in case of duplicates

        logger : `~logging.Logger`, optional
            Log messages to this logger.

        Returns
        -------
        n : int
            Number of inserted or updated rows.

        """

        # autoincrement work around for Postgres
        if self.engine.dialect.name == 'postgresql':
            self.session.execute('''
            SELECT setval('obs_obsid_seq', MAX(obsid)) FROM obs
            ''')

        n = 0
        for obs in observations:
            if update:
                self.session.merge(obs)
            else:
                self.session.add(obs)
            n += 1

        try:
            self.session.commit()
        except Exception as e:
            if logger is not None:
                logger.debug(str(e))
            self.session.rollback()

        return n

    def clean_ephemeris(self, objid, jd_start, jd_stop):
        """Remove ephemeris between dates (inclusive).

        Parameters
        ----------
        objid: int
            Object ID.

        jd_start, jd_stop: float or ``None``
            Julian date range(inclusive), or ``None`` for unbounded.

        Returns
        -------
        n: int
            Number of removed rows.

        """

        eph = self.session.query(Eph).filter_by(objid=objid)
        eph = util.filter_by_date_range(eph, jd_start, jd_stop, Eph.jd).all()

        count = len(eph)
        for e in eph:
            self.session.delete(e)
        self.session.commit()

        return count

    def clean_found(self, objid, jd_start, jd_stop):
        """Remove found objects.

        Parameters
        ----------
        objid : int
            Object ID or ``None`` for all objects.

        jd_start, jd_stop : float or ``None``
            Julian date range(inclusive), or ``None`` for unbounded.

        Returns
        -------
        count : int
            Number of rows removed.

        """

        found = self.session.query(Found)
        if objid is not None:
            found = found.filter_by(objid=objid)
        found = util.filter_by_date_range(
            found, jd_start, jd_stop, Found.jd)

        count = 0
        for f in found:
            count += 1
            self.session.delete(f)
        self.session.commit()

        return count

    def get_ephemeris(self, objid, jd_start, jd_stop):
        """Get ephemeris data from database.

        Parameters
        ----------
        objid : int or ``None``
            Database object ID or ``None`` for any object.

        jd_start, jd_stop : float
            Julian date range (inclusive).

        Returns
        -------
        eph : sqlalchemy Query

        """

        eph = self.session.query(Eph)

        if objid is not None:
            eph = eph.filter_by(objid=objid)

        eph = util.filter_by_date_range(eph, jd_start, jd_stop, Eph.jd)

        return eph

    def get_ephemeris_date_range(self, objids=None):
        """Ephemeris date limits.

        Parameters
        ----------
        objids : list, optional
            Limit query to these object IDs.

        Returns
        -------
        jd_min, jd_max : float

        """

        eph = self.session.query(sa.func.min(Eph.jd).label('jd_min'),
                                 sa.func.max(Eph.jd).label('jd_max'))

        if objids:
            eph = eph.filter(Eph.objid.in_(objids))

        jd_min, jd_max = eph.one()
        if None in [jd_min, jd_max]:
            raise NoEphemerisError('No ephemerides for {}'.format(objids))

        return jd_min, jd_max

    def get_ephemeris_interp(self, objid, epochs):
        """Get ephemeris at specific epochs by interpolation.

        Parameters
        ----------
        objid : int
            Object ID.

        epochs : array-like
            Compute ephemeris at these epochs.  Must be floats (for
            Julian date) or parsable by `~astropy.time.Time`.

        Returns
        -------
        eph : RADec
            Ephemeris.
        vmag : ndarray
            Apparent magnitude estimate.

        """

        ra, dec, vmag = [], [], []
        for jd0 in util.epochs_to_jd(epochs):
            # get two nearest points to epoch
            dt = sa.func.abs(Eph.jd - jd0).label('dt')
            eph = (
                self.session.query(Eph, dt)
                .filter(Eph.objid == objid)
                .filter(dt < 5)
                .order_by(dt)
                .limit(2)
                .all()
            )

            jd = [row[0].jd for row in eph]
            _ra = [row[0].ra for row in eph]
            _dec = [row[0].dec for row in eph]
            _vmag = [row[0].vmag for row in eph]
            dt = [row[1] for row in eph]

            i = np.argmin(jd)
            j = np.argmax(jd)

            a = RADec(_ra[i], _dec[i], unit='rad')
            b = RADec(_ra[j], _dec[j], unit='rad')

            c = util.spherical_interpolation(a, b, jd[i], jd[j], jd0)
            ra.append(c.ra)
            dec.append(c.dec)

            vmag.append(np.interp(jd0, jd, _vmag))

        return RADec(ra, dec, unit='rad'), np.array(vmag)

    def get_found(self, obj=None, start=None, stop=None):
        """Get found objects by object and/or date range.

        Parameters
        ----------
        obj : int or string, optional
            Find detections of this object.

        start, stop : string or `~astropy.time.Time`, optional
            Date range to search, inclusive.  ``None`` for unbounded
            limit.

        Returns
        -------
        found : sqalchemy Query
            Matching found objects.

        """

        start, stop = util.epochs_to_jd((start, stop))

        found = self.session.query(Found)

        if obj is not None:
            objid = self.resolve_object(obj)[0]
            found = found.join(Obj).filter(Found.objid == objid)

        if start is not None:
            found = found.filter(Found.jd >= start)

        if stop is not None:
            found = found.filter(Found.jd <= stop)

        return found

    def get_found_by_id(self, foundids):
        """Get found objects by found ID.

        Parameters
        ----------
        foundids : array-like, optional
            Found IDs to retrieve.

        Returns
        -------
        found : sqalchemy Query
            Matching found objects.

        """

        found = self.session.query(Found).filter(
            Found.foundid.in_(foundids))
        return found

    def get_found_by_obsid(self, obsids):
        """Get found objects by observation ID.

        Parameters
        ----------
        obsids : array-like
            Observation IDs to search.

        Returns
        -------
        found : sqalchemy Query
            Matching found objects.

        """

        found = self.session.query(Found).filter(
            Found.objid.in_(obsids))
        return found

    def get_objects(self):
        """Return all objects.

        Returns
        -------
        objects : sqlalchemy Query
            All objects.

        """

        return self.session.query(Obj)

    def get_observation_date_range(self, source=Obs):
        """Observation date limits.

        Parameters
        ----------
        source : sqlalchemy mapping, optional
            Source for observations.  Use ``schema.Obs`` to search all
            sources.  Otherwise, pass the specific source object.

        Returns
        -------
        jd_min, jd_max : float

        Raises
        ------
        SourceNotFoundError

        """

        query = self.session.query(
            sa.func.min(source.jd_start),
            sa.func.max(source.jd_stop))

        jd_min, jd_max = query.one()
        if None in [jd_min, jd_max]:
            if source:
                msg = 'No observations for source: ' + source.__tablename__
            else:
                msg = 'No observations in database'
            raise SourceNotFoundError(msg)

        return jd_min, jd_max

    def get_observations_by_id(self, obsids):
        """Get observations by observation ID.

        Parameters
        ----------
        obsids: array-like
            Observation IDs to retrieve.

        Returns
        -------
        obs : sqlalchemy Query
            Matched observations.

        """

        obs = self.session.query(Obs).filter(Obs.obsid.in_(obsids))
        return obs

    def get_observations_by_date(self, start=None, stop=None, source=Obs):
        """Get observations by observation date.

        Parameters
        ----------
        start, stop : string or `~astropy.time.Time`, optional
            Date range to search, inclusive.

        source : sqlalchemy mapping, optional
            Source for observations.  Use ``schema.Obs`` to search all
            sources.  Otherwise, pass the specific source object.

        Returns
        -------
        obs : sqlalchemy Query
            Matched observations.

        """

        start, stop = util.epochs_to_jd((start, stop))

        obs = self.session.query(source)

        if start is not None:
            obs = obs.filter(source.jd_stop >= start)

        if stop is not None:
            obs = obs.filter(source.jd_start <= stop)

        return obs

    def get_observations_covering(self, shape, start=None, stop=None,
                                  source=Obs):
        """Find observations covering the given shape.

        Parameters
        ----------
        shape : string
            PostGIS WKT spatial object.

        start, stop: float or `~astropy.time.Time`, optional
            Search this time range.  Floats are Julian dates.

        source : sqlalchemy mapping, optional
            Source for observations.  Use ``schema.Obs`` to search all
            sources.  Otherwise, pass the specific source object.

        Returns
        -------
        obs : sqlalchemy Query
            Matching observations.

        """

        #obs = self.get_observations_by_date(start, stop, source=source)
        #obs = obs.filter(source.fov.ST_Covers(shape))

        start, stop = util.epochs_to_jd((start, stop))

        obs = self.session.query(source).filter(source.fov.ST_Covers(shape))

        if start is not None:
            obs = obs.filter(source.jd_stop >= start)

        if stop is not None:
            obs = obs.filter(source.jd_start <= stop)

        return obs

    def get_observations_intersecting(self, shape, start=None, stop=None,
                                      source=Obs):
        """Find observations intersecting the given shape.

        Parameters
        ----------
        shape : string or list of strings
            PostGIS WKT spatial object(s).

        start, stop: float or `~astropy.time.Time`, optional
            Search this time range.  Floats are Julian dates.  If `shape` is a
            list, then `start` and/or `stop` may be a list.

        source : sqlalchemy mapping, optional
            Source for observations.  Use ``schema.Obs`` to search all
            sources.  Otherwise, pass the specific source object.

        Returns
        -------
        obs : sqlalchemy Query
            Matching observations.

        """

        if np.iterable(start):
            start = util.epochs_to_jd(start)
        else:
            start = util.epochs_to_jd([start])[0]
        if np.iterable(stop):
            stop = util.epochs_to_jd(stop)
        else:
            stop = util.epochs_to_jd([stop])[0]
        obs = self.session.query(source)

        filters = []
        if isinstance(shape, (list, tuple)):
            _starts = start if np.iterable(start) else cycle(start)
            _stops = stop if np.iterable(stop) else cycle(stop)
            for _shape, _start, _stop in zip(shape, _starts, _stops):
                filt = [source.fov.ST_Intersects(_shape)]
                if _start is not None:
                    filt.append(source.jd_stop >= _start)

                if _stop is not None:
                    filt.append(source.jd_start <= _stop)

                filters.append(sa.and_(*filt))

            obs = obs.filter(sa.or_(*filters))
        else:
            obs = obs.filter(source.fov.ST_Intersects(shape))
            if start is not None:
                obs = obs.filter(source.jd_stop >= start)

            if stop is not None:
                obs = obs.filter(source.jd_start <= stop)

        return obs

    def observation_covers(self, obs, c):
        """Test if the observation covers the coordinate.

        Parameters
        ----------
        obs : Obs
            Observation to test.

        c : RADec
            Coordinate to test.

        Returns
        -------
        covered : bool

        """

        point = str(Point(c))
        covered = self.session.query(obs.fov.ST_Covers(point)).scalar()
        return covered

    def resolve_objects(self, objects):
        """Resolve objects to database object ID and designation.

        Parmeters
        ---------
        obj: array of str or int
            Objects to resolve: use strings for designation, ints for
            object ID.

        Returns
        -------
        objects : tuple
            (objid, desg) where object ID is ``None`` if not found.

        Raises
        ------
        ``BadObjectID`` if an object ID is provided but not in the
        database.

        """

        return tuple((self.resolve_object(obj) for obj in objects))

    def resolve_object(self, obj):
        """Resolve object to database object ID and designation.

        Parmeters
        ---------
        obj: str or int
            Object to resolve: use strings for designation, int for
            object ID.

        Returns
        -------
        objid: int
            Object ID or ``None`` if a designation is queried but not
            in the database.

        desg: str
            Object designation.

        Raises
        ------
        ``BadObjectID`` if an object ID is provided but not in the
        database.

        """

        query = self.session.query(Obj.objid, Obj.desg)

        if isinstance(obj, str):
            query = query.filter(Obj.desg == obj)
        else:
            query = query.filter(Obj.objid == obj)

        try:
            objid, desg = query.one()
        except NoResultFound:
            if isinstance(obj, str):
                objid = None
                desg = obj
            else:
                raise BadObjectID('{} not found in database'.format(obj))

        return objid, desg

import numpy,import astropy,from astropy