#!/usr/bin/env python
from astropy.io import fits
import numpy as np
from scipy import ndimage, signal, optimize, linalg, stats
import multiprocessing
import re
import glob
import os
import sys
import time


class Consumer(multiprocessing.Process):

    def __init__(self, task_queue, result_queue):
        multiprocessing.Process.__init__(self)
        self.task_queue = task_queue
        self.result_queue = result_queue

    def run(self):
        proc_name = self.name
        while True:
            next_task = self.task_queue.get()
            if next_task is None:
                # Poison pill means we should exit
                break
            self.result_queue.put(next_task())
        return


class Task(object):
    def __init__(self, index, func, args):
        self.index = index
        self.func = func
        self.args = args

    def __call__(self):
        return self.index, self.func(*self.args)


def _smooth(im, ivar, sig=1, spline_filter=False):
    '''
    Private function _smooth smooths an image accounting for the
    inverse variance.  It optionally spline filters the result to save
    time in later calls to ndimage.map_coordinates.

    Args:
        im: ndarray, 2D image to smooth
        ivar: ndarray, 2D inverse variance, shape should match im
        sig: float, standard deviation of Gaussian smoothing kernel, default 1
        spline_filter: boolean, spline filter the result?  Default False.

    Returns:
        imsmooth : ndarray, smoothed image of the same size as im
    '''

    if not isinstance(im, np.ndarray) or not isinstance(ivar, np.ndarray):
        raise TypeError("image and ivar passed to _smooth must be ndarrays")
    if im.shape != ivar.shape or len(im.shape) != 2:
        raise ValueError("image and ivar must be 2D ndarrays of the same shape")

    nx = int(sig * 4 + 1) * 2 + 1
    x = np.arange(nx) - nx // 2
    x, y = np.meshgrid(x, x)

    window = np.exp(-(x ** 2 + y ** 2) / (2 * sig ** 2))
    imsmooth = signal.convolve2d(im * ivar, window, mode='same')
    imsmooth /= signal.convolve2d(ivar, window, mode='same') + 1e-35
    imsmooth *= im != 0

    if spline_filter:
        imsmooth = ndimage.spline_filter(imsmooth)

    return imsmooth


def _spotloc(phi, sep, pitch=15, D=8.2):
    '''
    Private function _spotloc computes the location of four
    satellite spots in units of lambda

    Args:
        phi: float, the angle of spots in degrees
        sep: float, the separation of the spots in units of lambda/D
        pitch: float, lenslet pitch in units of milliarcseconds. Default 15
        D: float, telescope effective aperture in meters. Default 8.2

    Returns:
        r : ndarray, array of (identical) separations in units of lenslets/microns
        phi : ndarray, array of spot angles in radians
    '''

    phi = np.arange(4) * np.pi / 2 + phi
    r = sep * 1e-6 / D * 3600 * 180 / np.pi / (pitch * 1e-3)
    r = r * np.ones(phi.shape)

    return r, phi


def _par_to_dx_dy(p, lam):
    '''
    Private function _par_to_dx_dy.  Use the parameters given by
    argument p (as a list), together with the wavelength array, to
    compute the offsets as a function of wavelength given as a
    second-order polynomial in wavelength.

    Args:
        p: list of floats, the coefficients of the polynomial fit to the centroid
        lam: wavelength(s), either as a number or as an array

    Returns:
        dx, dy : tuple of 1D arrays, each the same shape as lam, with the wavelength-dependent offsets

    '''

    order = (len(p) - 2) // 2
    dx = (lam * 0 + 1.) * p[0]
    dy = (lam * 0 + 1.) * p[1 + order]

    for i in range(1, order + 1):
        dx += p[i] * lam ** i
        dy += p[1 + i + order] * lam ** i

    return dx, dy


def _spotintens(p, cube, lam):
    '''
    Private function _spotintens computes the negative of the sum of
    the intensity of the four satellite spots induced by the SCExAO DM
    in XYdiag mode.  It is intended to be passed to a minimization
    routine (to maximize sum of the intensities).

    Args:
        p: list of floats
           p[0] is the angle of spots in degrees
           p[1] is the separation in lambda/D
           p[2] - p[-1] are the coefficients of the polynomial fit to the centroid
        cube: 3D ndarray, input data cube, assumed to be smoothed and spline filtered
        lam: 1D ndarray, wavelength array in microns corresponding to the first axis of cube

    Returns:
        sumval : float, the negative of the sum of the spot intensities at the four locations
                 given by p, with integer multiples of pi/2 added to p[0].

    '''

    if not isinstance(cube, np.ndarray):
        raise TypeError("cube must be a 3-dimensional ndarray")
    if len(cube.shape) != 3:
        raise TypeError("cube must be a 3-dimensional ndarray")
    if len(cube) != len(lam):
        raise ValueError("Dimensions of lam and cube must match")

    r, phi = _spotloc(p[0], p[1])
    rcosphi = r * np.cos(phi)
    rsinphi = r * np.sin(phi)
    sumval = 0

    dx, dy = _par_to_dx_dy(p[2:], lam)

    for i in range(cube.shape[0]):
        x = rcosphi * lam[i] + cube.shape[2] // 2 + dx[i]
        y = rsinphi * lam[i] + cube.shape[1] // 2 + dy[i]

        sumval -= np.sum(ndimage.map_coordinates(cube[i], [y, x], prefilter=False))

    return sumval


def _resid(im1, im2):
    '''
    Private function _resid takes two images, and subtracts a constant
    and a scaled version of the second image from the first.  It
    returns the sum of the squared residuals.

    Args:
        im1: ndarray, image
        im2: ndarray, image of the same dimensions as im1

    Returns:
        chisq : float, sum of the squared residuals after subtracting the
                best fit constant+im2 model from im1

    '''

    if not isinstance(im1, np.ndarray) or not isinstance(im2, np.ndarray):
        raise TypeError("Images passed to _resid must be ndarrays.")
    if im1.shape != im2.shape:
        raise ValueError("Dimensions of images in _resid must match.")

    Sxx = np.sum(im2 ** 2)
    Sx = np.sum(im2)
    Sy = np.sum(im1)
    Sxy = np.sum(im1 * im2)
    S = im2.shape[0]
    norm = (S * Sxy - Sx * Sy) / (S * Sxx - Sx ** 2)
    zeropt = (Sxx * Sy - Sx * Sxy) / (S * Sxx - Sx ** 2)

    return np.sum((im1 - norm * im2 - zeropt) ** 2)


def _resid_xy(p, im1, im2, x, y):
    '''
    Private function _resid_xy computes the residual between im1 and
    im2, but with an offset applied to the positions of im2 before
    finding the residual.  The residual is only computed at certain
    points given by ndarrays x and y.  This function is intended to be
    passed to a minimization routine.

    The function calls map_coordinates to compute the reference and
    comparison images at the appropriate points, then calls _resid to
    fit the model and compute the residual.

    Args:
        p: list of two floats
           p[0] is the x offset between im1 and im2
           p[1] is the y offset
        im1: 2D ndarray, reference image, assumed to be spline_filtered
        im2: 2D ndarray, comparison image, assumed to be spline_filtered
        x: 2D ndarray, x coordinates of the points at which to fit im2 to im1
        y: 2D ndarray, y coordinates of the points at which to fit im2 to im1

    Returns:
        chisq : float, square of the residuals between im1 and a model made by
                translating im2 and sampling at x and y.

    '''

    dx, dy = [p[0], p[1]]
    ref_im = ndimage.map_coordinates(im1, [y, x], prefilter=False)
    compare_im = ndimage.map_coordinates(im2, [y + dy, x + dx], prefilter=False)

    return _resid(ref_im, compare_im)


def _cc_resid(p, pp, cube, lam, x, y, retarr=False):
    '''
    Private function _cc_resid cross-correlates wavelength slices of a
    data cube to a reference wavelength.  The output is the sum of the
    squares of the lenslet-by-lenslet differences between the
    reference and scaled/shifted slices (after optimizing the relative
    scale), or the scaled/shifted slices themselves (as a template).

    Args:
        p: list of floats, the coefficients of the polynomial fit to the centroid
        pp: list of floats, the coefficients of the reference polynomial fit to the
            centroid (used to keep the reference center in place when optimizing p)
        cube: 3D ndarray, data cube
        lam: 1D ndarray, wavelengths
        x: ndarray with x lenslet coordinates of the cube to use in the calculation
        y: ndarray with y lenslet coordinates of the cube to use in the calculation
        retarr: return the scaled/shifted template rather than the
                           cross-correlation score? Default False.

    Returns:
        If retarr is False (default), return the sum of the squared differences
        of the various wavelength channels with the reference channel.  If
        retarr is True, return the template image, the average of all of the
        scaled and shifted images at the input x and y coordinates.

    '''

    dx1, dy1 = _par_to_dx_dy(p, lam)
    ny, nx = [cube.shape[1], cube.shape[2]]

    # Centroid of the reference wavelength slice

    iref = cube.shape[0] // 2
    dx2, dy2 = _par_to_dx_dy(pp, lam[iref])
    xx = x + dx2
    yy = y + dy2

    if retarr:
        interparr = np.zeros(tuple([cube.shape[0]] + list(x.shape)))
    else:
        ref = ndimage.map_coordinates(cube[iref], [yy + ny // 2, xx + nx // 2], prefilter=False)
        chisq = 0

    for i in range(cube.shape[0]):

        _x = (xx - dx1[i]) * lam[i] / lam[iref] + dx1[i] + nx // 2
        _y = (yy - dy1[i]) * lam[i] / lam[iref] + dy1[i] + ny // 2

        compare = ndimage.map_coordinates(cube[i], [_y, _x], prefilter=False)
        if retarr:
            interparr[i] = compare
        else:
            chisq += _resid(ref, compare)

    if retarr:
        return np.mean(interparr, axis=0)
    else:
        return chisq


def get_sats_satf(p, cube, lam):
    '''
    retrieves the pixel locations of all four satellite spots at each wavelength,
    and the negative sum of the four spot intensities at each wavelength

    Args:
        p: list of floats, the coefficients of the polynomial fit to the centroid
        cube: ndarray, data cube for which the centroid is fitted
        lam: ndarray, wavelengths for the datacube

    Returns:
        sats: pixel locations (in [x,y] format) of all four satellite spots at each wavelength, shape (wvs, 4, 2)
        satf: float, negative sum of the four spot intensities at each wavelength
    '''

    if not isinstance(cube, np.ndarray):
        raise TypeError("data cube must be a 3-dimensional ndarray")
    if len(cube.shape) != 3:
        raise TypeError("data cube must be a 3-dimensional ndarray")
    if len(cube) != len(lam):
        raise ValueError("Dimensions of lam and cube must match")

    r, phi = _spotloc(p[0], p[1])
    rcosphi = r * np.cos(phi)
    rsinphi = r * np.sin(phi)

    dx, dy = _par_to_dx_dy(p[2:], lam)

    sats = np.zeros((len(lam), 4, 2))
    satf = np.zeros((len(lam), 4))

    for i in range(cube.shape[0]):
        x = rcosphi * lam[i] + cube.shape[2] // 2 + dx[i]
        y = rsinphi * lam[i] + cube.shape[1] // 2 + dy[i]

        # note that for sats, indices are recorded as [x,y], not [y,x], to be consistent with pyklip's convention
        sats[i] = np.stack((x, y), axis=1)
        satf[i] = ndimage.map_coordinates(cube[i], [y, x], prefilter=False)

    return sats, satf


def recen(p, cube, lam, sats, satf, n=None, scale=False, head=None, outfile=None,
          mask=None, data_HDU1=True):
    '''
    Function recen recenters the input data according to the offset
    parameters given in the argument p.  Optionally scale the data by
    wavelength to undo the scaling of diffraction.  Return the
    recentered cube.

    Args:
        p: list of floats, the coefficients of the polynomial fit to the centroid
        cube: 3D ndarray, data cube
        lam: 1D array of wavelengths
        sats: fitted satellite spot indices before recentering
        satf: fitted satellite spot fluxes
        n : integer, spatial dimension of recentered cube.  Default original cube size
        scale: boolean, rescale by wavelength?  Default False
        head: fits header for output file.  Default None
        outfile: string or None, name of output file.  Default None
        mask: boolean lenslet mask
        data_HDU1: boolean, write data to HDU1 and leave HDU0 with no data?  Default True.

    Returns:
        3D ndarray, nlam x n x n, recentered (optionally scaled by wavelength) data cube.

    '''

    if n is None:
        n = cube.shape[1]

    ny, nx = [cube.shape[1], cube.shape[2]]
    dx, dy = _par_to_dx_dy(p[2:], lam)

    x = np.arange(n) - n // 2
    x, y = np.meshgrid(x, x)

    cencube = np.zeros((lam.shape[0], x.shape[0], x.shape[1]))
    for i in range(lam.shape[0]):
        if scale:
            _y = y * lam[i] / lam[len(lam) / 2] + dy[i] + ny // 2
            _x = x * lam[i] / lam[len(lam) / 2] + dx[i] + nx // 2
            print('fitted spots for scaled recen cubes not tested for correctness yet')
            sats[i, :] -= [nx // 2, ny // 2]
            sats[i, :, 1] = sats[i, :, 1] * lam[i] / lam[len(lam) / 2] + ny // 2 - dy[i]
            sats[i, :, 0] = sats[i, :, 0] * lam[i] / lam[len(lam) / 2] + nx // 2 - dx[i]
        else:
            _y = y + dy[i] + ny // 2
            _x = x + dx[i] + nx // 2
            sats[i, :, 1] -= dy[i]
            sats[i, :, 0] -= dx[i]

        cencube[i] = ndimage.map_coordinates(cube[i], [_y, _x], prefilter=False)

    if mask is not None:
        k = (mask.shape[0] - n) // 2
        for i in range(cencube.shape[0]):
            cencube[i] *= mask[k:-k, k:-k]

    if data_HDU1:
        out = fits.HDUList(fits.PrimaryHDU(None, head))
        out.append(fits.PrimaryHDU(cencube))
        out.append(fits.PrimaryHDU(np.ones(cencube.shape)))
    else:
        out = fits.HDUList(fits.PrimaryHDU(cencube, head))

    if outfile is not None:
        out.writeto(outfile, overwrite=True)
        savespot_hdr(sats, satf, outfile, hdu_num=int(data_HDU1))
    else:
        print('output directory not specified, recentered cubes not saved...')

    return cencube


def fitrelcen(image1, image2, x, y, method='Powell'):
    '''
    Function fitrelcen fits for the offset between two images without
    any wavelength dependence by calling _resid_xy, minimizing the sum
    of the squared differences between the images (after optimizing
    over the wavelength-dependent relative normalization).

    Args:
        image1: 2D ndarray, first image
        image2:2D ndarray, second image
        x: ndarray, x coordinates of pixels/lenslets to use
        y: ndarray, y coordinates of pixels/lenslets to use
        method : method passed to scipy.optimize.minimize. Default 'Powell'.

    Returns:
        xc, yc : two floating point numbers giving the best-fit offset between
                 image1 and image2

    '''

    xc, yc = optimize.minimize(_resid_xy, [0, 0],
                               (image1, image2, x, y), method=method).x

    return [xc, yc]


def fitcen(cube, ivar, lam, spotsep=None, guess_center_loc=None, i1=1, i2=-1, r1=15, r2=35, spot_dx=4):
    '''
    Function fitcen.  Fit for the center of a CHARIS data cube using
    the satellite spots by maximizing the agreement between scaled
    cube slices around the spot locations.  If no spot locations are
    provided, use only the diffraction pattern itself in an annulus
    around the image center.

    Args:
        cube: 3D ndarray, CHARIS data cube
        ivar: 3D ndarray, inverse variance of the CHARIS data cube
        lam: 1D ndarray, wavelengths
        spotsep: float or None.  If float, separation of the satellite spots in units
                 of lambda/D.  If None, only use the diffraction pattern in an annulus
                 between r1 and r2.
        guess_center_loc: manually specify initial location of image center if necessary, in [x, y] format
        i1: int, first slice of the data cube to use.  Default 1 (skip slice 0)
        i2: int, high limit of slices to use.  Default -1 (skip last slice)
        r1: float, minimum separation from approximate center for the annulus of the
            diffraction pattern to use in centroiding.  Default 15
        r2: float, maximum separation from approximate center for the annulus of the
            diffraction pattern to use in centroiding.  Default 35
        spot_dx: float, radius around spot location to cut out in order to match the
                 spot location as a function of wavelength.  Default 4

    Returns:
        p : list of floats
            p[0] is the angle of spots in degrees
            p[1] is the separation in lambda/D
            p[2] - p[-1] are the coefficients of the polynomial fit to the centroid

    '''

    ####################################################################
    # Lightly smooth the cube before starting.
    ####################################################################

    cubesmooth = cube.copy()
    mask = np.any(cubesmooth, axis=0) != 0
    iref = i1 + len(lam[i1:i2]) // 2

    for i in range(cube.shape[0]):
        cubesmooth[i] = _smooth(cube[i], ivar[i], lam[i] / 3., spline_filter=False)
        cubesmooth[i] *= mask
        cubesmooth[i] = ndimage.spline_filter(cubesmooth[i])
    # print(np.mean(cube), np.mean(cubesmooth), np.std(cube), np.std(cubesmooth))

    if spotsep is not None:
        if np.abs(spotsep - 15.9) < 1:
            phi = -18 * np.pi / 180
        elif np.abs(spotsep - 10) < 1:
            phi = 27 * np.pi / 180
        else:
            print("Must call fitcen with a valid separation for the satellite spots.")
            print("Currently accepted separations: 10 and 15.5 lambda/D.")
            return None

    ####################################################################
    # Initial center using inner region of the cube (not the spots)
    ####################################################################


    x = np.arange(cube.shape[1]) - cube.shape[1] // 2
    x, y = np.meshgrid(x, x)
    indx = np.where((x ** 2 + y ** 2 > r1 ** 2) * (x ** 2 + y ** 2 < r2 ** 2))

    if guess_center_loc is None:
        xc0, yc0 = optimize.minimize(_cc_resid, [0, 0], ([0, 0], cubesmooth[i1:i2], lam[i1:i2], x[indx], y[indx], False),
                                     method='Powell').x
    else:
        xc0 = guess_center_loc[0]
        yc0 = guess_center_loc[1]

    if spotsep is None:
        return [xc0, yc0]

    p0 = [phi, spotsep, xc0, 0, yc0, 0]

    ####################################################################
    # Now estimate the center by fitting for the spot locations at
    # each wavelength
    ####################################################################

    p1 = optimize.minimize(_spotintens, p0, (cubesmooth[i1:i2], lam[i1:i2]), method='Powell').x
    phi, spotsep = [p1[0], p1[1]]

    ####################################################################
    # Use the areas immediately around the spots to get a center, now
    # by cross-correlation.
    ####################################################################

    _r, _phi = _spotloc(phi, spotsep)
    rcosphi = _r * np.cos(_phi) * lam[iref]
    rsinphi = _r * np.sin(_phi) * lam[iref]

    ok = np.zeros(x.shape)
    for i in range(len(_phi)):
        ok += (rcosphi[i] - x) ** 2 + (rsinphi[i] - y) ** 2 < spot_dx ** 2
    indx = np.where(ok != 0)

    p2 = optimize.minimize(_cc_resid, p0[2:], (p1[2:], cubesmooth[i1:i2], lam[i1:i2], x[indx], y[indx]),
                           method='Powell').x

    ####################################################################
    # Pull out the region of the image corresponding to the average spots.
    ####################################################################

    avgspots = _cc_resid(p2, p2, cubesmooth[i1:i2], lam[i1:i2], x, y, retarr=True)

    lam0 = lam[iref]

    ####################################################################
    # Revise the center location one more time, use this new center
    # (derived from the spots) to shift the coefficients of the fit
    ####################################################################

    p3 = optimize.minimize(_spotintens, [phi, spotsep, 0, 0], (np.asarray([avgspots]), np.asarray([lam0])),
                           method='Powell').x

    ####################################################################
    # Use the center derived from the spots to shift the coefficients
    # of the fit
    ####################################################################

    # center:

    # x: (p2[0] + p3[2] + p2[1]*lam0 - p2[0] - p2[1]*lam)*lam/lam0 + p2[0] + p2[1]*lam
    # y: (p2[2] + p3[3] + p2[3]*lam0 - p2[2] - p2[3]*lam)*lam/lam0 + p2[2] + p2[3]*lam

    # x: p2[0] + (p3[2]/lam0 + 2*p2[1])*lam - (p2[1]/lam0)*lam**2
    # y: p2[2] + (p3[3]/lam0 + 2*p2[3])*lam - (p2[3]/lam0)*lam**2

    phi, spotsep = [p3[0], p3[1]]
    cen_coef = [p2[0], p3[2] / lam0 + 2 * p2[1], -p2[1] / lam0,
                p2[2], p3[3] / lam0 + 2 * p2[3], -p2[3] / lam0]

    return list([phi, spotsep]) + cen_coef


def fitcen_parallel(infiles, astrogrid_status=None, astrogrid_sep=None, smooth_coef=True,
                    guess_center_loc=None, maxcpus=multiprocessing.cpu_count() // 2):
    '''
    Function fitcen_parallel.  Centroid a series of CHARIS data cubes
    in parallel using fitcen.  By default, get the wavelengths and
    astrogrid parameters from the headers.  This might fail on early
    CHARIS data before the standardization of headers.

    Args:
        infiles: list of file names for CHARIS data cubes
        astrogrid_status: None or list of astrogrid configurations for SCExAO.
                          If None, try to read the astrogrid configuration from the header.
                          If this fails, assume there is no astrogrid and centroid using the
                          general diffraction pattern.  Default None.
        astrogrid_sep: None or list of astrogrid spot separations in units of lambda/D.
                       If None, try to read from the header.
                       If that fails, centroid using the general diffraction pattern.
                       Default None.
        smooth_coef: boolean.  smooth the nonlinear coefficients of the centroid fit (the terms
                     proportional to lambda and lambda^2) over the sequence of cubes?
                     Default True.
        guess_center_loc: manually specify initial location of image center if necessary, in [x, y] format
        maxcpus: int, maximum number of CPUs to use in parallelization

    Returns:
        [centroid_params, x, y, mask]

        centroid_params : 2D array of centroid parameters, first dimension is the number of files.
                          Second dimension is the length of theh wavelength-dependent model.
        x : x-coordinates of the centroid at the middle wavelength
        y : y-coordinates of the centroid at the middle wavelength
        mask : 1D boolean array, True if astrogrid was on, or None if the astrogrid was never on.

    '''

    ####################################################################
    # First try to load the astrogrid status and spot separations from
    # the FITS headers
    ####################################################################

    if astrogrid_status is None:
        astrogrid_status = []
        astrogrid_sep = []
        for ifile in infiles:
            head = fits.open(ifile)[0].header
            try:
                astrogrid_status += [head['X_GRDST']]
                astrogrid_sep += [head['X_GRDSEP']]
            except:
                astrogrid_status += [None]
                astrogrid_sep += [None]

    tasks = multiprocessing.Queue()
    results = multiprocessing.Queue()
    ncpus = min(multiprocessing.cpu_count(), maxcpus)
    consumers = [Consumer(tasks, results)
                 for i in range(ncpus)]
    for w in consumers:
        w.start()

    fids = [int(re.sub('.*CRSA', '', re.sub('_cube.fits', '', infile)))
            for infile in infiles]

    lamlist = []
    grid_on = np.zeros(len(fids), np.int32)

    ####################################################################
    # Load each file, get the wavelengths, and pass it to fitcen.
    ####################################################################

    for i in range(len(fids)):

        hdulist = fits.open(infiles[i])
        cube = hdulist[1].data
        ivar = hdulist[2].data
        head = hdulist[0].header
        lam = head['lam_min'] * np.exp(np.arange(cube.shape[0]) * head['dloglam'])
        lam *= 1e-3
        lamlist += [lam]

        if astrogrid_status[i] == 'XYdiag':
            spotsep = astrogrid_sep[i]
            grid_on[i] = 1
        elif astrogrid_status[i] is None:
            spotsep = 15.5
        else:
            spotsep = None

        tasks.put(Task(i, fitcen, (cube, ivar, lam, spotsep, guess_center_loc, 1, -1)))

    for i in range(ncpus):
        tasks.put(None)

    centroid_params = None

    for fid in fids:
        index, result = results.get()
        if centroid_params is None:
            centroid_params = np.zeros((len(fids), len(result)))
        centroid_params[index] = result

    x = None
    y = None

    if np.sum(grid_on) == 0:
        mask = None
    else:
        mask = grid_on

    ####################################################################
    # If desired, smooth the nonlinear coefficients (in wavelength) of
    # the centroiding polynomial so that it varies smoothly with time.
    # Do not smooth the offsets, i.e., do not change the centers at a
    # given reference wavelength.  These can be refined later.
    ####################################################################

    if smooth_coef:

        x1 = polyfit(fids, centroid_params[:, 3], mask=mask, return_y=False)
        x2 = polyfit(fids, centroid_params[:, 4], mask=mask, return_y=False)
        y1 = polyfit(fids, centroid_params[:, 6], mask=mask, return_y=False)
        y2 = polyfit(fids, centroid_params[:, 7], mask=mask, return_y=False)
        x = np.zeros(centroid_params.shape[0])
        y = np.zeros(centroid_params.shape[0])

        for i in range(len(fids)):
            p = centroid_params[i]
            lamref = lamlist[i][len(lamlist[i]) // 2]
            dxref = p[3] * lamref + p[4] * lamref ** 2
            dyref = p[6] * lamref + p[7] * lamref ** 2
            dxnew = x1[i] * lamref + x2[i] * lamref ** 2
            dynew = y1[i] * lamref + y2[i] * lamref ** 2

            x[i] = dxref + p[2]
            y[i] = dyref + p[5]

            centroid_params[i, 2] += dxref - dxnew
            centroid_params[i, 5] += dyref - dynew

        centroid_params[:, 3] = x1
        centroid_params[:, 4] = x2
        centroid_params[:, 6] = y1
        centroid_params[:, 7] = y2

    return [centroid_params, x, y, mask]


def fitallrelcen(infiles, r1=15, r2=50, maxcpus=multiprocessing.cpu_count() // 2):
    '''
    Function fitallrelcen.  Fit for the relative centroids between all
    pairs of frames at the central wavelength using the PSF in an
    annulus around the center.  Return the best-fit relative offets.

    Args:
        infiles: list of file names
        r1: int, minimum separation in lenslets from the image center for annular reference region.  Default 15.
        r2: int, maximum separation in lenslets from the image center for annular reference region.  Default 50.
        maxcpus: int, maximum number of CPUs to allocate for parallelization. Default 1/2 of the available CPUs.

    Returns:
        xsol : 1D ndarray of the relative centers in x
        ysol : 1D ndarray of the relative centers in y

    '''

    ncpus = min(multiprocessing.cpu_count(), maxcpus)

    shape = fits.open(infiles[0])[1].data.shape
    iref = shape[0] // 2

    ####################################################################
    # Lightly smooth all images first.
    ####################################################################

    allims = np.zeros((len(infiles), shape[1], shape[2]))

    for i in range(len(infiles)):
        im = fits.open(infiles[i])[1].data[iref]
        ivar = fits.open(infiles[i])[2].data[iref]
        allims[i] = _smooth(im, ivar, 0.5, True)

    tasks = multiprocessing.Queue()
    results = multiprocessing.Queue()
    consumers = [Consumer(tasks, results)
                 for k in range(ncpus)]
    for w in consumers:
        w.start()

    ####################################################################
    # Find the relative centers by cross-correlation, using only
    # lenslets >r1 and <r2 from the nominal center.
    ####################################################################

    ny, nx = [allims.shape[1], allims.shape[2]]
    x = np.arange(nx) - nx // 2
    x, y = np.meshgrid(x, x)
    r = np.sqrt(x ** 2 + y ** 2)
    x += nx // 2
    y += ny // 2

    ok = np.where((r > r1) * (r < r2))
    y = y[ok].copy()
    x = x[ok].copy()

    for i in range(len(infiles)):
        for j in range(i + 1, len(infiles)):
            index = i * (len(infiles)) + j
            tasks.put(Task(index, fitrelcen, (allims[i], allims[j], x, y)))

    for k in range(ncpus):
        tasks.put(None)

    ####################################################################
    # Once we have all offsets between pairs of frames, solve for the
    # best-fit positions.  This is only defined up to a constant, so
    # also add the constraint that the mean position is zero.  We can
    # add an offset later based on the absolute centroiding routine.
    ####################################################################

    arr_x = np.zeros((len(infiles), len(infiles)))
    arr_y = np.zeros((len(infiles), len(infiles)))

    for i in range(len(infiles)):
        for j in range(i + 1, len(infiles)):
            index, result = results.get()
            fid2 = index % len(infiles)
            fid1 = index // len(infiles)

            arr_x[fid1, fid2] = result[0]
            arr_x[fid2, fid1] = -result[0]
            arr_y[fid1, fid2] = result[1]
            arr_y[fid2, fid1] = -result[1]

    n = len(infiles)
    b = np.zeros((n + 1))
    b[:n] = np.sum(arr_x, axis=0)
    arr = np.ones((n, n + 1))
    arr[:, :n] *= -1
    arr[:, :n] += n * np.identity(n)
    xsol = np.linalg.lstsq(arr.T, b, rcond=-1)[0]

    b[:n] = np.sum(arr_y, axis=0)
    ysol = np.linalg.lstsq(arr.T, b, rcond=-1)[0]

    return xsol, ysol


def polyfit(x, y, order=2, clip=2.5, niter=5, mask=None, return_y=True):
    '''
    Smooth a series of points with a polynomial, iteratively clipping
    outliers.

    Args:
        x: 1D ndarray of x coordinates for the polynomial fit
        y: 1D ndarray of y coordinates for the polynomial fit
        order: int, order of the polynomial to fit.  Default 2.
        clip: float, number of sigma outliers to clip.  Default 2.5.
        niter: int, number of iterations of sigma clipping.  Default 5.
        mask: boolean ndarray or None: mask each y value?  Default None
        return_y: boolean, return smoothed y values (as opposed to coefficients of the polynomial fit)?
                  Default True.

    Returns:
        y_smoothed : 1D array, if return_y=True
        coef : array of the polynomial coefficients if return_y=False

    '''

    if len(x) <= order and not return_y:
        return y

    arr = np.ones((len(x), order + 1))
    for i in range(1, order + 1):
        arr[:, i] = (x - np.median(x)) ** i

    ok = np.ones(y.shape)
    fit_vals = y.copy()

    for i in range(niter):

        _arr = arr.copy()
        if mask is not None:
            ok *= mask
        for j in range(order + 1):
            _arr[:, j] *= ok

        fitcoef = np.linalg.lstsq(_arr, y * ok, rcond=-1)[0]

        resid = y.copy()
        fit_vals[:] = 0
        for j in range(order + 1):
            resid -= fitcoef[j] * arr[:, j]
            fit_vals += fitcoef[j] * arr[:, j]

        std = np.sqrt(np.mean(resid[np.where(ok)] ** 2))
        ok = resid ** 2 < clip ** 2 * std ** 2

    if return_y:
        return fitcoef
    else:
        return fit_vals


def specphotcal(infiles, cencoef, aperture=1.):
    '''
    Function specphotcal.  Computes approximate photometry from the
    satellite spots (using aperture photometry) and scale each
    wavelength to this photometric value.  This should crudely put the
    cubes in units of contrast, though it omits the scaling of
    satellite spot intensity with 1/lambda^2.

    Args:
        infiles: list of file names with CHARIS data cubes
        cencoef: 2D ndarray with coefficeitns
        aperture: float, radius of aperture for photometry in units of lambda/D

    Returns:
        all_phot: photocalibration coefficients

    '''

    fids = [int(re.sub('.*CRSA', '', re.sub('_cube.fits', '', infile)))
            for infile in infiles]

    phi = polyfit(fids, cencoef[:, 0], return_y=False)
    sep = polyfit(fids, cencoef[:, 1], return_y=False)

    all_phot = []
    all_x = []
    all_y = []

    for i in range(len(fids)):

        hdulist = fits.open(infiles[i], mode='update')
        im = hdulist[1].data
        head = hdulist[0].header
        lam = head['lam_min'] * np.exp(np.arange(im.shape[0]) * head['dloglam'])
        lam *= 1e-3

        ny, nx = [im.shape[1], im.shape[2]]

        dx, dy = _par_to_dx_dy(cencoef[i, 2:], lam)
        phot = np.zeros(lam.shape)

        x = np.arange(nx)
        y = np.arange(ny)
        x, y = np.meshgrid(x, y)
        _r, _phi = _spotloc(phi[i], sep[i])

        # Manual offsets to account for the fact that two of the spots
        # lie atop spiders.  Check to see if the first spot is atop a
        # spider or not.

        if np.abs(_phi[0] % np.pi - 2.8) < np.pi / 4:
            _dphi = -np.pi / 2 + 12 * np.pi / 180 * np.asarray([1, -1, 1, -1])
        else:
            _dphi = -np.pi / 2 + 12 * np.pi / 180 * np.asarray([-1, 1, -1, 1])

        caltable = []

        for j in range(lam.shape[0]):

            xcoord = np.cos(_phi) * _r * lam[j] + dx[j] + nx // 2
            ycoord = np.sin(_phi) * _r * lam[j] + dy[j] + ny // 2
            xref = np.cos(_phi + _dphi) * _r * lam[j] + dx[j] + nx // 2
            yref = np.sin(_phi + _dphi) * _r * lam[j] + dy[j] + ny // 2

            # CHARIS is approximately critically sampled at 1.15 microns
            radius = 2 * aperture * (lam[j] / 1.15)


            for k in range(len(xcoord)):
                dist = np.sqrt((x - xcoord[k]) ** 2 + (y - ycoord[k]) ** 2)
                distref = np.sqrt((x - xref[k]) ** 2 + (y - yref[k]) ** 2)

                ref = np.sum(im[j][distref < radius])
                phot[j] += np.sum(im[j][dist < radius]) - ref

            caltable += [(dx[j] + nx // 2, dy[j] + ny // 2, phot[j])]
        all_phot += [phot]
        all_x += [dx + nx // 2]
        all_y += [dy + ny // 2]

    return all_phot