import pyklip.klip as klip
import pyklip.spectra_management as spec
import pyklip.fakes as fakes
import pyklip.kpp.stat.stat_utils as stat_utils
import pyklip.empca as empca
import multiprocessing as mp
import ctypes
import numpy as np
import cProfile
import os
import itertools
import copy
import warnings
import astropy.io.fits as fits
import scipy.interpolate as interp
from scipy.stats import norm
import scipy.ndimage as ndi

#Logic to test mkl exists
try:
    import mkl
    mkl_exists = True
except ImportError:
    mkl_exists = False

# Turns parallelism off for debugging purposes
global parallel
debug = False

if debug is True:
    print("You turned on debug mode, so parallelism is off. Code may be slower.")


def _tpool_init(original_imgs, original_imgs_shape, aligned_imgs, aligned_imgs_shape, output_imgs, output_imgs_shape,
                pa_imgs, wvs_imgs, centers_imgs, filenums_imgs, psf_library, psf_library_shape):
    """
    Initializer function for the thread pool that initializes various shared variables. Main things to note that all
    except the shapes are shared arrays (mp.Array).

    Args:
        original_imgs: original images from files to read and align&scale.
        original_imgs_shape: (N,y,x), N = number of frames = num files * num wavelengths
        aligned: aligned and scaled images for processing.
        aligned_imgs_shape: (wv, N, y, x), wv = number of wavelengths per datacube
        output_imgs: output images after KLIP processing
        output_imgs_shape: (N, y, x, b), b = number of different KL basis cutoffs for KLIP routine
                           after KLIP finishes fitting, the final output_imgs is assigned to sub_imgs,
                           in which the b-dimension is swapped to the leading dimension: (b, N, y, x)
        pa_imgs, wvs_imgs: arrays of size N with the PA and wavelength
        centers_img: array of shape (N,2) with [x,y] image center for image frame
        filenums_imgs (np.array): array of size N with the filenumber corresponding to each image. 
        psf_library: array of shape (N_lib, y, x) with N_lib PSF library images
    """
    global original, original_shape, aligned, aligned_shape, output, output_shape, img_pa, img_wv, img_center, img_filenums, \
        psf_lib, psf_lib_shape
    # original images from files to read and align&scale. Shape of (N,y,x)
    original = original_imgs
    original_shape = original_imgs_shape
    # aligned and scaled images for processing. Shape of (wv, N, y, x)
    aligned = aligned_imgs
    aligned_shape = aligned_imgs_shape
    # output images after KLIP processing
    output = output_imgs
    output_shape = output_imgs_shape
    # parameters for each image (PA, wavelegnth, image center, image number)
    img_pa = pa_imgs
    img_wv = wvs_imgs
    img_center = centers_imgs
    img_filenums = filenums_imgs
    psf_lib = psf_library
    psf_lib_shape = psf_library_shape


def _save_spectral_cubes(dataset, pixel_weights, time_collapse, numbasis, flux_cal, outputdirpath, fileprefix):

    '''
    Saves spectral cubes by collapsing dataset along time dimension

    Args:
        dataset: an instance of CHARISData class
        pixel_weights: pixel weights for dataset.output, same shape as dataset.output
        time_collapse: method for time collapse. Support: 'mean', 'weighted-mean', 'median'
        numbasis: number of KL basis vectors to use (can be a scalar or list like). Length of b
        flux_cal: option to calibrate flux
        outputdirpath: output directory
        fileprefix: file prefix

    Returns:
        saves collapsed spectral cubes to output
    '''

    print('time collapsing reduced data of shape (b, N, wv, y, x):{}'.format(dataset.output.shape))

    spectral_cubes = klip.collapse_data(dataset.output, pixel_weights, axis=1, collapse_method=time_collapse)

    if flux_cal:
        spectral_cubes = np.array([dataset.calibrate_output(cube, spectral=True)
                                        for cube in spectral_cubes])

    for KLcutoff, spectral_cube in zip(numbasis, spectral_cubes):
        filename = '{}-KL{}-speccube.fits'.format(fileprefix, KLcutoff)
        filepath = os.path.join(outputdirpath, filename)
        dataset.savedata(filepath, spectral_cube, klipparams=dataset.klipparams.format(numbasis=KLcutoff),
                         filetype="PSF Subtracted Spectral Cube")

    return


def _save_wv_collapsed_images(dataset, pixel_weights, numbasis, time_collapse, wv_collapse, num_wvs,
                              spectrum, spectra_template, flux_cal, outputdirpath, fileprefix, verbose = True):
    """
    Saves KLmode cube, shape (b, y, x), each slice is a 2D image collapsed along both time and
    wavelength dimension for a specific numbasis

    Args:
        dataset: an instance of CHARISData
        pixel_weights: pixel weights for dataset.output, same shape as dataset.output
        numbasis: number of KL basis vectors to use (can be a scalar or list like). Length of b
        time_collapse: method for time collapse. Support: 'median', 'mean', 'weighted-mean', 'weighted-median'
        wv_collapse: method for wavelength collapse. Supports: 'median', 'mean', 'weighted-mean', 'trimmed-mean'
        num_wvs: number of wavelength slices
        spectrum: refer to klip_dataset() docstring
        spectra_template:
        flux_cal: option to calibrate flux
        outputdirpath: output directory
        fileprefix: file prefix

    Returns:
        saves wavelength collapsed images to output
    """
    if verbose is True:
        print('wavelength collapsing reduced data of shape (b, N, wv, y, x):{}'.format(dataset.output.shape))

    spectral_cubes = klip.collapse_data(dataset.output, pixel_weights, axis=1, collapse_method=time_collapse) # spectral_cubes shape (b, wv, y, x)

    if spectrum is None or num_wvs == 1:
        KLmode_cube = klip.collapse_data(spectral_cubes, axis=1, collapse_method=wv_collapse)
    else:
        # if spectrum weighting is carried out, the collapse method is default to mean collapse weighted by spectrum
        # which is an exact copy of the original code before this refactoring
        # TODO: make spectral weighting work with all collapse methods
        print('spectrum weighting turned on...\n'
              'default to mean collapse weighted by pixel_weights * spectra_template...')
        spectra_template = spectra_template.reshape(dataset.output.shape[1:3])  # make same shape as dataset.output
        KLmode_cube = np.nanmean(pixel_weights * dataset.output * spectra_template[None, :, :, None, None], axis=(1, 2))
        KLmode_cube /= np.nanmean(spectra_template[None, :, :, None, None] * pixel_weights, axis=(1, 2))

    if flux_cal:
        KLmode_cube = dataset.calibrate_output(KLmode_cube, spectral=False)
    filename = '{}-KLmodes-all.fits'.format(fileprefix)
    filepath = os.path.join(outputdirpath, filename)
    numbasis_str = '[' + " ".join(str(basis) for basis in numbasis) + ']'
    dataset.savedata(filepath, KLmode_cube,
                     klipparams=dataset.klipparams.format(numbasis=numbasis_str), filetype="KL Mode Cube",
                     zaxis=numbasis)

    return


def _arraytonumpy(shared_array, shape=None, dtype=None):
    """
    Covert a shared array to a numpy array
    Args:
        shared_array: a multiprocessing.Array array
        shape: a shape for the numpy array. otherwise, will assume a 1d array
        dtype: data type of the arrays. Should be either ctypes.c_float(default) or ctypes.c_double

    Returns:
        numpy_array: numpy array for vectorized operation. still points to the same memory!
                     returns None is shared_array is None
    """
    if dtype is None:
        dtype = ctypes.c_float

    # if you passed in nothing you get nothing
    if shared_array is None:
        return None

    numpy_array = np.frombuffer(shared_array.get_obj(), dtype=dtype)
    if shape is not None:
        numpy_array.shape = shape

    return numpy_array


def _align_and_scale_per_image(img_index, aligned_center, ref_wv, dtype=None):
    """
    Aligns and scales the an individual image (used for pyklip lite)

    Args:
        img_index: index of image for the shared arrays
        algined_center: center to align things to
        ref_wv: wavelength to scale images to
        dtype: data type of the arrays. Should be either ctypes.c_float(default) or ctypes.c_double

    Returns:
        None
    """
    if dtype is None:
        dtype = ctypes.c_float

    original_imgs = _arraytonumpy(original, original_shape,dtype=dtype)
    wvs_imgs = _arraytonumpy(img_wv,dtype=dtype)
    centers_imgs = _arraytonumpy(img_center, (np.size(wvs_imgs),2),dtype=dtype)
    aligned_imgs = _arraytonumpy(aligned, aligned_shape,dtype=dtype)

    aligned_imgs[img_index,:,:] = klip.align_and_scale(original_imgs[img_index], aligned_center,
                                                       centers_imgs[img_index], ref_wv/wvs_imgs[img_index])
    return


def _align_and_scale(iterable_arg):
    """
    Aligns and scales the set of original images about a reference center and scaled to a reference wavelength.
    Note: is a helper function to only be used after initializing the threadpool!

    Args:
        iterable_arg: a tuple of three elements:
            ref_wv_iter: a tuple of two elements. First is the index of the reference wavelength (between 0 and 36).
                         second is the value of the reference wavelength. This is to determine scaling
            ref_center: a two-element array with the [x,y] center position to align all the images to.
            dtype: Should be equal to float or np.float32. Define the data type of the arrays.
                    float is actually the default double.

    Returns:
        just returns ref_wv_iter again
    """

    # extract out arguments from the iteration argument
    ref_wv_iter = iterable_arg[0]
    ref_center = iterable_arg[1]
    dtype = iterable_arg[2]
    ref_wv_index = ref_wv_iter[0]
    ref_wv = ref_wv_iter[1]

    original_imgs = _arraytonumpy(original, original_shape,dtype=dtype)
    wvs_imgs = _arraytonumpy(img_wv,dtype=dtype)
    centers_imgs = _arraytonumpy(img_center, (np.size(wvs_imgs),2),dtype=dtype)

    aligned_imgs = _arraytonumpy(aligned, aligned_shape,dtype=dtype)
    aligned_imgs[ref_wv_index, :, :, :] =  np.array([klip.align_and_scale(frame, ref_center, old_center, ref_wv/old_wv,dtype=dtype)
                                                     for frame, old_center, old_wv in zip(original_imgs, centers_imgs, wvs_imgs)])

    return ref_wv_index, ref_wv


def _klip_section(img_num, parang, wavelength, wv_index, numbasis, radstart, radend, phistart, phiend, minmove,
                  ref_center, dtype=None):
    """
    DEPRECIATED. Still being preserved in case we want to change size of atomization. But will need some fixing

    Runs klip on a section of an image as given by the geometric parameters. Helper fucntion of klip routines and
    requires thread pool to be initialized! Currently is designed only for ADI+SDI. Not yet that flexible.

    Args:
        img_num: file index for the science image to process
        parang: PA of science iamge
        wavelength: wavelength of science image
        wv_index: array index of the wavelength of the science image
        numbasis: number of KL basis vectors to use (can be a scalar or list like). Length of b
        avg_rad: average radius of this annulus
        radstart: inner radius of the annulus (in pixels)
        radend: outer radius of the annulus (in pixels)
        phistart: lower bound in CCW angle from x axis for the start of the section
        phiend: upper boundin CCW angle from y axis for the end of the section
        minmove: minimum movement between science image and PSF reference image to use PSF reference image (in pixels)
        ref_center: 2 element list for the center of the science frames. Science frames should all be aligned.
        dtype: data type of the arrays. Should be either ctypes.c_float(default) or ctypes.c_double

    Returns:
        Returns True on success and False on failure.
        Output images are stored in output array as defined by _tpool_init()
    """
    global output, aligned

    if dtype is None:
        dtype = ctypes.c_float

    #create a coordinate system
    x, y = np.meshgrid(np.arange(original_shape[2] * 1.0), np.arange(original_shape[1] * 1.0))
    x.shape = (x.shape[0] * x.shape[1])
    y.shape = (y.shape[0] * y.shape[1])
    r = np.sqrt((x - ref_center[0])**2 + (y - ref_center[1])**2)
    phi = np.arctan2(y - ref_center[1], x - ref_center[0])

    #grab the pixel location of the section we are going to anaylze
    section_ind = np.where((r >= radstart) & (r < radend) & (phi >= phistart) & (phi < phiend))
    if np.size(section_ind) == 0:
        if verbose is True:
            print("section is empty, skipping...")
        return False

    #grab the files suitable for reference PSF
    #load shared arrays for wavelengths and PAs
    wvs_imgs = _arraytonumpy(img_wv,dtype=dtype)
    pa_imgs = _arraytonumpy(img_pa,dtype=dtype)
    #calculate average movement in this section
    avg_rad = (radstart + radend) / 2.0
    moves = klip.estimate_movement(avg_rad, parang, pa_imgs, wavelength, wvs_imgs)
    file_ind = np.where(moves >= minmove)
    if np.size(file_ind[0]) < 1:
        if verbose is True:
            print("less than 1 reference PSFs available for minmove={0}, skipping...".format(minmove))
        return False

    #load aligned images and make reference PSFs
    aligned_imgs = _arraytonumpy(aligned, (aligned_shape[0], aligned_shape[1], aligned_shape[2]*aligned_shape[3]),dtype=dtype)[wv_index]
    ref_psfs = aligned_imgs[file_ind[0], :]
    ref_psfs = ref_psfs[:,  section_ind[0]]
    #ref_psfs = ref_psfs[:, section_ind]
    #print(img_num, avg_rad, ref_psfs.shape)
    #print(sub_imgs.shape)
    #print(sub_imgs[img_num, section_ind, :].shape)

    #write to output
    output_imgs = _arraytonumpy(output, (output_shape[0], output_shape[1]*output_shape[2], output_shape[3]),dtype=dtype)
    klipped = klip.klip_math(aligned_imgs[img_num, section_ind], ref_psfs, numbasis)
    output_imgs[img_num, section_ind, :] = klipped
    return True


def _klip_section_profiler(img_num, parang, wavelength, wv_index, numbasis, radstart, radend, phistart, phiend, minmove,
                           ref_center=None):
    """
    DEPRECIATED. Still being preserved in case we want to change size of atomization. But will need some fixing

    Profiler wrapper for _klip_section. Outputs a file openable by pstats.Stats for each annulus wavelength.
    However there is no guarentee which wavelength and which subsection of the annulus is saved to disk.

    Args: Same arguments as _klip_section
    """
    cProfile.runctx("_klip_section(img_num, parang, wavelength, wv_index, numbasis, radstart, radend, phistart, phiend,"
                    " minmove, ref_center)", globals(), locals(), 'profile-{0}.out'.format(int(radstart+radend)/2))
    return True


def _klip_section_multifile_profiler(scidata_indices, wavelength, wv_index, numbasis, radstart, radend, phistart,
                                     phiend, minmove, ref_center=None, minrot=0):
    """
    Profiler wrapper for _klip_section_multifile. Outputs a file openable by pstats.Stats for each annulus wavelength.
    However there is no guarentee which wavelength and which subsection of the annulus is saved to disk. There
    is the ability to output a profiler file for each subsection and wavelength but it's too many files and who
    actually looks at all of them.

    Args: Same arguments as _klip_section_multifile()
    """
    cProfile.runctx("_klip_section_multifile(scidata_indices, wavelength, wv_index, numbasis, radstart, radend, "
                    "phistart, phiend, minmove, ref_center, minrot)", globals(), locals(),
                    'profile-{0}.out'.format(int(radstart + radend)/2))
    return True


def _klip_section_multifile(scidata_indices, wavelength, wv_index, numbasis, maxnumbasis, radstart, radend, phistart,
                            phiend, minmove, ref_center, minrot, maxrot, spectrum, mode, corr_smooth=1, psflib_good=None,
                            psflib_corr=None, lite=False, dtype=None, algo='klip'):
    """
    Runs klip on a section of the image for all the images of a given wavelength.
    Bigger size of atomization of work than _klip_section but saves computation time and memory. Currently no need to
    break it down even smaller when running on machines on the order of 32 cores.

    Args:
        scidata_indices: array of file indicies that are the science images for this wavelength
        wavelength: value of the wavelength we are processing
        wv_index: index of the wavelenght we are processing
        numbasis: number of KL basis vectors to use (can be a scalar or list like). Length of b
        maxnumbasis: if not None, maximum number of KL basis/correlated PSFs to use for KLIP. Otherwise, use max(numbasis)    
        radstart: inner radius of the annulus (in pixels)
        radend: outer radius of the annulus (in pixels)
        phistart: lower bound in CCW angle from x axis for the start of the section
        phiend: upper boundin CCW angle from y axis for the end of the section
        minmove: minimum movement between science image and PSF reference image to use PSF reference image (in pixels)
        ref_center: 2 element list for the center of the science frames. Science frames should all be aligned.
        minrot: minimum PA rotation (in degrees) to be considered for use as a reference PSF (good for disks)
        maxrot: maximum PA rotation (in degrees) to be considered for use as a reference PSF (temporal variability)
        spectrum: if not None, a array of length N with the flux of the template spectrum at each wavelength. Uses
                    minmove to determine the separation from the center of the segment to determine contamination
                    (e.g. minmove=3, checks how much containmination is within 3 pixels of the hypothetical source)
                    if smaller than 10%, (hard coded quantity), then use it for reference PSF
        mode: one of ['ADI', 'SDI', 'ADI+SDI'] for ADI, SDI, or ADI+SDI
        corr_smooth (float): size of sigma of Gaussian smoothing kernel (in pixels) when computing most correlated PSFs. If 0, no smoothing
        lite: if True, use low memory footprint mode
        dtype: data type of the arrays. Should be either ctypes.c_float(default) or ctypes.c_double
        algo (str): algorithm to use ('klip', 'nmf', 'empca')

    Returns:
        returns True on success, False on failure. Does not return whether KLIP on each individual image was sucessful.
        Saves data to output array as defined in _tpool_init()
    """
    if dtype is None:
        dtype = ctypes.c_float

    #create a coordinate system. Can use same one for all the images because they have been aligned and scaled
    x, y = np.meshgrid(np.arange(original_shape[2] * 1.0), np.arange(original_shape[1] * 1.0))
    x.shape = (x.shape[0] * x.shape[1]) #Flatten
    y.shape = (y.shape[0] * y.shape[1])
    r, phi = klip.make_polar_coordinates(x, y, ref_center) #flattened polar coordinates

    #grab the pixel location of the section we are going to anaylze
    section_ind = np.where((r >= radstart) & (r < radend) & (phi >= phistart) & (phi < phiend))
    if np.size(section_ind) <= 1:
        if verbose is True:
            print("section is too small ({0} pixels), skipping...".format(np.size(section_ind)))
        return False

    #export some of klip.klip_math functions to here to minimize computation repeats

    #load aligned images for this wavelength
    #if lite memory, the aligend array has a different size
    if lite:
        aligned_imgs = _arraytonumpy(aligned, (aligned_shape[0], aligned_shape[1] * aligned_shape[2]),dtype=dtype)
    else:
        aligned_imgs = _arraytonumpy(aligned, (aligned_shape[0], aligned_shape[1], aligned_shape[2] * aligned_shape[3]),dtype=dtype)[wv_index]

    ref_psfs = aligned_imgs[:,  section_ind[0]]

    if algo.lower() == 'empca':

        #TODO: include scidata_indices selection in here
        try:
            full_model = np.zeros(ref_psfs.shape)
            ref_psfs[np.isnan(ref_psfs)] = 0.
            good_ind = np.sum(ref_psfs > 0., axis=0) > numbasis[-1]

            # set weights for empca
            # TODO: change the inner and outer suppression angle to not be hard coded
            rflat = np.reshape(r[section_ind[0]], -1) # 1D array of radii, size=number of pixels in the section
            weights = empca.set_pixel_weights(ref_psfs[:, good_ind], rflat[good_ind], mode='standard', inner_sup=17,
                                              outer_sup=66)

            # run empca reduction
            output_imgs_np = _arraytonumpy(output, (output_shape[0], output_shape[1] * output_shape[2], output_shape[3]), dtype=dtype)
            for i, rank in enumerate(numbasis):
                # get indices of the image section that have enough finite values along the time dimension
                good_ind_model = empca.weighted_empca(ref_psfs[:, good_ind], weights=weights, niter=15, nvec=rank)
                full_model[:, good_ind] = good_ind_model
                output_imgs_np[:, section_ind[0], i] = aligned_imgs[:, section_ind[0]] - full_model

        except (ValueError, RuntimeError, TypeError) as err:
            print(err.args)
            return False

        return True

    
    #do the same for the reference PSFs
    #playing some tricks to vectorize the subtraction of the mean for each row
    ref_psfs_mean_sub = ref_psfs - np.nanmean(ref_psfs, axis=1)[:, None]
    ref_psfs_mean_sub[np.where(np.isnan(ref_psfs_mean_sub))] = 0

    #calculate the covariance matrix for the reference PSFs
    #note that numpy.cov normalizes by p-1 to get the NxN covariance matrix
    #we have to correct for that in the klip.klip_math routine when consturcting the KL
    #vectors since that's not part of the equation in the KLIP paper
    covar_psfs = np.cov(ref_psfs_mean_sub)
    if ref_psfs_mean_sub.shape[0] == 1:
        # EDGE CASE: if there's only 1 image, we need to reshape to covariance matrix into a 2D matrix
        covar_psfs = covar_psfs.reshape((1,1))

    if corr_smooth > 0:
        # calcualte the correlation matrix, with possible smoothing  
        aligned_imgs_3d = aligned_imgs.reshape([aligned_imgs.shape[0], aligned_shape[-2], aligned_shape[-1]]) # make a cube that's not flattened in spatial dimension
        # smooth only the square that encompasses the segment
        # we need to figure where that is
        # figure out the smallest square that encompasses this sector
        blank_img = np.ones(aligned_shape[-2:]) * np.nan
        blank_img.ravel()[section_ind] = 0
        y_good, x_good = np.where(~np.isnan(blank_img))
        ymin = np.min(y_good)
        ymax = np.max(y_good)
        xmin = np.min(x_good)
        xmax = np.max(x_good)
        blank_img_crop = blank_img[ymin:ymax+1, xmin:xmax+1]
        section_ind_smooth_crop = np.where(~np.isnan(blank_img_crop))
        # now that we figured out only the region of interest for each image to smooth, let's smooth that region'
        ref_psfs_smoothed = []
        for aligned_img_2d in aligned_imgs_3d:
            smooth_sigma = 1
            smoothed_square_crop = ndi.gaussian_filter(aligned_img_2d[ymin:ymax+1, xmin:xmax+1], smooth_sigma)
            smoothed_section = smoothed_square_crop[section_ind_smooth_crop]
            smoothed_section[np.isnan(smoothed_section)] = 0
            ref_psfs_smoothed.append(smoothed_section)
    
        corr_psfs = np.corrcoef(ref_psfs_smoothed)
        if ref_psfs_mean_sub.shape[0] == 1:
            # EDGE CASE: if there's only 1 image, we need to reshape the correlation matrix into a 2D matrix
            corr_psfs = corr_psfs.reshape((1,1))
            
        # smoothing could have caused some ref images to have all 0s
        # which would give them correlation matrix entries of NaN
        # 0 them out for now.
        corr_psfs[np.where(np.isnan(corr_psfs))] = 0
    else:
        # if we don't smooth, we can use the covariance matrix to calculate the correlation matrix. It'll be slightly faster
        covar_diag = np.diagflat(1./np.sqrt(np.diag(covar_psfs)))
        corr_psfs = np.dot( np.dot(covar_diag, covar_psfs ), covar_diag)
        

    #grab the parangs
    parangs = _arraytonumpy(img_pa, dtype=dtype)
    filenums = _arraytonumpy(img_filenums, dtype=dtype)

    for file_index, parang, filenum in zip(scidata_indices, parangs[scidata_indices], filenums[scidata_indices]):
        try:
            _klip_section_multifile_perfile(file_index, section_ind, ref_psfs_mean_sub, covar_psfs, corr_psfs,
                                            parang, filenum, wavelength, wv_index, (radstart + radend) / 2.0, numbasis,
                                            maxnumbasis, minmove, minrot, maxrot, mode, psflib_good=psflib_good,
                                            psflib_corr=psflib_corr, spectrum=spectrum, lite=lite, dtype=dtype,
                                            algo=algo)
        except (ValueError, RuntimeError, TypeError) as err:
            print(err.args)
            return False


    #   [_klip_section_multifile_perfile(file_index, section_ind, ref_psfs_mean_sub, covar_psfs,
    #                                   parang, wavelength, wv_index, (radstart + radend) / 2.0, numbasis, minmove)
    #      for file_index,parang in zip(scidata_indices, parangs[scidata_indices])]

    return True


def _klip_section_multifile_perfile(img_num, section_ind, ref_psfs, covar,  corr, parang, filenum, wavelength, wv_index, avg_rad,
                                    numbasis, maxnumbasis, minmove, minrot, maxrot, mode,
                                    psflib_good=None, psflib_corr=None,
                                    spectrum=None, lite=False, dtype=None, algo='klip'):
    """
    Imitates the rest of _klip_section for the multifile code. Does the rest of the PSF reference selection and runs KLIP.

    Args:
        img_num: file index for the science image to process
        section_ind: np.where(pixels are in this section of the image). Note: coordinate system is collapsed into 1D
        ref_psfs: reference psf images of this section
        covar: the covariance matrix of the reference PSFs. Shape of (N,N)
        corr: the correlation matrix of the refernece PSFs. Shape of (N,N)
        parang: PA of science iamage
        filenum (int): file number of science image
        wavelength: wavelength of science image
        wv_index: array index of the wavelength of the science image
        avg_rad: average radius of this annulus
        numbasis: number of KL basis vectors to use (can be a scalar or list like). Length of b
        maxnumbasis: if not None, maximum number of KL basis/correlated PSFs to use for KLIP. Otherwise, use max(numbasis)           
        minmove: minimum movement between science image and PSF reference image to use PSF reference image (in pixels)
        mode: one of ['ADI', 'SDI', 'ADI+SDI'] for ADI, SDI, or ADI+SDI
        psflib_good: array of size N_lib indicating which N_good are good are selected in the passed in corr matrix
        psflib_corr: matrix of size N_sci x N_good with correlation between the target franes and the good RDI PSFs
        spectrum: if not None, a array of length N with the flux of the template spectrum at each wavelength. Uses
                    minmove to determine the separation from the center of the segment to determine contamination and
                    the size of the PSF (TODO: make PSF size another quanitity)
                    (e.g. minmove=3, checks how much containmination is within 3 pixels of the hypothetical source)
                    if smaller than 10%, (hard coded quantity), then use it for reference PSF
        lite: if True, in memory-lite mode
        dtype: data type of the arrays. Should be either ctypes.c_float(default) or ctypes.c_double

    Returns:
        return True on success, False on failure.
        Saves image to output array defined in _tpool_init()
    """
    if dtype is None:
        dtype = ctypes.c_float
        
    # grab the files suitable for reference PSF
    # load shared arrays for wavelengths, PAs, and filenumbers
    wvs_imgs = _arraytonumpy(img_wv, dtype=dtype)
    pa_imgs = _arraytonumpy(img_pa, dtype=dtype)
    filenums_imgs = _arraytonumpy(img_filenums, dtype=dtype)
    # calculate average movement in this section for each PSF reference image w.r.t the science image
    moves = klip.estimate_movement(avg_rad, parang, pa_imgs, wavelength, wvs_imgs, mode)
    # check all the PSF selection criterion
    # enough movement of the astrophyiscal source
    if spectrum is None:
        goodmv = (moves >= minmove)
    else:
        # optimize the selection based on the spectral template rather than just an exclusion principle
        goodmv = (spectrum * norm.sf(moves  -minmove/2.355, scale=minmove/2.355) <= 0.1 * spectrum[wv_index])

    # enough field rotation
    if minrot > 0:
        goodmv = (goodmv) & (np.abs(pa_imgs - parang) >= minrot)

    # if no SDI, don't use other wavelengths
    if "SDI" not in mode.upper():
        goodmv = (goodmv) & (wvs_imgs == wavelength)
    # if no ADI, don't use other parallactic angles
    if "ADI" not in mode.upper():
        goodmv = (goodmv) & (filenums_imgs == filenum)
    # if both aren't in here, we shouldn't be using any frames in the sequence
    if "ADI" not in mode.upper() and "SDI" not in mode.upper():
        goodmv = (goodmv) & False
    include_rdi = "RDI" in mode.upper()

    good_file_ind = np.where(goodmv)
    # Remove reference psfs if they are mostly nans
    ref2rm = np.where(np.nansum(np.isfinite(ref_psfs[good_file_ind[0], :]),axis=1) < 5)[0]
    good_file_ind = (np.delete(good_file_ind[0],ref2rm),)
    if (np.size(good_file_ind[0]) < 1) and (not include_rdi):
        if verbose is True:
            print("less than 1 reference PSFs available for minmove={0}, skipping...".format(minmove))
        return False
    # pick out a subarray. Have to play around with indicies to get the right shape to index the matrix
    covar_files = covar[good_file_ind[0].reshape(np.size(good_file_ind), 1), good_file_ind[0]]

    # pick only the most correlated reference PSFs if there's more than enough PSFs
    if maxnumbasis is None:
        maxnumbasis = np.max(numbasis)
    maxbasis_possible = np.size(good_file_ind)
    # if RDI, also include the size of the RDI PSF library
    # and load in PSF library
    if include_rdi:
        num_good_rdi = np.size(psflib_good)
        maxbasis_possible += num_good_rdi
        psf_library = _arraytonumpy(psf_lib, (psf_lib_shape[0], psf_lib_shape[1]*psf_lib_shape[2]), dtype=dtype)

    # load input/output data
    if lite:
        aligned_imgs = _arraytonumpy(aligned, (aligned_shape[0], aligned_shape[1] * aligned_shape[2]),dtype=dtype)
    else:
        aligned_imgs = _arraytonumpy(aligned, (aligned_shape[0], aligned_shape[1], aligned_shape[2] * aligned_shape[3]),dtype=dtype)[wv_index]
    numpix = np.size(section_ind[0])

    # do we want to downselect out of all the possible references
    if maxbasis_possible > maxnumbasis:
        # grab the x-correlation with the sci img for valid PSFs
        xcorr = corr[img_num, good_file_ind[0]]
        if include_rdi:
            # calculate real xcorr between image and RDI PSFs for this sector for only the maxnumbasis
            # best reference PSFs.
            # grab the maxnumbasis most correlated PSFs from the library
            num_rdi_psfs_first_downselect = np.min([maxnumbasis, num_good_rdi])
            rdi_best_corr_max_possbile_indices = np.argsort(psflib_corr[img_num, psflib_good])[-num_rdi_psfs_first_downselect:]
            # grab these PSFs
            rdi_best_corr_max_possible = psf_library[psflib_good[rdi_best_corr_max_possbile_indices]]
            rdi_best_corr_max_possible = rdi_best_corr_max_possible[:, section_ind[0]]
            # recalculate their correlations in this sector
            sci_img = aligned_imgs[img_num, section_ind[0]].reshape(1, numpix)
            # to calculate correlation, first subtract off mean for each image
            sci_img_mean_sub = sci_img - np.mean(sci_img, axis=1)[:,None]
            rdi_best_corr_max_possible_mean_sub = rdi_best_corr_max_possible - np.mean(rdi_best_corr_max_possible, axis=1)[:,None]
            # we will then calcualte the covariance between the science image with the PSF Library
            sci_x_rdi_best_covar = np.dot(sci_img_mean_sub, rdi_best_corr_max_possible_mean_sub.T ) / (numpix - 1)
            # to convert from covariance to correlation matrix, we just need to divide by normalization of the datasets
            sci_norm = np.sum(sci_img_mean_sub*sci_img_mean_sub, axis=1)
            rdi_best_norm = np.sum(rdi_best_corr_max_possible_mean_sub*rdi_best_corr_max_possible_mean_sub, axis=1)
            # convert form covaraince matrix to correlation matrix
            sci_x_rdi_best_corr = sci_x_rdi_best_covar / np.sqrt(rdi_best_norm) / np.sqrt(sci_norm)
            # now we've recalculated the correlation for this sector for the best RDI PSFs

            # indicate which are RDI frames (remember we are using onyl the top correlated ones that we just comptued
            # correlations for
            is_rdi_psf = np.append(np.repeat(False, np.size(xcorr)), np.repeat(True, num_rdi_psfs_first_downselect))
            # indices for both the dataset and PSF library arrays squished together
            psfindices = np.append(np.arange(np.size(xcorr)), psflib_good[rdi_best_corr_max_possbile_indices])
            # cross correlation now includes both
            xcorr = np.append(xcorr, sci_x_rdi_best_corr)

        sort_ind = np.argsort(xcorr)
        closest_matched = sort_ind[-maxnumbasis:]  # sorted smallest first so need to grab from the end
        if include_rdi:
            # separate out the RDI ones
            rdi_selected = np.where(is_rdi_psf[closest_matched])
            rdi_closest_matched = psfindices[closest_matched[rdi_selected]]
            # remove the RDI ones from closest_matched to imitate non-RDI behavior
            closest_matched = psfindices[closest_matched[np.where(~is_rdi_psf[closest_matched])]]

        # grab smaller set of reference PSFs
        ref_psfs_selected = ref_psfs[good_file_ind[0][closest_matched], :]
        # grab the new and smaller covariance matrix
        covar_files = covar_files[closest_matched.reshape(np.size(closest_matched), 1), closest_matched]

        if include_rdi:
            rdi_psfs_selected = psf_library[rdi_closest_matched]
            rdi_psfs_selected = rdi_psfs_selected[:, section_ind[0]]
    else:
        # else just grab the reference PSFs for all the valid files
        ref_psfs_selected = ref_psfs[good_file_ind[0], :]

        if include_rdi:
            rdi_psfs_selected = psf_library[psflib_good][:, section_ind[0]]
    
    # add PSF library to reference psf list and covariance matrix if needed
    if include_rdi:

        #subctract the mean and remove the Nans from the RDI PSFs before measuring the covariance.
        # this was already done in _klip_section_multifile for the other PSFs)
        rdi_psfs_selected = rdi_psfs_selected - np.nanmean(rdi_psfs_selected, axis=1)[:, None]
        rdi_psfs_selected[np.where(np.isnan(rdi_psfs_selected))] = 0

        # compute covariances. I could just grab these from ~20 lines above, but too lazy
        rdi_covar = np.cov(rdi_psfs_selected) # N_rdi_sel x N_rdi_sel
        # compute cross term
        # cross term has shape N_dataset_ref x N_rdi_selected
        covar_ref_x_rdi = np.dot((ref_psfs_selected - np.nanmean(ref_psfs_selected, axis=1)[:,None]),
                                 (rdi_psfs_selected - np.nanmean(rdi_psfs_selected, axis=1)[:,None]).T) / (numpix - 1)
        # piece together covariance matrix. It should looke like
        # [ cov_ref, cov_ref_x_rdi ]
        # [ cov_rdi_x_ref, cov_rdi ]
        # first append the horizontal component to get shape of N_all_refs x N_dataset_ref
        covar_files = np.append(covar_files, covar_ref_x_rdi, axis=1)
        # now append the bottom half
        covar_files_bottom = np.append(covar_ref_x_rdi.T, rdi_covar, axis=1)
        covar_files = np.append(covar_files, covar_files_bottom, axis=0)

        # append the rdi psfs to the reference PSFs
        ref_psfs_selected = np.append(ref_psfs_selected, rdi_psfs_selected, axis=0)


    # output_images has shape (N, y*x, b) and not (N, y, x, b) as normal
    output_imgs = _arraytonumpy(output, (output_shape[0], output_shape[1]*output_shape[2], output_shape[3]),dtype=dtype)

    # run KLIP
    try:
        if algo.lower() == 'klip':
            klipped = klip.klip_math(aligned_imgs[img_num, section_ind[0]], ref_psfs_selected, numbasis, covar_psfs=covar_files)
        elif algo.lower() == 'nmf':
            import pyklip.nmf_imaging as nmf_imaging
            klipped = nmf_imaging.nmf_math(aligned_imgs[img_num, section_ind].ravel(), ref_psfs, componentNum=numbasis[0])
            klipped = klipped.reshape(klipped.shape[0], 1)
        elif algo.lower() == "none":
            klipped = np.array([aligned_imgs[img_num, section_ind[0]] for _ in range(len(numbasis))]) # duplicate by requested numbasis
            klipped = klipped.T # retrun in shape (p, b) as expected
    except (ValueError, RuntimeError, TypeError) as err:
        print(err.args)
        return False

    # write to output 
    output_imgs[img_num, section_ind[0], :] = klipped

    return True


def rotate_imgs(imgs, angles, centers, new_center=None, numthreads=None, flipx=False, hdrs=None,
                disable_wcs_rotation = False,pool=None):
    """
    derotate a sequences of images by their respective angles

    Args:
        imgs: array of shape (N,y,x) containing N images
        angles: array of length N with the angle to rotate each frame. Each angle should be CCW in degrees.
        centers: array of shape N,2 with the [x,y] center of each frame
        new_centers: a 2-element array with the new center to register each frame. Default is middle of image
        numthreads: number of threads to be used
        flipx: flip the x axis after rotation if desired
        hdrs: array of N wcs astrometry headers

    Returns:
        derotated: array of shape (N,y,x) containing the derotated images
    """
    if pool is None:
        tpool = mp.Pool(processes=numthreads)
    else:
        tpool = pool

    # klip.rotate(img, -angle, oldcenter, [152,152]) for img, angle, oldcenter
    # multithreading the rotation for each image
    if hdrs is None:
        tasks = [tpool.apply_async(klip.rotate, args=(img, angle, center, new_center, flipx, None))
                 for img, angle, center in zip(imgs, angles, centers)]
    else:
        tasks = [tpool.apply_async(klip.rotate, args=(img, angle, center, new_center, flipx, None))
                 for img, angle, center in zip(imgs, angles, centers)]
        # lazy hack around the fact that wcs objects don't preserve wcs.cd fields when sent to other processes
        # so let's just do it manually outside of the rotation
        if not disable_wcs_rotation:
            for angle, astr_hdr in zip(angles, hdrs):
                if astr_hdr is None:
                    continue
                klip._rotate_wcs_hdr(astr_hdr, angle, flipx=flipx)

    # reform back into a giant array
    derotated = np.array([task.get() for task in tasks])

    if pool is None:
        tpool.close()
        tpool.join()

    return derotated


def high_pass_filter_imgs(imgs, numthreads=None, filtersize=10, pool=None):
    """
    filters a sequences of images using a FFT

    Args:
        imgs: array of shape (N,y,x) containing N images
        numthreads: number of threads to be used
        filtersize: size in Fourier space of the size of the space. In image space, size=img_size/filtersize
        pool: multiprocessing thread pool (optional). To avoid repeatedly creating one when processing a list of images.

    Returns:
        filtered: array of shape (N,y,x) containing the filtered images
    """

    if pool is None:
        tpool = mp.Pool(processes=numthreads)
    else:
        tpool = pool

    tasks = [tpool.apply_async(klip.high_pass_filter, args=(img, filtersize)) for img in imgs]

    #reform back into a giant array
    filtered = np.array([task.get() for task in tasks])

    if pool is None:
        tpool.close()
        tpool.join()

    return filtered


def generate_noise_maps(imgs, aligned_center, dr, IWA=None, OWA=None, numthreads=None, pool=None):
    """
    Create a noise map for each image.  
    The noise levels are computed using azimuthally averaged noise in the images

    Args: 
        imgs: array of shape (N,y,x) containing N images
        aligned_center: [x,y] location of the center. All images should be aligned to common center
        dr (float): how mnay pixels wide the annulus to compute the noise should be
        IWA (float): inner working angle (how close to the center of the image to start). If none, this is 0
        OWA (float): outer working angle (if None, it is the entire image.)
        numthreads: number of threads to be used
        pool: multiprocessing thread pool (optional). To avoid repeatedly creating one when processing a list of images.

    Returns:
        noise_maps: array of shape (N,y,x) containing N noise maps
    """
    if pool is None:
        tpool = mp.Pool(processes=numthreads)
    else:
        tpool = pool

    if IWA is None:
        IWA = 0
    if OWA is None:
        OWA = np.max(imgs.shape[1:])

    tasks = [tpool.apply_async(stat_utils.get_image_stat_map, args=(img,),
                                kwds={'IOWA' : (IWA,OWA), 'centroid': aligned_center, 'N' : None, 'Dr' : dr, 'type' : 'stddev' }) 
            for img in imgs]

    #reform back into a giant array
    noise_maps = np.array([task.get() for task in tasks])

    if pool is None:
        tpool.close()
        tpool.join()

    return noise_maps


def klip_parallelized_lite(imgs, centers, parangs, wvs, filenums, IWA, OWA=None, mode='ADI+SDI', annuli=5, subsections=4,
                           movement=3, numbasis=None, aligned_center = None, numthreads=None, minrot=0, maxrot=360,
                           annuli_spacing="constant", maxnumbasis=None, corr_smooth=1, 
                           spectrum=None, dtype=None, algo='klip', compute_noise_cube=False, **kwargs):
    """
    multithreaded KLIP PSF Subtraction, has a smaller memory foot print than the original

    Args:
        imgs: array of 2D images for ADI. Shape of array (N,y,x)
        centers: N by 2 array of (x,y) coordinates of image centers
        parangs: N length array detailing parallactic angle of each image
        wvs: N length array of the wavelengths
        filenums (np.array): array of length N of the filenumber for each image
        IWA: inner working angle (in pixels)
        OWA: outer working angle (in pixels)
        mode: one of ['ADI', 'SDI', 'ADI+SDI'] for ADI, SDI, or ADI+SDI
        anuuli: number of annuli to use for KLIP
        subsections: number of sections to break each annuli into
        movement: minimum amount of movement (in pixels) of an astrophysical source
                  to consider using that image for a refernece PSF
        numbasis: number of KL basis vectors to use (can be a scalar or list like). Length of b
        annuli_spacing: how to distribute the annuli radially. Currently three options. Constant (equally spaced), 
                log (logarithmical expansion with r), and linear (linearly expansion with r)
        maxnumbasis: if not None, maximum number of KL basis/correlated PSFs to use for KLIP. Otherwise, use max(numbasis) 
        corr_smooth (float): size of sigma of Gaussian smoothing kernel (in pixels) when computing most correlated PSFs. If 0, no smoothing          
        aligned_center: array of 2 elements [x,y] that all the KLIP subtracted images will be centered on for image
                        registration
        numthreads: number of threads to use. If none, defaults to using all the cores of the cpu
        minrot: minimum PA rotation (in degrees) to be considered for use as a reference PSF (good for disks)
        maxrot: maximum PA rotation (in degrees) to be considered for use as a reference PSF (temporal variability)
        spectrum: if not None, a array of length N with the flux of the template spectrum at each wavelength. Uses
                    minmove to determine the separation from the center of the segment to determine contamination and
                    the size of the PSF (TODO: make PSF size another quanitity)
                    (e.g. minmove=3, checks how much containmination is within 3 pixels of the hypothetical source)
                    if smaller than 10%, (hard coded quantity), then use it for reference PSF
        kwargs: in case you pass it stuff that we don't use in the lite version
        dtype: data type of the arrays. Should be either ctypes.c_float (default) or ctypes.c_double
        algo (str): algorithm to use ('klip', 'nmf', 'empca')
        compute_noise_cube:  if True, compute the noise in each pixel assuming azimuthally uniform noise

    Returns:
        sub_imgs: array of [array of 2D images (PSF subtracted)] using different number of KL basis vectors as
                    specified by numbasis. Shape of (b,N,y,x).
    """

    ################## Interpret input arguments ####################

    # default numbasis if none
    if numbasis is None:
        totalimgs = imgs.shape[0]
        maxbasis = np.min([totalimgs, 100]) # only going up to 100 KL modes by default
        numbasis = np.arange(1, maxbasis + 10, 10)
        print("KL basis not specified. Using default.", numbasis)
    else:
        if hasattr(numbasis, "__len__"):
            numbasis = np.array(numbasis)
        else:
            numbasis = np.array([numbasis])

    # default aligned_center if none:
    if aligned_center is None:
        aligned_center = [np.mean(centers[:,0]), np.mean(centers[:,1])]

    # save all bad pixels
    allnans = np.where(np.isnan(imgs))

    # use first image to figure out how to divide the annuli
    dims = imgs.shape
    x, y = np.meshgrid(np.arange(dims[2] * 1.0), np.arange(dims[1] * 1.0))
    nanpix = np.where(np.isnan(imgs[0]))

    # if user didn't supply how to define IWA
    if OWA is None:
        full_image = True # reduce the full image
        # define OWA as either the closest NaN pixel or edge of image if no NaNs exist
        if np.size(nanpix) == 0:
            OWA = np.sqrt(np.max((x - centers[0][0]) ** 2 + (y - centers[0][1]) ** 2))
        else:
            # grab the NaN from the 1st percentile (this way we drop outliers)
            OWA = np.sqrt(np.percentile((x[nanpix] - centers[0][0]) ** 2 + (y[nanpix] - centers[0][1]) ** 2, 1))
    else:
        full_image = False # don't reduce the full image, only up the the IWA

    #error checking for too small of annuli go here

    #calculate the annuli ranges
    rad_bounds = klip.define_annuli_bounds(annuli, IWA, OWA, annuli_spacing)

    # if OWA wasn't passed in, we're going to assume we reduce the full image, so last sector emcompasses everything
    if full_image:
        # last annulus should mostly emcompass everything
        rad_bounds[annuli - 1] = (rad_bounds[annuli - 1][0], imgs[0].shape[0])

    #divide annuli into subsections
    dphi = 2 * np.pi / subsections
    phi_bounds = [[dphi * phi_i - np.pi, dphi * (phi_i + 1) - np.pi] for phi_i in range(subsections)]
    phi_bounds[-1][1] = np.pi


    #calculate how many iterations we need to do
    global tot_iter
    tot_iter = np.size(np.unique(wvs)) * len(phi_bounds) * len(rad_bounds)

    #before we start, create the output array in flattened form
    #sub_imgs = np.zeros([dims[0], dims[1] * dims[2], numbasis.shape[0]])

    ########################### Create Shared Memory ###################################

    if dtype is None:
        dtype = ctypes.c_float
    mp_data_type = dtype

    #implement the thread pool
    #make a bunch of shared memory arrays to transfer data between threads
    #make the array for the original images and initalize it
    original_imgs = mp.Array(mp_data_type, np.size(imgs))
    original_imgs_shape = imgs.shape
    original_imgs_np = _arraytonumpy(original_imgs, original_imgs_shape,dtype=dtype)
    original_imgs_np[:] = imgs
    #make array for recentered/rescaled image (only big enough for one wavelength at a time)
    unique_wvs = np.unique(wvs)
    recentered_imgs = mp.Array(mp_data_type, np.size(imgs))
    recentered_imgs_shape = imgs.shape
    #make output array which also has an extra dimension for the number of KL modes to use
    output_imgs = mp.Array(mp_data_type, np.size(imgs)*np.size(numbasis))
    output_imgs_np = _arraytonumpy(output_imgs,dtype=dtype)
    output_imgs_np[:] = np.nan
    output_imgs_shape = imgs.shape + numbasis.shape
    #remake the PA, wv, and center arrays as shared arrays
    pa_imgs = mp.Array(mp_data_type, np.size(parangs))
    pa_imgs_np = _arraytonumpy(pa_imgs,dtype=dtype)
    pa_imgs_np[:] = parangs
    wvs_imgs = mp.Array(mp_data_type, np.size(wvs))
    wvs_imgs_np = _arraytonumpy(wvs_imgs,dtype=dtype)
    wvs_imgs_np[:] = wvs
    centers_imgs = mp.Array(mp_data_type, np.size(centers))
    centers_imgs_np = _arraytonumpy(centers_imgs, centers.shape,dtype=dtype)
    centers_imgs_np[:] = centers
    filenums_imgs = mp.Array(mp_data_type, np.size(filenums))
    filenums_imgs_np = _arraytonumpy(filenums_imgs,dtype=dtype)
    filenums_imgs_np[:] = filenums

    tpool = mp.Pool(processes=numthreads, initializer=_tpool_init,
                    initargs=(original_imgs, original_imgs_shape, recentered_imgs, recentered_imgs_shape, output_imgs,
                              output_imgs_shape, pa_imgs, wvs_imgs, centers_imgs, filenums_imgs, None, None), maxtasksperchild=50)

    # SINGLE THREAD DEBUG PURPOSES ONLY
    if debug:
        _tpool_init(original_imgs, original_imgs_shape, recentered_imgs, recentered_imgs_shape, output_imgs,
                              output_imgs_shape, pa_imgs, wvs_imgs, centers_imgs, filenums_imgs, None, None)

    print("Total number of tasks for KLIP processing is {0}".format(tot_iter))
    jobs_complete = 0
    #align and scale the images for each image. Use map to do this asynchronously
    for wv_index, this_wv in enumerate(np.unique(wvs)):
        print("Begin processing of wv {0:.4} with index {1}".format(this_wv, wv_index))
        print("Aligning and scaling imgs")
        recentered_imgs_np = _arraytonumpy(recentered_imgs, recentered_imgs_shape,dtype=dtype)
        
        #multitask this
        tasks = [tpool.apply_async(_align_and_scale_per_image, args=(img_index, aligned_center, this_wv,dtype))
                 for img_index in range(recentered_imgs_shape[0])]

        #save it to shared memory
        for img_index, aligned_img_task in enumerate(tasks):
            aligned_img_task.wait()

        #list to store each threadpool task
        outputs = []
    
        print("Wavelength {1:.4} with index {0} has finished align and scale. Queuing for KLIP".format(wv_index, this_wv))

        #pick out the science images that need PSF subtraction for this wavelength
        scidata_indices = np.where(wvs == this_wv)[0]

        # commented out code to do _klip_section instead of _klip_section_multifile
        # outputs += [tpool.apply_async(_klip_section_profiler, args=(file_index, parang, wv_value, wv_index, numbasis,
        #                                                   radstart, radend, phistart, phiend, movement))
        #                     for phistart,phiend in phi_bounds
        #                 for radstart, radend in rad_bounds
        #             for file_index,parang in zip(scidata_indices, parangs[scidata_indices])]

        #perform KLIP asynchronously for each group of files of a specific wavelength and section of the image
        lite = True

        if not debug:
            outputs += [tpool.apply_async(_klip_section_multifile,
                                          args=(scidata_indices, this_wv, wv_index, numbasis,
                                                maxnumbasis,
                                                radstart, radend, phistart, phiend, movement,
                                                aligned_center, minrot, maxrot, spectrum,
                                                mode, corr_smooth, None, None, lite, dtype, algo))
                        for phistart,phiend in phi_bounds
                        for radstart, radend in rad_bounds]
        else:
            outputs += [_klip_section_multifile(scidata_indices, this_wv, wv_index, numbasis,
                                               maxnumbasis,
                                               radstart, radend, phistart, phiend, movement,
                                               aligned_center, minrot, maxrot, spectrum,
                                               mode, corr_smooth, None, None, lite, dtype, algo)
                        for phistart,phiend in phi_bounds
                        for radstart, radend in rad_bounds]

        #harness the data!
        #check make sure we are completely unblocked before outputting the data
        if not debug:
            for out in outputs:
                out.wait()
                if (jobs_complete + 1) % 10 == 0:
                    print("{0:.4}% done ({1}/{2} completed)".format((jobs_complete+1)*100.0/tot_iter, jobs_complete, tot_iter))
                jobs_complete += 1



    #close to pool now and make sure there's no processes still running (there shouldn't be or else that would be bad)
    print("Closing threadpool")
    tpool.close()
    tpool.join()

    #finished. Let's reshape the output images
    #move number of KLIP modes as leading axis (i.e. move from shape (N,y,x,b) to (b,N,y,x)
    sub_imgs = _arraytonumpy(output_imgs, output_imgs_shape,dtype=dtype)
    sub_imgs = np.rollaxis(sub_imgs.reshape((dims[0], dims[1], dims[2], numbasis.shape[0])), 3)

    #restore bad pixels
    sub_imgs[:, allnans[0], allnans[1], allnans[2]] = np.nan

    # calculate weights for weighted mean if necessary
    if compute_noise_cube:
        print("Computing weights for weighted collapse")
        # figure out ~how wide to make it
        annuli_widths = [annuli_bound[1] - annuli_bound[0] for annuli_bound in rad_bounds]
        dr_spacing = np.min(annuli_widths)
        # generate all teh noise maps. We need to collapse the sub_imgs into 3-D to easily do this
        sub_imgs_shape = sub_imgs.shape
        sub_imgs_flatten = sub_imgs.reshape([sub_imgs_shape[0]*sub_imgs_shape[1], sub_imgs_shape[2], sub_imgs_shape[3]])
        noise_imgs = generate_noise_maps(sub_imgs_flatten, aligned_center, dr_spacing, IWA=IWA, OWA=rad_bounds[-1][1], numthreads=numthreads)
        # reform the 4-D cubes
        noise_imgs = noise_imgs.reshape(sub_imgs_shape) # reshape into a cube with same shape as sub_imgs
    else:
        noise_imgs = np.array([1.])

    return sub_imgs, aligned_center, noise_imgs


def klip_parallelized(imgs, centers, parangs, wvs, filenums, IWA, OWA=None, mode='ADI+SDI', annuli=5, subsections=4, movement=3,
                      numbasis=None, aligned_center=None, numthreads=None, minrot=0, maxrot=360, 
                      annuli_spacing="constant", maxnumbasis=None, corr_smooth=1,
                      spectrum=None, psf_library=None, psf_library_good=None, psf_library_corr=None,
                      save_aligned = False, restored_aligned = None, dtype=None, algo='klip', compute_noise_cube=False, verbose = True):
    """
    Multitprocessed KLIP PSF Subtraction

    Args:
        imgs: array of 2D images for ADI. Shape of array (N,y,x)
        centers: N by 2 array of (x,y) coordinates of image centers
        parangs: N length array detailing parallactic angle of each image
        wvs: N length array of the wavelengths
        filenums (np.array): array of length N of the filenumber for each image
        IWA: inner working angle (in pixels)
        OWA: outer working angle (in pixels)
        mode: one of ['ADI', 'SDI', 'ADI+SDI'] for ADI, SDI, or ADI+SDI
        anuuli: number of annuli to use for KLIP
        subsections: number of sections to break each annuli into
        movement: minimum amount of movement (in pixels) of an astrophysical source
                  to consider using that image for a refernece PSF
        numbasis: number of KL basis vectors to use (can be a scalar or list like). Length of b
        aligned_center: array of 2 elements [x,y] that all the KLIP subtracted images will be centered on for image
                        registration
        numthreads: number of threads to use. If none, defaults to using all the cores of the cpu
        minrot: minimum PA rotation (in degrees) to be considered for use as a reference PSF (good for disks)
        maxrot: maximum PA rotation (in degrees) to be considered for use as a reference PSF (temporal variability)
        annuli_spacing: how to distribute the annuli radially. Currently three options. Constant (equally spaced), 
                        log (logarithmical expansion with r), and linear (linearly expansion with r)
        maxnumbasis: if not None, maximum number of KL basis/correlated PSFs to use for KLIP. Otherwise, use max(numbasis)
        corr_smooth (float): size of sigma of Gaussian smoothing kernel (in pixels) when computing most correlated PSFs. If 0, no smoothing
        spectrum: if not None, a array of length N with the flux of the template spectrum at each wavelength. Uses
                    minmove to determine the separation from the center of the segment to determine contamination and
                    the size of the PSF (TODO: make PSF size another quanitity)
                    (e.g. minmove=3, checks how much containmination is within 3 pixels of the hypothetical source)
                    if smaller than 10%, (hard coded quantity), then use it for reference PSF
        psf_library: array of (N_lib, y, x) with N_lib PSF library PSFs
        psf_library_good: array of size N_lib indicating which N_good are good are selected in the passed in corr matrix
        psf_library_corr: matrix of size N_sci x N_good with correlation between the target franes and the good RDI PSFs
        save_aligned:	Save the aligned and scaled images (as well as various wcs information), True/False
        restore_aligned: The aligned and scaled images from a previous run of klip_dataset
        				(usually restored_aligned = dataset.aligned_and_scaled)
        dtype: data type of the arrays. Should be either ctypes.c_float(default) or ctypes.c_double
        algo (str): algorithm to use ('klip', 'nmf', 'empca')
        compute_noise_cube:  if True, compute the noise in each pixel assuming azimuthally uniform noise

    Returns:
        sub_imgs: array of [array of 2D images (PSF subtracted)] using different number of KL basis vectors as
                    specified by numbasis. Shape of (b,N,y,x).
        aligned_center: (x,y) specifying the common center the output images are aligned to
    """

    ################## Interpret input arguments ####################

    #defaullt numbasis if none
    if numbasis is None:
        totalimgs = imgs.shape[0]
        maxbasis = np.min([totalimgs, 100]) #only going up to 100 KL modes by default
        numbasis = np.arange(1, maxbasis + 5, 10)
        if verbose is True:
            print("KL basis not specified. Using default.", numbasis)
    else:
        if hasattr(numbasis, "__len__"):
            numbasis = np.array(numbasis)
        else:
            numbasis = np.array([numbasis])

    # check which algo we will use and whether the inputs are correct
    if algo.lower() == 'klip':
        pass
    elif algo.lower() == 'empca':
        pass
    elif algo.lower() == 'nmf':
        # check to see the correct nmf packages are installed 
        import pyklip.nmf_imaging as nmf_imaging
        if np.size(numbasis) > 1:
            raise ValueError("NMF can only be run with one basis")
    elif algo.lower() == 'none':
        pass
    else:
        raise ValueError("Algo {0} is not supported".format(algo))


    #default aligned_center if none:
    if aligned_center is None:
        #aligned_center = [int(imgs.shape[2]//2), int(imgs.shape[1]//2)]
        aligned_center = [np.mean(centers[:,0]), np.mean(centers[:,1])]

    # validate RDI has an RDI Library with supporting cast
    if "RDI" in mode.upper():
        if psf_library is None:
            raise ValueError("Need to pass in PSF library fields if you want to do RDI.")
    if psf_library is not None:
        if psf_library_corr is None or psf_library_good is None:
            raise ValueError("Need to pass in correlatoin matrix and good selection array for PSF library")


    #save all bad pixels
    allnans = np.where(np.isnan(imgs))

    #use first image to figure out how to divide the annuli
    #TODO: what to do with OWA
    #need to make the next 10 lines or so much smarter
    dims = imgs.shape
    x, y = np.meshgrid(np.arange(dims[2] * 1.0), np.arange(dims[1] * 1.0))
    nanpix = np.where(np.isnan(imgs[0]))
    # if user didn't supply how to define NaNs
    if OWA is None:
        full_image = True # reduce the full image
        # define OWA as either the closest NaN pixel or edge of image if no NaNs exist
        if np.size(nanpix) == 0:
            OWA = np.sqrt(np.max((x - centers[0][0]) ** 2 + (y - centers[0][1]) ** 2))
        else:
            # grab the NaN from the 1st percentile (this way we drop outliers)
            OWA = np.sqrt(np.percentile((x[nanpix] - centers[0][0]) ** 2 + (y[nanpix] - centers[0][1]) ** 2, 1))
    else:
        full_image = False # don't reduce the full image, only up the the IWA

    #error checking for too small of annuli go here


    #calculate the annuli ranges
    rad_bounds = klip.define_annuli_bounds(annuli, IWA, OWA, annuli_spacing)

    # if OWA wasn't passed in, we're going to assume we reduce the full image, so last sector emcompasses everything
    if full_image:
        # last annulus should mostly emcompass everything
        rad_bounds[annuli - 1] = (rad_bounds[annuli - 1][0], imgs[0].shape[0])

    #divide annuli into subsections
    dphi = 2 * np.pi / subsections
    phi_bounds = [[dphi * phi_i - np.pi, dphi * (phi_i + 1) - np.pi] for phi_i in range(subsections)]
    phi_bounds[-1][1] = np.pi

    # print(rad_bounds)

    #calculate how many iterations we need to do
    global tot_iter
    tot_iter = np.size(np.unique(wvs)) * len(phi_bounds) * len(rad_bounds)

    #before we start, create the output array in flattened form
    #sub_imgs = np.zeros([dims[0], dims[1] * dims[2], numbasis.shape[0]])


    ########################### Create Shared Memory ###################################

    # default to using float precision
    if dtype is None:
        dtype = ctypes.c_float
    # we should use the same datatype for both
    mp_data_type = dtype

    #implement the thread pool
    #make a bunch of shared memory arrays to transfer data between threads
    #make the array for the original images and initalize it
    original_imgs = mp.Array(mp_data_type, np.size(imgs))
    original_imgs_shape = imgs.shape
    original_imgs_np = _arraytonumpy(original_imgs, original_imgs_shape,dtype=dtype)
    original_imgs_np[:] = imgs
    #make array for recentered/rescaled image for each wavelength
    unique_wvs = np.unique(wvs)
    recentered_imgs = mp.Array(mp_data_type, np.size(imgs)*np.size(unique_wvs))
    recentered_imgs_shape = (np.size(unique_wvs),) + imgs.shape
    #make output array which also has an extra dimension for the number of KL modes to use
    output_imgs = mp.Array(mp_data_type, np.size(imgs)*np.size(numbasis))
    output_imgs_np = _arraytonumpy(output_imgs,dtype=dtype)
    output_imgs_np[:] = np.nan
    output_imgs_shape = imgs.shape + numbasis.shape
    #remake the PA, wv, filenums, and center arrays as shared arrays
    pa_imgs = mp.Array(mp_data_type, np.size(parangs))
    pa_imgs_np = _arraytonumpy(pa_imgs,dtype=dtype)
    pa_imgs_np[:] = parangs
    wvs_imgs = mp.Array(mp_data_type, np.size(wvs))
    wvs_imgs_np = _arraytonumpy(wvs_imgs,dtype=dtype)
    wvs_imgs_np[:] = wvs
    centers_imgs = mp.Array(mp_data_type, np.size(centers))
    centers_imgs_np = _arraytonumpy(centers_imgs, centers.shape,dtype=dtype)
    centers_imgs_np[:] = centers
    filenums_imgs = mp.Array(mp_data_type, np.size(filenums))
    filenums_imgs_np = _arraytonumpy(filenums_imgs,dtype=dtype)
    filenums_imgs_np[:] = filenums
    if psf_library is not None:
        psf_lib = mp.Array(mp_data_type, np.size(psf_library))
        psf_lib_np = _arraytonumpy(psf_lib, psf_library.shape, dtype=dtype)
        psf_lib_np[:] = psf_library
        psf_lib_shape = psf_library.shape
    else:
        psf_lib = None
        psf_lib_shape = None

    if restored_aligned is not None:
        recentered_imgs_np = _arraytonumpy(recentered_imgs, recentered_imgs_shape,dtype=dtype)
        recentered_imgs_np[:] = restored_aligned

    tpool = mp.Pool(processes=numthreads, initializer=_tpool_init,
                    initargs=(original_imgs, original_imgs_shape, recentered_imgs, recentered_imgs_shape, output_imgs,
                            output_imgs_shape, pa_imgs, wvs_imgs, centers_imgs, filenums_imgs, psf_lib, psf_lib_shape),
                    maxtasksperchild=50)

    # # SINGLE THREAD DEBUG PURPOSES ONLY
    if debug:
        _tpool_init(original_imgs, original_imgs_shape, recentered_imgs, recentered_imgs_shape, output_imgs,
                            output_imgs_shape, pa_imgs, wvs_imgs, centers_imgs, filenums_imgs, psf_lib, psf_lib_shape)


    if restored_aligned is None:
        #align and scale the images for each image. Use map to do this asynchronously
        if verbose is True:
            print("Begin align and scale images for each wavelength")
        realigned_index = tpool.imap_unordered(_align_and_scale, zip(enumerate(unique_wvs), itertools.repeat(aligned_center),itertools.repeat(dtype)))
    else:
        #align and scale the images for each image. Use map to do this asynchronously
        realigned_index = enumerate(unique_wvs)

    #list to store each threadpool task
    outputs = []
    #as each is finishing, queue up the aligned data to be processed with KLIP
    for wv_index, wv_value in realigned_index:
        if verbose is True:
            print("Wavelength {1:.4} with index {0} has finished align and scale. Queuing for KLIP".format(wv_index, wv_value))

        #pick out the science images that need PSF subtraction for this wavelength
        scidata_indices = np.where(wvs == wv_value)[0]

        # commented out code to do _klip_section instead of _klip_section_multifile
        # outputs += [tpool.apply_async(_klip_section_profiler, args=(file_index, parang, wv_value, wv_index, numbasis,
        #                                                   radstart, radend, phistart, phiend, movement))
        #                     for phistart,phiend in phi_bounds
        #                 for radstart, radend in rad_bounds
        #             for file_index,parang in zip(scidata_indices, parangs[scidata_indices])]

        #perform KLIP asynchronously for each group of files of a specific wavelength and section of the image
        lite = False

        if not debug:
            outputs += [tpool.apply_async(_klip_section_multifile, (scidata_indices, wv_value, wv_index, numbasis,
                                                                        maxnumbasis,
                                                                        radstart, radend, phistart, phiend, movement,
                                                                        aligned_center, minrot, maxrot, spectrum,
                                                                        mode, corr_smooth, 
                                                                        psf_library_good, psf_library_corr, False,
                                                                        dtype, algo))
                        for phistart,phiend in phi_bounds
                        for radstart, radend in rad_bounds]
        else:
            outputs += [_klip_section_multifile(scidata_indices, wv_value, wv_index, numbasis,
                                                maxnumbasis,
                                                radstart, radend, phistart, phiend, movement,
                                                aligned_center, minrot, maxrot, spectrum,
                                                mode, corr_smooth,
                                                psf_library_good, psf_library_corr, False,
                                                dtype, algo)
                        for phistart,phiend in phi_bounds
                        for radstart, radend in rad_bounds]

    #harness the data!
    #check make sure we are completely unblocked before outputting the data
    if not debug:
        if verbose is True:
            print("Total number of tasks for KLIP processing is {0}".format(tot_iter))
        for index, out in enumerate(outputs):
            out.wait()
            if (index + 1) % 10 == 0:
                print("{0:.4}% done ({1}/{2} completed)".format((index+1)*100.0/tot_iter, index, tot_iter))


    #close to pool now and make sure there's no processes still running (there shouldn't be or else that would be bad)
    if verbose is True:
        print("Closing threadpool")
    tpool.close()
    tpool.join()

    #finished. Let's reshape the output images
    #move number of KLIP modes as leading axis (i.e. move from shape (N,y,x,b) to (b,N,y,x)
    sub_imgs = _arraytonumpy(output_imgs, output_imgs_shape,dtype=dtype)
    sub_imgs = np.rollaxis(sub_imgs.reshape((dims[0], dims[1], dims[2], numbasis.shape[0])), 3)

    #restore bad pixe
    sub_imgs[:, allnans[0], allnans[1], allnans[2]] = np.nan

    # calculate weights for weighted mean if necessary
    if compute_noise_cube:
        print("Computing weights for weighted collapse")
        # figure out ~how wide to make it
        annuli_widths = [annuli_bound[1] - annuli_bound[0] for annuli_bound in rad_bounds]
        dr_spacing = np.min(annuli_widths)
        # generate all teh noise maps. We need to collapse the sub_imgs into 3-D to easily do this
        sub_imgs_shape = sub_imgs.shape
        sub_imgs_flatten = sub_imgs.reshape([sub_imgs_shape[0]*sub_imgs_shape[1], sub_imgs_shape[2], sub_imgs_shape[3]])
        noise_imgs = generate_noise_maps(sub_imgs_flatten, aligned_center, dr_spacing, IWA=IWA, OWA=rad_bounds[-1][1], numthreads=numthreads)
        # reform the 4-D cubes
        noise_imgs = noise_imgs.reshape(sub_imgs_shape) # reshape into a cube with same shape as sub_imgs
    else:
        noise_imgs = np.array([1.])

    if save_aligned:
        aligned_and_scaled = _arraytonumpy(recentered_imgs, recentered_imgs_shape, dtype=dtype)
        return sub_imgs, aligned_center, noise_imgs, aligned_and_scaled
    else:
        return sub_imgs, aligned_center, noise_imgs


def klip_dataset(dataset, mode='ADI+SDI', outputdir=".", fileprefix="", annuli=5, subsections=4, movement=3,
                 numbasis=None, numthreads=None, minrot=0, calibrate_flux=False, aligned_center=None,
                 annuli_spacing="constant", maxnumbasis=None, corr_smooth=1, spectrum=None, psf_library=None, 
                 highpass=False, lite=False, save_aligned = False, restored_aligned = None, dtype=None, algo='klip',
                 time_collapse="mean", wv_collapse='mean', verbose = True):
    """
    run klip on a dataset class outputted by an implementation of Instrument.Data

    Args:
        dataset:        an instance of Instrument.Data (see instruments/ subfolder)
        mode:           some combination of ADI, SDI, and RDI (e.g. "ADI+SDI", "RDI")
        outputdir:      directory to save output files
        fileprefix:     filename prefix for saved files
        anuuli:         number of annuli to use for KLIP
        subsections:    number of sections to break each annuli into
        movement:       minimum amount of movement (in pixels) of an astrophysical source
                        to consider using that image for a refernece PSF
        numbasis:       number of KL basis vectors to use (can be a scalar or list like). Length of b
        numthreads:     number of threads to use. If none, defaults to using all the cores of the cpu
        minrot:         minimum PA rotation (in degrees) to be considered for use as a reference PSF (good for disks)
        calibrate_flux: if True calibrate flux of the dataset, otherwise leave it be
        aligned_center: array of 2 elements [x,y] that all the KLIP subtracted images will be centered on for image
                        registration
        annuli_spacing: how to distribute the annuli radially. Currently three options. Constant (equally spaced), 
                        log (logarithmical expansion with r), and linear (linearly expansion with r)
        maxnumbasis: if not None, maximum number of KL basis/correlated PSFs to use for KLIP. Otherwise, use max(numbasis)
        corr_smooth (float): size of sigma of Gaussian smoothing kernel (in pixels) when computing most correlated PSFs. If 0, no smoothing
        spectrum:       (only applicable for SDI) if not None, optimizes the choice of the reference PSFs based on the
                        spectrum shape.
                        - an array: of length N with the flux of the template spectrum at each wavelength.
                        - a string: Currently only supports "methane" between 1 and 10 microns.
                        Uses minmove to determine the separation from the center of the segment to determine contamination and
                        the size of the PSF (TODO: make PSF size another quantity)
                        (e.g. minmove=3, checks how much contamination is within 3 pixels of the hypothetical source)
                        if smaller than 10%, (hard coded quantity), then use it for reference PSF
        psf_library:    if not None, a rdi.PSFLibrary object with a PSF Library for RDI
        highpass:       if True, run a Gaussian high pass filter (default size is sigma=imgsize/10)
                            can also be a number specifying FWHM of box in pixel units

        lite:           if True, run a low memory version of the alogirhtm

        save_aligned:	Save the aligned and scaled images (as well as various wcs information), True/False
        restore_aligned: The aligned and scaled images from a previous run of klip_dataset
        				(usually restored_aligned = dataset.aligned_and_scaled)
        dtype:          data type of the arrays. Should be either ctypes.c_float(default) or ctypes.c_double
        algo (str):     algorithm to use ('klip', 'nmf', 'empca', 'none'). None will run no PSF subtraction. 
        time_collapse:  how to collapse the data in time. Currently support: "mean", "weighted-mean", 'median', "weighted-median"
        wv_collapse:    how to collapse the data in wavelength. Currently support: 'median', 'mean', 'trimmed-mean'
        verbose (bool): if True, print KLIP processes. 

    Returns
        Saved files in the output directory
        Returns: nothing, but saves to dataset.output: (b, N, wv, y, x) 5D cube of KL cutoff modes (b), number of images
                            (N), wavelengths (wv), and spatial dimensions. Images are derotated.
                            For ADI only, the wv is omitted so only 4D cube
    """
    ######### Check inputs ##########

    # empca currently does not support movement or minrot
    if algo.lower() == 'empca' and (minrot != 0 or movement != 0):
        raise ValueError('empca currently does not support movement, minrot selection criteria, '
                         'must be set to 0')
    elif algo.lower() == 'none':
        # remove some psfsubtraction params
        movement = 0
        minmove = 0
        numbasis = [1]

    # defaullt numbasis if none
    if numbasis is None:
        totalimgs = dataset.input.shape[0]
        maxbasis = np.min([totalimgs, 100]) # only going up to 100 KL modes by default
        numbasis = np.arange(1, maxbasis + 5, 10)
        if verbose is True:
            print("KL basis not specified. Using default.", numbasis)
    else:
        if hasattr(numbasis, "__len__"):
            numbasis = np.array(numbasis)
        else:
            numbasis = np.array([numbasis])


    time_collapse = time_collapse.lower()
    weighted = "weighted" in time_collapse # boolean whether to use weights

    if corr_smooth < 0:
        raise ValueError("corr_smooth needs be non-negative. Supplied value is {0}".format(corr_smooth))

    # RDI Sanity Checks to make sure PSF Library is properly configured
    if "RDI" in mode:
        if lite:
            raise ValueError("RDI is currently not supported in memory lite mode.")
        if psf_library is None:
            raise ValueError("You need to pass in a psf_library if you want to run RDI")
        if psf_library.dataset is dataset:
            raise ValueError("The PSF Library is not prepared for this dataset. Run psf_library.prepare_library()")
        if aligned_center is not None:
            if not np.array_equal(aligned_center, psf_library.aligned_center): 
                raise ValueError("The images need to be aligned to the same center as the RDI Library")

        else:
            aligned_center = psf_library.aligned_center
        # good rdi_library
        master_library = psf_library.master_library
        rdi_corr_matrix = psf_library.correlation
        rdi_good_psfs = psf_library.isgoodpsf
    else:
        master_library = None
        rdi_corr_matrix = None
        rdi_good_psfs = None

    ######### End chcking inputs ########

    # Save the WCS and centers info, incase we need it again!
    if save_aligned is True:
        dataset.old_wcs = copy.deepcopy(dataset.wcs)
        dataset.old_centers = copy.deepcopy(dataset.centers)
        dataset.old_PAs = copy.deepcopy(dataset.PAs)

    # select memory-lite version if requested
    # If lite, we are not running iwth save_aligned and restore_aligned keywords. Throw error if this happens
    if lite:
        klip_function = klip_parallelized_lite
        if (save_aligned is True) or (restored_aligned is True):
            raise ValueError('save_aligned and restored_aligned are not compatible with lite mode')
        # save_aligned = False
        # restored_aligned = None
    else:
        klip_function = klip_parallelized
    
    # If re-running KLIP with same data, restore centers to old value
    # We don't need to highpass the data if the aligned and scaled images are being restored
    if restored_aligned is not None:
        dataset.centers = copy.deepcopy(dataset.old_centers)
        dataset.wcs = copy.deepcopy(dataset.old_wcs)
        dataset.PAs = copy.deepcopy(dataset.old_PAs)
    else:
        if isinstance(highpass, bool):
            if highpass:
                dataset.input = high_pass_filter_imgs(dataset.input, numthreads=numthreads)
        else:
            # should be a number
            if isinstance(highpass, (float, int)):
                highpass = float(highpass)
                fourier_sigma_size = (dataset.input.shape[1]/(highpass)) / (2*np.sqrt(2*np.log(2)))
                dataset.input = high_pass_filter_imgs(dataset.input, numthreads=numthreads, filtersize=fourier_sigma_size)

    # if no outputdir specified, then current working directory (don't want to write to '/'!)
    if outputdir == "":
        outputdir = "."

    if spectrum is not None:
        if isinstance(spectrum,np.ndarray):
            spectrum_name = "custom"
            if np.size(spectrum) == np.size(dataset.wvs):
                spectra_template = spectrum
            else:
                raise ValueError("{0} is not a valid spectral template. Length of spectrum must be {1}."
                                 .format(spectrum,np.size(dataset.wvs)))
        if isinstance(spectrum,str):
            spectrum_name = spectrum
            if spectrum.lower() == "methane":
                pykliproot = os.path.dirname(os.path.realpath(__file__))
                spectrum_dat = np.loadtxt(os.path.join(pykliproot,"spectra","t800g100nc.flx"))[:160] #skip wavelegnths longer of 10 microns
                spectrum_wvs = spectrum_dat[:,1]
                spectrum_fluxes = spectrum_dat[:,3]
                spectrum_interpolation = interp.interp1d(spectrum_wvs, spectrum_fluxes, kind='cubic')

                spectra_template = spectrum_interpolation(dataset.wvs)
            else:
                raise ValueError("{0} is not a valid spectral template. Only currently supporting 'methane'"
                                 .format(spectrum))
    else:
        spectra_template = None
        spectrum_name = None

    # save klip parameters as a string
    maxbasis_str = maxnumbasis if maxnumbasis is not None else np.max(numbasis) # prefer to use maxnumbasis if possible
    klipparams = "mode={mode},annuli={annuli},subsect={subsections},minmove={movement}," \
                 "numbasis={numbasis}/{maxbasis},minrot={minrot},calibflux={calibrate_flux},spectrum={spectrum}," \
                 "highpass={highpass}, time_collapse={weighted}".format(mode=mode, annuli=annuli, 
                                              subsections=subsections, movement=movement,
                                              numbasis="{numbasis}", maxbasis=maxbasis_str, minrot=minrot,
                                              calibrate_flux=calibrate_flux, spectrum=spectrum_name, highpass=highpass,
                                              weighted=time_collapse)
    dataset.klipparams = klipparams

    # set all the klip_parallelized.py args here
    pyklip_args = {'OWA':dataset.OWA, 'mode':mode, 'annuli':annuli, 'subsections':subsections, 'movement':movement,
                    'numbasis':numbasis, 'numthreads':numthreads, 'minrot':minrot, 'aligned_center':aligned_center,
                    'annuli_spacing':annuli_spacing, 'maxnumbasis':maxnumbasis, 'corr_smooth':corr_smooth,
                    'spectrum':spectra_template, 'psf_library':master_library,
                    'psf_library_corr':rdi_corr_matrix, 'psf_library_good':rdi_good_psfs,
                    'save_aligned' : save_aligned, 'restored_aligned' : restored_aligned, 'dtype':dtype,
                    'algo':algo, 'compute_noise_cube':weighted, 'verbose':verbose}

    #Set MLK parameters
    if mkl_exists:
        old_mkl = mkl.get_max_threads()
        mkl.set_num_threads(1)

    # some bookkeeping
    unique_wvs = np.unique(dataset.wvs)
    num_wvs = int(np.size(unique_wvs))
    number_of_klmodes = np.size(numbasis)
    num_cubes = np.size(dataset.wvs) // num_wvs

    # run KLIP
    # For SDI(+ADI)(+RDI) reductions
    if "SDI" in mode:
        if verbose is True:
            print("Beginning {0} KLIP".format(mode))

        # Actually run the PSF Subtraction with all the arguments
        klip_outputs = klip_function(dataset.input, dataset.centers, dataset.PAs, dataset.wvs, dataset.filenums,
                                     dataset.IWA, **pyklip_args)

        # parse the output of klip. Normally, it is just the klipped_imgs,
        # but some optional arguments return more things
        if not save_aligned:
            klipped_imgs, klipped_center, stddev_frames = klip_outputs
        else:
            klipped_imgs, klipped_center, stddev_frames, dataset.aligned_and_scaled = klip_outputs

        # save output and image center for each output
        dataset.output = klipped_imgs
        dataset.output_centers = np.array([klipped_center for _ in range(klipped_imgs.shape[1])])
        # construct the output wcs info, but it's currently just a copy of the input one until we rotate it
        dataset.output_wcs = np.array([w.deepcopy() if w is not None else None for w in dataset.wcs])

    # For ADI only datasets, can run KLIP on each wavelenght separately
    else:
        # set up output, output centers, and output wcs variables but they are the same as the input for now
        dataset.output_centers = np.copy(dataset.centers)
        dataset.output_wcs = np.array([w.deepcopy() if w is not None else None for w in dataset.wcs])

        # append output to a list at first since we are running it a bunch of times
        dataset.output = []
        stddev_frames = []

        # save the output of aligned_and_scaled optionally in a list
        if save_aligned:
            dataset.aligned_and_scaled = []

        for wvindex,unique_wv in enumerate(unique_wvs):
            if (num_wvs > 1) and (verbose is True):
                print("Running KLIP ADI on slice {0}/{1}: {2:.3f} um".format(wvindex+1, num_wvs, unique_wv))
            thiswv = np.where(dataset.wvs == unique_wv)

            if restored_aligned is not None:
                restored_aligned_thiswv = restored_aligned[np.where(unique_wv == unique_wvs)]
            else:
                restored_aligned_thiswv = None

            klip_output = klip_function(dataset.input[thiswv], dataset.centers[thiswv], dataset.PAs[thiswv], dataset.wvs[thiswv],
                                        dataset.filenums[thiswv],
                                        dataset.IWA, **pyklip_args)
            
            klipped_imgs = klip_output[0]
            klipped_center = klip_output[1]
            noise_frames = klip_output[2]
            
            # save data for this wavelength
            dataset.output.append(klipped_imgs)
            dataset.output_centers[thiswv[0], 0] = klipped_center[0]
            dataset.output_centers[thiswv[0], 1] = klipped_center[1]
            stddev_frames.append(noise_frames)

            # if we need to save the aligned_and_scaled images, put them in a list
            if save_aligned:
                dataset.aligned_and_scaled.append(klip_output[-1][0])

        # convert lists to numpy arrays
        dataset.output = np.array(dataset.output)
        stddev_frames = np.array(stddev_frames)

        if save_aligned:
            dataset.aligned_and_scaled = np.array(dataset.aligned_and_scaled)

        # reformat the output to be consistent with the other modes.
        # Currently shape is (wv,b,N,y,x). Switch to (b,N,wv,y,x), then flatten in wavelength dimension
        dataset.output = np.swapaxes(dataset.output, 0, 1) # shape of (b, wv, N, y, x)
        dataset.output = np.swapaxes(dataset.output, 1, 2) # shape of (b, N, wv, y, x)
        # then collapse N/wv together
        dataset.output = np.reshape(dataset.output, (dataset.output.shape[0], dataset.output.shape[1]*dataset.output.shape[2], dataset.output.shape[3], dataset.output.shape[4]) )

        if weighted:
            # do the same with the stddev frames
            stddev_frames = np.swapaxes(stddev_frames, 0, 1) # shape of (b, wv, N, y, x)
            stddev_frames = np.swapaxes(stddev_frames, 1, 2) # shape of (b, N, wv, y, x)
            # then collapse N/wv together
            stddev_frames = np.reshape(stddev_frames, (stddev_frames.shape[0], stddev_frames.shape[1]*stddev_frames.shape[2], stddev_frames.shape[3], stddev_frames.shape[4]) )
        else:
            stddev_frames = 1

    # TODO: handling of only a single numbasis
    # derotate all the images
    # flatten so it's just a 3D array (collapse KL and Nframes dimensions)
    oldshape = dataset.output.shape
    dataset.output = dataset.output.reshape(oldshape[0]*oldshape[1], oldshape[2], oldshape[3])
    if weighted:
        # do the same for the stddev frames
        stddev_frames = stddev_frames.reshape(oldshape[0]*oldshape[1], oldshape[2], oldshape[3])

    # we need to duplicate PAs and centers for the different KL mode cutoffs we supplied
    flattend_parangs = np.tile(dataset.PAs, oldshape[0])
    flattened_centers = np.tile(dataset.output_centers.reshape(oldshape[1]*2), oldshape[0]).reshape(oldshape[1]*oldshape[0],2)

    # align center to center of image if not specified
    # note that klip_parallelized aligns everything to the mean of the input centers, whereas now we will re align it
    # to the middle of the array for cosmetic purposes. 
    if aligned_center is None:
        aligned_center = [int(dataset.input.shape[2]//2), int(dataset.input.shape[1]//2)]

    # parallelized rotate images
    if verbose is True:
        print("Derotating Images...")
    rot_imgs = rotate_imgs(dataset.output, flattend_parangs, flattened_centers, numthreads=numthreads, flipx=dataset.flipx,
                           hdrs=dataset.output_wcs, new_center=aligned_center)
    # re-expand the images in num cubes/num wvs (num KLmode cutoffs, num cubes, num wvs, y, x)
    rot_imgs = rot_imgs.reshape(oldshape[0], oldshape[1]//num_wvs, num_wvs, oldshape[2], oldshape[3])

    # rotate the weights too if necessary
    if weighted:
        stddev_frames = rotate_imgs(stddev_frames, flattend_parangs, flattened_centers, numthreads=numthreads, flipx=dataset.flipx, new_center=aligned_center)
        stddev_frames = stddev_frames.reshape(oldshape[0], oldshape[1]//num_wvs, num_wvs, oldshape[2], oldshape[3])

    # save modified data and centers
    dataset.output = rot_imgs
    dataset.output_centers[:,0] = aligned_center[0]
    dataset.output_centers[:,1] = aligned_center[1]

   
    # valid output path and write iamges
    outputdirpath = os.path.realpath(outputdir)
    if verbose is True:
        print("Writing Images to directory {0}".format(outputdirpath))

    # create weights for each pixel. If we aren't doing weighted mean, weights are just ones
    pixel_weights = 1./stddev_frames**2

    if num_wvs > 1:
        _save_spectral_cubes(dataset, pixel_weights, time_collapse, numbasis, calibrate_flux, outputdirpath, fileprefix)

    _save_wv_collapsed_images(dataset, pixel_weights, numbasis, time_collapse, wv_collapse, num_wvs, spectrum,
                              spectra_template, calibrate_flux, outputdirpath, fileprefix, verbose)

    # Restore old setting
    if mkl_exists:
        mkl.set_num_threads(old_mkl)

    return
