"""Classes to compute cluster centers with different algorithms.

This file contains the generic centering class and various different
implementations of different centering algorithms.  The main algorithm is
CenteringWcenZred, but others are included as part of the centering training
procedure.
"""
import fitsio
import esutil
import numpy as np

from .utilities import gaussFunction
from .utilities import interpol

class Centering(object):
    """
    Generic centering base class for computing cluster centers.
    """

    def __init__(self, cluster, zlambda_corr=None):
        """
        Instantiate a Centering object for a cluster.

        Parameters
        ----------
        cluster: `redmapper.Cluster`
           Cluster to compute centering
        zlambda_corr: `redmapper.ZlambdaCorrectionPar`, optional
           z_lambda correction parameters, if desired.  Default is None.
        """
        # Reference to the cluster; may need to copy
        self.cluster = cluster

        # And the zlambda_corr structure
        self.zlambda_corr = zlambda_corr

        # For convenience, make references to these structures
        self.zredstr = cluster.zredstr
        self.config = cluster.config
        self.cosmo = cluster.cosmo

        # Reset values
        self.ra = np.zeros(self.config.percolation_maxcen) - 400.0
        self.dec = np.zeros(self.config.percolation_maxcen) - 400.0
        self.ngood = 0
        self.index = np.zeros(self.config.percolation_maxcen, dtype=np.int32) - 1
        self.maxind = -1
        self.lnlamlike = -1.0
        self.lnbcglike = -1.0
        self.p_cen = np.zeros(self.config.percolation_maxcen)
        self.q_cen = np.zeros(self.config.percolation_maxcen)
        self.p_fg = np.zeros(self.config.percolation_maxcen)
        self.q_miss = 0.0
        self.p_sat = np.zeros(self.config.percolation_maxcen)
        self.p_c = np.zeros(self.config.percolation_maxcen)

    def find_center(self):
        """
        Stub to override to find center
        """
        return False

class CenteringBCG(Centering):
    """
    Centering class using the brightest cluster galaxy (BCG) algorithm.
    """

    def find_center(self):
        """
        Find the center using the CenteringBCG algorithm.

        This algorithm takes the brightest member with pmem > 0.8 and calls it
        the central galaxy.

        Will set self.maxind (index of best center); self.ra, self.dec
        (position of best center); self.ngood (number of good candidates);
        self.index[:] (indices of all the candidates); self.p_cen[:] (pcen
        centering probabilities); self.q_cen[:] (qcen unused miss
        probabilities); self.p_sat[:] (p_sat satellite probabilities).

        Returns
        -------
        success: `bool`
           True when a center is successfully found. (Always True).
        """
        # This is somewhat arbitrary, and is not yet configurable
        pmem_cut = 0.8

        use, = np.where((self.cluster.neighbors.r < self.cluster.r_lambda) &
                        ((self.cluster.neighbors.pmem > pmem_cut) |
                         (np.abs(self.cluster.neighbors.zred - self.cluster.redshift) < 2.0 * self.cluster.neighbors.zred_e)))

        if use.size == 0:
            return False

        mind = np.argmin(self.cluster.neighbors.refmag[use])

        self.maxind = use[mind]
        self.ra = np.array([self.cluster.neighbors.ra[self.maxind]])
        self.dec = np.array([self.cluster.neighbors.dec[self.maxind]])
        self.ngood = 1
        self.index[0] = self.maxind
        self.p_cen[0] = 1.0
        self.q_cen[0] = 1.0
        self.p_sat[0] = 0.0

        return True

class CenteringWcenZred(Centering):
    """
    Centering class using the "wcen-zred" algorithm.

    This algorithm computes the primary centering likelihood algorithm by
    computing the connectivity of the members, as well as ensuring
    consistency between zred of the candidates and the cluster redshift.
    """

    def find_center(self):
        """
        Find the center using the CenteringWcenZred algorithm.

        This algorithm computes the primary centering likelihood algorithm by
        computing the connectivity of the members, as well as ensuring
        consistency between zred of the candidates and the cluster redshift.

        Will set self.maxind (index of best center); self.ra, self.dec
        (position of best center); self.ngood (number of good candidates);
        self.index[:] (indices of all the candidates); self.p_cen[:] (pcen
        centering probabilities); self.q_cen[:] (qcen unused miss
        probabilities); self.p_sat[:] (p_sat satellite probabilities).

        Returns
        -------
        success: `bool`
           True when a center is successfully found. (Always True).
        """
        # These are the galaxies considered as candidate centers
        use, = np.where((self.cluster.neighbors.r < self.cluster.r_lambda) &
                        (self.cluster.neighbors.pfree >= self.config.percolation_pbcg_cut) &
                        (self.cluster.neighbors.zred_chisq < self.config.wcen_zred_chisq_max) &
                        ((self.cluster.neighbors.pmem > 0.0) |
                         (np.abs(self.cluster.redshift - self.cluster.neighbors.zred) < 5.0 * self.cluster.neighbors.zred_e)))

        # Do the phi_cen filter
        mbar = self.cluster.mstar + self.config.wcen_Delta0 + self.config.wcen_Delta1 * np.log(self.cluster.Lambda / self.config.wcen_pivot)
        phi_cen = gaussFunction(self.cluster.neighbors.refmag[use],
                                1. / (np.sqrt(2. * np.pi) * self.config.wcen_sigma_m),
                                mbar,
                                self.config.wcen_sigma_m)

        if self.zlambda_corr is not None:
            zrmod = interpol(self.zlambda_corr.zred_uncorr, self.zlambda_corr.z, self.cluster.redshift)
            gz = gaussFunction(self.cluster.neighbors.zred[use],
                               1. / (np.sqrt(2. * np.pi) * self.cluster.neighbors.zred_e[use]),
                               zrmod,
                               self.cluster.neighbors.zred_e[use])
        else:
            gz = gaussFunction(self.cluster.neighbors.zred[use],
                               1. / (np.sqrt(2. * np.pi) * self.cluster.neighbors.zred_e[use]),
                               self.cluster.redshift,
                               self.cluster.neighbors.zred_e[use])

        # and the w filter.  We need w for each galaxy that is considered a candidate center.
        # Note that in order to calculate w we need to know all the galaxies that are
        # around it, but only within r_lambda *of that galaxy*.  This is tricky.

        u, = np.where(self.cluster.neighbors.p > 0.0)

        # This is the maximum radius in units of degrees (r_lambda is Mpc; mpc_scale is Mpc / degree)
        maxrad = 1.1 * self.cluster.r_lambda / self.cluster.mpc_scale

        htm_matcher = esutil.htm.Matcher(self.cluster.neighbors.depth,
                                         self.cluster.neighbors.ra[use],
                                         self.cluster.neighbors.dec[use])
        i2, i1, dist = htm_matcher.match(self.cluster.neighbors.ra[u],
                                         self.cluster.neighbors.dec[u],
                                         maxrad, maxmatch=0)

        subdifferent, = np.where(~(use[i1] == u[i2]))
        i1 = i1[subdifferent]
        i2 = i2[subdifferent]
        pdis = dist[subdifferent] * self.cluster.mpc_scale
        pdis = np.sqrt(pdis**2. + self.config.wcen_rsoft**2.)

        lum = 10.**((self.cluster.mstar - self.cluster.neighbors.refmag) / (2.5))

        # Put a floor on w when we have a strange candidate at the edge that doesn't
        # match any good galaxies
        w = np.zeros(use.size) + 1e-3
        for i in range(use.size):
            # need to filter on r_lambda...
            subgal, = np.where(i1 == i)
            if subgal.size > 0:
                inside, = np.where(pdis[subgal] < self.cluster.r_lambda)
                if inside.size > 0:
                    indices = u[i2[subgal[inside]]]
                    if self.config.wcen_uselum:
                        w[i] = np.log(np.sum(self.cluster.neighbors.p[indices] * lum[indices] /
                                             pdis[subgal[inside]]) /
                                      ((1. / self.cluster.r_lambda) *
                                       np.sum(self.cluster.neighbors.p[indices] * lum[indices])))
                    else:
                        w[i] = np.log(np.sum(self.cluster.neighbors.p[indices] /
                                             pdis[subgal[inside]]) /
                                      ((1. / self.cluster.r_lambda) *
                                       np.sum(self.cluster.neighbors.p[indices])))

        sigscale = np.sqrt((np.clip(self.cluster.Lambda, None, self.config.wcen_maxlambda) / self.cluster.scaleval) / self.config.wcen_pivot)

        # scale with richness for Poisson errors
        sig = self.config.lnw_cen_sigma / sigscale

        fw = gaussFunction(np.log(w),
                           1. / (np.sqrt(2. * np.pi) * sig),
                           self.config.lnw_cen_mean,
                           sig)

        ucen = phi_cen * gz * fw

        lo, = np.where(ucen < 1e-10)
        ucen[lo] = 0.0

        # and the satellite function
        maxmag = self.cluster.mstar - 2.5 * np.log10(self.config.lval_reference)
        phi_sat = self.cluster._calc_luminosity(maxmag, idx=use)

        satsig = self.config.lnw_sat_sigma / sigscale
        fsat = gaussFunction(np.log(w),
                             1. / (np.sqrt(2. * np.pi) * satsig),
                             self.config.lnw_sat_mean,
                             satsig)

        usat = phi_sat * gz * fsat

        lo, = np.where(usat < 1e-10)
        usat[lo] = 0.0

        # and the background/foreground
        fgsig = self.config.lnw_fg_sigma / sigscale
        ffg = gaussFunction(np.log(w),
                            1. / (np.sqrt(2. * np.pi) * fgsig),
                            self.config.lnw_fg_mean,
                            fgsig)

        # we want to divide out the r, and we don't want small r's messing this up
        rtest = np.zeros(use.size) + 0.1

        bcounts = ffg * (self.cluster.calc_zred_bkg_density(rtest,
                                                            self.cluster.neighbors.zred[use],
                                                            self.cluster.neighbors.refmag[use]) /
                         (2. * np.pi * rtest)) * np.pi * self.cluster.r_lambda**2.

        # The start of Pcen
        Pcen_basic = np.clip(self.cluster.neighbors.pfree[use] * (ucen / (ucen + (self.cluster.Lambda / self.cluster.scaleval - 1.0) * usat + bcounts)),None, 0.99999)

        # make sure we don't have any bad values
        bad, = np.where(~np.isfinite(Pcen_basic))
        Pcen_basic[bad] = 0.0

        okay, = np.where(Pcen_basic > 0.0)
        if okay.size == 0:
            # There are literally NO centers
            self.q_miss = 1.0

            # Set the same as the input galaxy...
            # We need this to be an array of length 1
            good = np.atleast_1d(np.argmin(self.cluster.neighbors.r[use]))

            maxind = use[good[0]]

            Pcen = np.zeros(use.size)
            Qcen = np.zeros(use.size)

        else:
            # Do the renormalization

            Pcen_unnorm = np.zeros(use.size)

            # Only consider centrals that have a non-zero probability
            ok, = np.where(Pcen_basic > 0)

            st = np.argsort(Pcen_basic[ok])[::-1]
            if st.size < self.config.percolation_maxcen:
                good = ok[st]
            else:
                good = ok[st[0: self.config.percolation_maxcen]]

            self.ngood = good.size

            for i in range(self.ngood):
                Pcen0 = Pcen_basic[good[i]]
                Pcen_basic[good[i]] = 0.0
                Pcen_unnorm[good[i]] = Pcen0 * np.prod(1.0 - Pcen_basic[good])
                Pcen_basic[good[i]] = Pcen0

            Qmiss = np.prod(1.0 - Pcen_basic[good])

            KQ = 1./(Qmiss + np.sum(Pcen_unnorm))
            KP = 1./np.sum(Pcen_unnorm)

            Pcen = KP * Pcen_unnorm
            Qcen = KQ * Pcen_unnorm

            mod1 = np.sum(np.log(ucen[good] + (self.cluster.Lambda - 1) * usat[good] + bcounts[good]))
            mod2 = np.sum(np.log(self.cluster.Lambda * usat[good] + bcounts[good]))

            # A new statistic that doesn't quite work
            Qmiss = -2.0 * np.sum(np.log((ucen[good] + (self.cluster.Lambda - 1) * usat[good] + bcounts[good]) / (self.cluster.Lambda * usat[good] + bcounts[good])))

            maxind = use[good[0]]

        Pfg_basic = bcounts[good] / ((self.cluster.Lambda - 1.0) * usat[good] + bcounts[good])
        inf, = np.where(~np.isfinite(Pfg_basic))
        Pfg_basic[inf] = 0.0

        Pfg = (1.0 - Pcen[good]) * Pfg_basic

        Psat_basic = (self.cluster.Lambda - 1.0) * usat[good] / ((self.cluster.Lambda - 1.0) * usat[good] + bcounts[good])
        inf, = np.where(~np.isfinite(Psat_basic))
        Psat_basic[inf] = 0.0

        Psat = (1.0 - Pcen[good]) * Psat_basic

        self.ra[0: good.size] = self.cluster.neighbors.ra[use[good]]
        self.dec[0: good.size] = self.cluster.neighbors.dec[use[good]]
        self.maxind = use[good[0]]
        self.index[0: good.size] = use[good]
        self.p_cen[0: good.size] = Pcen[good]
        self.q_cen[0: good.size] = Qcen[good]
        self.p_fg[0: good.size] = Pfg
        self.p_sat[0: good.size] = Psat
        self.p_c[0: good.size] = Pcen_basic[good]

        return True

class CenteringRandom(Centering):
    """
    Centering class using the random-position algorithm.

    This is used for filter calibration.
    """

    def find_center(self):
        """
        Find the center using the CenteringRandom algorithm.

        This algorithm takes a random position within cluster r_lambda and
        calls it the center.  It is not a very good centering algorithm.

        Will set self.maxind (index of best center); self.ra, self.dec
        (position of best center); self.ngood (number of good candidates);
        self.index[:] (indices of all the candidates); self.p_cen[:] (pcen
        centering probabilities); self.q_cen[:] (qcen unused miss
        probabilities); self.p_sat[:] (p_sat satellite probabilities).

        Returns
        -------
        success: `bool`
           True when a center is successfully found. (Always True).
        """
        r = self.cluster.r_lambda * np.sqrt(np.random.random(size=1))
        phi = 2. * np.pi * np.random.random(size=1)

        x = r * np.cos(phi) / (self.cluster.mpc_scale)
        y = r * np.sin(phi) / (self.cluster.mpc_scale)

        ra_cen = self.cluster.ra + x / np.cos(np.radians(self.cluster.dec))
        dec_cen = self.cluster.dec + y

        self.ra[0] = ra_cen
        self.dec[0] = dec_cen
        self.ngood = 1
        self.index[0] = -1
        self.maxind = -1
        self.p_cen[0] = 1.0
        self.q_cen[0] = 1.0
        self.p_sat[0] = 0.0
        self.p_fg[0] = 0.0
        self.p_c[0] = 1.0

        return True

class CenteringRandomSatellite(Centering):
    """
    Centering class using the random-satellite algorithm.

    This is used for filter calibration.
    """

    def find_center(self):
        """
        Find the center using the CenteringRandomSatellite algorithm.

        This algorithm takes a random member (weighted by member pmem) and
        calls it the center.  It is not a very good centering algorithm (but
        better than pure random!)

        Will set self.maxind (index of best center); self.ra, self.dec
        (position of best center); self.ngood (number of good candidates);
        self.index[:] (indices of all the candidates); self.p_cen[:] (pcen
        centering probabilities); self.q_cen[:] (qcen unused miss
        probabilities); self.p_sat[:] (p_sat satellite probabilities).

        Returns
        -------
        success: `bool`
           True when a center is successfully found. (Always True).
        """
        st = np.argsort(self.cluster.neighbors.pmem)[::-1]

        pdf = self.cluster.neighbors.pmem[st]
        pdf /= np.sum(pdf)
        cdf = np.cumsum(pdf, dtype=np.float64)
        cdfi = (cdf * st.size).astype(np.int32)

        rand = (np.random.uniform(size=1) * st.size).astype(np.int32)
        ind = np.where(cdfi >= rand[0])[0][0]
        maxind = st[ind]

        ra_cen = self.cluster.neighbors.ra[maxind]
        dec_cen = self.cluster.neighbors.dec[maxind]

        self.ra[0] = ra_cen
        self.dec[0] = dec_cen
        self.index[0] = maxind
        self.maxind = maxind
        self.ngood = 1
        self.p_cen[0] = 1.0
        self.q_cen[0] = 1.0
        self.p_sat[0] = 0.0
        self.p_fg[0] = 0.0
        self.p_c[0] = 1.0

        return True


# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Models for doing PSF/PRF fitting photometry on image data.
"""

import copy
import itertools
import warnings

import numpy as np
from astropy.nddata import NDData
from astropy.modeling import Parameter, Fittable2DModel
from astropy.utils.exceptions import (AstropyWarning,
                                      AstropyDeprecationWarning)

__all__ = ['NonNormalizable', 'FittableImageModel', 'EPSFModel',
           'GriddedPSFModel', 'IntegratedGaussianPRF', 'PRFAdapter']

class NonNormalizable(AstropyWarning):
    """
    Used to indicate that a :py:class:`FittableImageModel` model is
    non-normalizable.

    """
    pass

class FittableImageModel(Fittable2DModel):
    """
    A fittable 2D model of an image allowing for image intensity scaling
    and image translations.

    This class takes 2D image data and computes the
    values of the model at arbitrary locations (including at intra-pixel,
    fractional positions) within this image using spline interpolation
    provided by :py:class:`~scipy.interpolate.RectBivariateSpline`.

    The fittable model provided by this class has three model parameters:
    an image intensity scaling factor (`flux`) which is applied to
    (normalized) image, and two positional parameters (`x_0` and `y_0`)
    indicating the location of a feature in the coordinate grid on which
    the model is to be evaluated.

    If this class is initialized with `flux` (intensity scaling factor)
    set to `None`, then `flux` is going to be estimated as ``sum(data)``.

    Parameters
    ----------
    data : numpy.ndarray
        Array containing 2D image.

    origin : tuple, None, optional
        A reference point in the input image ``data`` array. When origin is
        `None`, origin will be set at the middle of the image array.

        If `origin` represents the location of a feature (e.g., the position
        of an intensity peak) in the input ``data``, then model parameters
        `x_0` and `y_0` show the location of this peak in an another target
        image to which this model was fitted. Fundamentally, it is the
        coordinate in the model's image data that should map to
        coordinate (`x_0`, `y_0`) of the output coordinate system on which the
        model is evaluated.

        Alternatively, when `origin` is set to ``(0,0)``, then model parameters
        `x_0` and `y_0` are shifts by which model's image should be translated
        in order to match a target image.

    normalize : bool, optional
        Indicates whether or not the model should be build on normalized
        input image data. If true, then the normalization constant (*N*) is
        computed so that

        .. math::
            N \\cdot C \\cdot \\Sigma_{i,j}D_{i,j} = 1,

        where *N* is the normalization constant, *C* is correction factor
        given by the parameter ``normalization_correction``, and
        :math:`D_{i,j}` are the elements of the input image ``data`` array.

    normalization_correction : float, optional
        A strictly positive number that represents correction that needs to
        be applied to model's data normalization (see *C* in the equation
        in the comments to ``normalize`` for more details).

        A possible application for this parameter is to account for aperture
        correction. Assuming model's data represent a PSF to be fitted to
        some target star, we set ``normalization_correction`` to the aperture
        correction that needs to be applied to the model. That is,
        ``normalization_correction`` in this case should be set to the
        ratio between the total flux of the PSF (including flux outside model's
        data) to the flux of model's data.
        Then, best fitted value of the `flux` model
        parameter will represent an aperture-corrected flux of the target star.

    fill_value : float, optional
        The value to be returned by the `evaluate` or
        ``astropy.modeling.Model.__call__`` methods
        when evaluation is performed outside the definition domain of the
        model.

    ikwargs : dict, optional

        Additional optional keyword arguments to be passed directly to the
        `compute_interpolator` method. See `compute_interpolator` for more
        details.

    """
    flux = Parameter(description='Intensity scaling factor for image data.',
                     default=1.0)
    x_0 = Parameter(description='X-position of a feature in the image in '
                    'the output coordinate grid on which the model is '
                    'evaluated.', default=0.0)
    y_0 = Parameter(description='Y-position of a feature in the image in '
                    'the output coordinate grid on which the model is '
                    'evaluated.', default=0.0)

    def __init__(self, data, flux=flux.default,
                 x_0=x_0.default, y_0=y_0.default,
                 normalize=False, normalization_correction=1.0,
                 origin=None, oversampling=1, fill_value=0.0, ikwargs={}):
        self._fill_value = fill_value
        self._img_norm = None
        self._normalization_status = 0 if normalize else 2
        self._store_interpolator_kwargs(ikwargs)
        self._set_oversampling(oversampling)

        if normalization_correction <= 0:
            raise ValueError("'normalization_correction' must be strictly "
                             "positive.")
        self._normalization_correction = normalization_correction

        self._data = np.array(data, copy=True, dtype=np.float64)

        if not np.all(np.isfinite(self._data)):
            raise ValueError("All elements of input 'data' must be finite.")

        # set input image related parameters:
        self._ny, self._nx = self._data.shape
        self._shape = self._data.shape
        if self._data.size < 1:
            raise ValueError("Image data array cannot be zero-sized.")

        # set the origin of the coordinate system in image's pixel grid:
        self.origin = origin

        if flux is None:
            if self._img_norm is None:
                self._img_norm = self._compute_raw_image_norm(self._data)
            flux = self._img_norm

        self._compute_normalization(normalize)

        super().__init__(flux, x_0, y_0)

        # initialize interpolator:
        self.compute_interpolator(ikwargs)

    def _compute_raw_image_norm(self, data):
        """
        Helper function that computes the uncorrected inverse normalization
        factor of input image data. This quantity is computed as the
        *sum of all pixel values*.

        .. note::
            This function is intended to be overriden in a subclass if one
            desires to change the way the normalization factor is computed.

        """
        return np.sum(self._data, dtype=np.float64)

    def _compute_normalization(self, normalize):
        """
        Helper function that computes (corrected) normalization factor
        of the original image data. This quantity is computed as the
        inverse "raw image norm" (or total "flux" of model's image)
        corrected by the ``normalization_correction``:

        .. math::
            N = 1/(\\Phi * C),

        where :math:`\\Phi` is the "total flux" of model's image as
        computed by `_compute_raw_image_norm` and *C* is the
        normalization correction factor. :math:`\\Phi` is computed only
        once if it has not been previously computed. Otherwise, the
        existing (stored) value of :math:`\\Phi` is not modified as
        :py:class:`FittableImageModel` does not allow image data to be
        modified after the object is created.

        .. note::
            Normally, this function should not be called by the
            end-user. It is intended to be overriden in a subclass if
            one desires to change the way the normalization factor is
            computed.
        """

        self._normalization_constant = 1.0 / self._normalization_correction

        if normalize:
            # compute normalization constant so that
            # N*C*sum(data) = 1:
            if self._img_norm is None:
                self._img_norm = self._compute_raw_image_norm(self._data)

            if self._img_norm != 0.0 and np.isfinite(self._img_norm):
                self._normalization_constant /= self._img_norm
                self._normalization_status = 0

            else:
                self._normalization_constant = 1.0
                self._normalization_status = 1
                warnings.warn("Overflow encountered while computing "
                              "normalization constant. Normalization "
                              "constant will be set to 1.", NonNormalizable)

        else:
            self._normalization_status = 2

    @property
    def oversampling(self):
        """
        The factor by which the stored image is oversampled.  I.e., an input
        to this model is multipled by this factor to yield the index into the
        stored image.
        """
        return self._oversampling

    def _set_oversampling(self, value):
        """
        This is a private method because it's used in the initializer but the
        ``oversampling``
        """
        try:
            value = float(value)
        except ValueError:
            raise ValueError('Oversampling factor must be a scalar')
        if value <= 0:
            raise ValueError('Oversampling factor must be greater than 0')

        self._oversampling = value

    @property
    def data(self):
        """ Get original image data. """
        return self._data

    @property
    def normalized_data(self):
        """ Get normalized and/or intensity-corrected image data. """
        return (self._normalization_constant * self._data)

    @property
    def normalization_constant(self):
        """ Get normalization constant. """
        return self._normalization_constant

    @property
    def normalization_status(self):
        """
        Get normalization status. Possible status values are:

        - 0: **Performed**. Model has been successfuly normalized at
          user's request.
        - 1: **Failed**. Attempt to normalize has failed.
        - 2: **NotRequested**. User did not request model to be normalized.

        """
        return self._normalization_status

    @property
    def normalization_correction(self):
        """
        Set/Get flux correction factor.

        .. note::
            When setting correction factor, model's flux will be adjusted
            accordingly such that if this model was a good fit to some target
            image before, then it will remain a good fit after correction
            factor change.

        """
        return self._normalization_correction

    @normalization_correction.setter
    def normalization_correction(self, normalization_correction):
        old_cf = self._normalization_correction
        self._normalization_correction = normalization_correction
        self._compute_normalization(normalize=self._normalization_status != 2)

        # adjust model's flux so that if this model was a good fit to some
        # target image, then it will remain a good fit after correction factor
        # change:
        self.flux *= normalization_correction / old_cf

    @property
    def shape(self):
        """A tuple of dimensions of the data array in numpy style (ny, nx)."""
        return self._shape

    @property
    def nx(self):
        """Number of columns in the data array."""
        return self._nx

    @property
    def ny(self):
        """Number of rows in the data array."""
        return self._ny

    @property
    def origin(self):
        """
        A tuple of ``x`` and ``y`` coordinates of the origin of the coordinate
        system in terms of pixels of model's image.

        When setting the coordinate system origin, a tuple of two `int` or
        `float` may be used. If origin is set to `None`, the origin of the
        coordinate system will be set to the middle of the data array
        (``(npix-1)/2.0``).

        .. warning::
            Modifying `origin` will not adjust (modify) model's parameters
            `x_0` and `y_0`.
        """
        return (self._x_origin, self._y_origin)

    @origin.setter
    def origin(self, origin):
        if origin is None:
            self._x_origin = (self._nx - 1) / 2.0
            self._y_origin = (self._ny - 1) / 2.0
        elif hasattr(origin, '__iter__') and len(origin) == 2:
            self._x_origin, self._y_origin = origin
        else:
            raise TypeError("Parameter 'origin' must be either None or an "
                            "iterable with two elements.")

    @property
    def x_origin(self):
        """X-coordinate of the origin of the coordinate system."""
        return self._x_origin

    @property
    def y_origin(self):
        """Y-coordinate of the origin of the coordinate system."""
        return self._y_origin

    @property
    def fill_value(self):
        """Fill value to be returned for coordinates outside of the domain of
        definition of the interpolator. If ``fill_value`` is `None`, then
        values outside of the domain of definition are the ones returned
        by the interpolator.

        """
        return self._fill_value

    @fill_value.setter
    def fill_value(self, fill_value):
        self._fill_value = fill_value

    def _store_interpolator_kwargs(self, ikwargs):
        """
        This function should be called in a subclass whenever model's
        interpolator is (re-)computed.
        """
        self._interpolator_kwargs = copy.deepcopy(ikwargs)

    @property
    def interpolator_kwargs(self):
        """
        Get current interpolator's arguments used when interpolator was
        created.
        """
        return self._interpolator_kwargs

    def compute_interpolator(self, ikwargs={}):
        """
        Compute/define the interpolating spline. This function can be overriden
        in a subclass to define custom interpolators.

        Parameters
        ----------
        ikwargs : dict, optional

            Additional optional keyword arguments. Possible values are:

            - **degree** : int, tuple, optional
                Degree of the interpolating spline. A tuple can be used to
                provide different degrees for the X- and Y-axes.
                Default value is degree=3.

            - **s** : float, optional
                Non-negative smoothing factor. Default value s=0 corresponds to
                interpolation.
                See :py:class:`~scipy.interpolate.RectBivariateSpline` for more
                details.

        Notes
        -----
            * When subclassing :py:class:`FittableImageModel` for the
              purpose of overriding :py:func:`compute_interpolator`,
              the :py:func:`evaluate` may need to overriden as well depending
              on the behavior of the new interpolator. In addition, for
              improved future compatibility, make sure
              that the overriding method stores keyword arguments ``ikwargs``
              by calling ``_store_interpolator_kwargs`` method.

            * Use caution when modifying interpolator's degree or smoothness in
              a computationally intensive part of the code as it may decrease
              code performance due to the need to recompute interpolator.

        """
        from scipy.interpolate import RectBivariateSpline

        if 'degree' in ikwargs:
            degree = ikwargs['degree']
            if hasattr(degree, '__iter__') and len(degree) == 2:
                degx = int(degree[0])
                degy = int(degree[1])
            else:
                degx = int(degree)
                degy = int(degree)
            if degx < 0 or degy < 0:
                raise ValueError("Interpolator degree must be a non-negative "
                                 "integer")
        else:
            degx = 3
            degy = 3

        if 's' in ikwargs:
            smoothness = ikwargs['s']
        else:
            smoothness = 0

        x = np.arange(self._nx, dtype=np.float)
        y = np.arange(self._ny, dtype=np.float)
        self.interpolator = RectBivariateSpline(
            x, y, self._data.T, kx=degx, ky=degy, s=smoothness
        )

        self._store_interpolator_kwargs(ikwargs)

    def evaluate(self, x, y, flux, x_0, y_0, use_oversampling=True):
        """
        Evaluate the model on some input variables and provided model
        parameters.

        Parameters
        ----------
        use_oversampling : bool, optional
            Whether to use the oversampling factor to calculate the
            model pixel indices.  The default is `True`, which means the
            input indices will be multipled by this factor.
        """

        if use_oversampling:
            xi = self._oversampling * (np.asarray(x) - x_0)
            yi = self._oversampling * (np.asarray(y) - y_0)
        else:
            xi = np.asarray(x) - x_0
            yi = np.asarray(y) - y_0

        xi += self._x_origin
        yi += self._y_origin

        f = flux * self._normalization_constant
        evaluated_model = f * self.interpolator.ev(xi, yi)

        if self._fill_value is not None:
            # find indices of pixels that are outside the input pixel grid and
            # set these pixels to the 'fill_value':
            invalid = (((xi < 0) | (xi > self._nx - 1)) |
                       ((yi < 0) | (yi > self._ny - 1)))
            evaluated_model[invalid] = self._fill_value

        return evaluated_model

class EPSFModel(FittableImageModel):
    """
    A subclass of `FittableImageModel`.

    Parameters
    ----------
    pixel_scale : float, tuple of two floats or `None`, optional
        .. warning::

            The ``pixel_scale`` keyword is now deprecated (since v0.6)
            and will likely be removed in v0.7.  Use the
            ``oversampling`` keyword instead.

        The pixel scale (in arbitrary units) of the ePSF.  The
        ``pixel_scale`` can either be a single float or tuple of two
        floats of the form ``(x_pixscale, y_pixscale)``.  If
        ``pixel_scale`` is a scalar then the pixel scale will be the
        same for both the x and y axes.  The default is `None`, which
        means it will be set to the inverse of the ``oversampling``
        factor.

        The ePSF ``pixel_scale`` is used only when building the ePSF and
        when fitting the ePSF to `Star` objects with `EPSFFitter`.  In
        those cases, the ``pixel_scale`` is used in conjunction with the
        `Star` pixel scale when building and fitting the ePSF.  This
        allows for building (and fitting) a ePSF using images of stars
        with different pixel scales (e.g. velocity aberrations).  The
        ``oversampling`` factor is ignored in these cases.

        If you are not using `EPSFBuilder` or `EPSFFitter`, then you
        must set the ``oversampling`` factor.  The ``pixel_scale`` will
        be ignored.
    """

    def __init__(self, data, flux=1.0, x_0=0, y_0=0, normalize=True,
                 normalization_correction=1.0, origin=None, oversampling=1.,
                 pixel_scale=None, fill_value=0., ikwargs={}):

        if pixel_scale is None:
            pixel_scale = 1. / oversampling
        else:
            warnings.warn('The pixel_scale keyword is deprecated and will '
                          'likely be removed in v0.7.  Use the oversampling '
                          'keyword instead.', AstropyDeprecationWarning)

        super().__init__(
            data=data, flux=flux, x_0=x_0, y_0=y_0, normalize=normalize,
            normalization_correction=normalization_correction, origin=origin,
            oversampling=oversampling, fill_value=fill_value, ikwargs=ikwargs)

        self._pixel_scale = pixel_scale

    @property
    def pixel_scale(self):
        """
        The ``(x, y)`` pixel scale (in arbitrary units) of the PSF.
        """

        return self._pixel_scale

    @pixel_scale.setter
    def pixel_scale(self, pixel_scale):
        if pixel_scale is not None:
            pixel_scale = np.atleast_1d(pixel_scale)
            if len(pixel_scale) == 1:
                pixel_scale = np.repeat(pixel_scale, 2).astype(float)
            elif len(pixel_scale) > 2:
                raise ValueError('pixel_scale must be a scalar or tuple '
                                 'of two floats.')

        self._pixel_scale = pixel_scale

class GriddedPSFModel(Fittable2DModel):
    """
    A fittable 2D model containing a grid PSF models defined at specific
    locations that are interpolated to evaluate a PSF at an arbitrary
    (x, y) position.

    Parameters
    ----------
    data : `~astropy.nddata.NDData`
        An `~astropy.nddata.NDData` object containing the grid of
        reference PSF arrays.  The data attribute must contain a 3D
        `~numpy.ndarray` containing a stack of the 2D PSFs (the data
        shape should be (N_psf, PSF_ny, PSF_nx)).  The meta
        attribute must be `dict` containing the following:

            * ``'grid_xypos'``:  A list of the (x, y) grid positions of
              each reference PSF.  The order of positions should match
              the first axis of the 3D `~numpy.ndarray` of PSFs.  In
              other words, ``grid_xypos[i]`` should be the (x, y)
              position of the reference PSF defined in ``data[i]``.
            * ``'oversampling'``:  The integer oversampling factor of the
               PSF.

        The meta attribute may contain other properties such as the
        telescope, instrument, detector, and filter of the PSF.
    """

    flux = Parameter(description='Intensity scaling factor for the PSF '
                     'model.', default=1.0)
    x_0 = Parameter(description='x position in the output coordinate grid '
                    'where the model is evaluated.', default=0.0)
    y_0 = Parameter(description='y position in the output coordinate grid '
                    'where the model is evaluated.', default=0.0)

    def __init__(self, data, flux=flux.default, x_0=x_0.default,
                 y_0=y_0.default, fill_value=0.0):

        if not isinstance(data, NDData):
            raise TypeError('data must be an NDData instance.')

        if data.data.ndim != 3:
            raise ValueError('The NDData data attribute must be a 3D numpy '
                             'ndarray')

        if 'grid_xypos' not in data.meta:
            raise ValueError('"grid_xypos" must be in the nddata meta '
                             'dictionary.')
        if len(data.meta['grid_xypos']) != data.data.shape[0]:
            raise ValueError('The length of grid_xypos must match the number '
                             'of input PSFs.')

        if 'oversampling' not in data.meta:
            raise ValueError('"oversampling" must be in the nddata meta '
                             'dictionary.')
        if not np.isscalar(data.meta['oversampling']):
            raise ValueError('oversampling must be a scalar value')

        self.data = np.array(data.data, copy=True, dtype=np.float)
        self.meta = data.meta
        self.grid_xypos = data.meta['grid_xypos']
        self.oversampling = data.meta['oversampling']

        self._grid_xpos, self._grid_ypos = np.transpose(self.grid_xypos)
        self._xgrid = np.unique(self._grid_xpos)  # also sorts values
        self._ygrid = np.unique(self._grid_ypos)  # also sorts values

        if (len(list(itertools.product(self._xgrid, self._ygrid))) !=
                len(self.grid_xypos)):
            raise ValueError('"grid_xypos" must form a regular grid.')

        self._xgrid_min = self._xgrid[0]
        self._xgrid_max = self._xgrid[-1]
        self._ygrid_min = self._ygrid[0]
        self._ygrid_max = self._ygrid[-1]

        super().__init__(flux, x_0, y_0)

    @staticmethod
    def _find_bounds_1d(data, x):
        """
        Find the index of the lower bound where ``x`` should be inserted
        into ``a`` to maintain order.

        The index of the upper bound is the index of the lower bound
        plus 2.  Both bound indices must be within the array.

        Parameters
        ----------
        data : 1D `~numpy.ndarray`
            The 1D array to search.

        x : float
            The value to insert.

        Returns
        -------
        index : int
            The index of the lower bound.
        """

        idx = np.searchsorted(data, x)
        if idx == 0:
            idx0 = 0
        elif idx == len(data):  # pragma: no cover
            idx0 = idx - 2
        else:
            idx0 = idx - 1

        return idx0

    def _find_bounding_points(self, x, y):
        """
        Find the indices of the grid points that bound the input
        ``(x, y)`` position.

        Parameters
        ----------
        x, y : float
            The ``(x, y)`` position where the PSF is to be evaluated.

        Returns
        -------
        indices : list of int
            A list of indices of the bounding grid points.
        """

        if not np.isscalar(x) or not np.isscalar(y):  # pragma: no cover
            raise TypeError('x and y must be scalars')

        if (x < self._xgrid_min or x > self._xgrid_max or
                y < self._ygrid_min or y > self._ygrid_max):  # pragma: no cover
            raise ValueError('(x, y) position is outside of the region '
                             'defined by grid of PSF positions')

        x0 = self._find_bounds_1d(self._xgrid, x)
        y0 = self._find_bounds_1d(self._ygrid, y)
        points = list(itertools.product(self._xgrid[x0:x0 + 2],
                                        self._ygrid[y0:y0 + 2]))

        indices = []
        for xx, yy in points:
            indices.append(np.argsort(np.hypot(self._grid_xpos - xx,
                                               self._grid_ypos - yy))[0])

        return indices

    @staticmethod
    def _bilinear_interp(xyref, zref, xi, yi):
        """
        Perform bilinear interpolation of four 2D arrays located at
        points on a regular grid.

        Parameters
        ----------
        xyref : list of 4 (x, y) pairs
            A list of 4 ``(x, y)`` pairs that form a rectangle.

        refdata : 3D `~numpy.ndarray`
            A 3D `~numpy.ndarray` of shape ``(4, nx, ny)``.  The first
            axis corresponds to ``xyref``, i.e. ``refdata[0, :, :]`` is
            the 2D array located at ``xyref[0]``.

        xi, yi : float
            The ``(xi, yi)`` point at which to perform the
            interpolation.  The ``(xi, yi)`` point must lie within the
            rectangle defined by ``xyref``.

        Returns
        -------
        result : 2D `~numpy.ndarray`
            The 2D interpolated array.
        """

        if len(xyref) != 4:
            raise ValueError('xyref must contain only 4 (x, y) pairs')

        if zref.shape[0] != 4:
            raise ValueError('zref must have a length of 4 on the first '
                             'axis.')

        xyref = [tuple(i) for i in xyref]
        idx = sorted(range(len(xyref)), key=xyref.__getitem__)
        xyref = sorted(xyref)   # sort by x, then y
        (x0, y0), (_x0, y1), (x1, _y0), (_x1, _y1) = xyref

        if x0 != _x0 or x1 != _x1 or y0 != _y0 or y1 != _y1:
            raise ValueError('The refxy points do not form a rectangle.')

        if not np.isscalar(xi):
            xi = xi[0]
        if not np.isscalar(yi):
            yi = yi[0]

        if not x0 <= xi <= x1 or not y0 <= yi <= y1:
            raise ValueError('The (x, y) input is not within the rectangle '
                             'defined by xyref.')

        data = np.asarray(zref)[idx]
        weights = np.array([(x1 - xi) * (y1 - yi), (x1 - xi) * (yi - y0),
                            (xi - x0) * (y1 - yi), (xi - x0) * (yi - y0)])
        norm = (x1 - x0) * (y1 - y0)

        return np.sum(data * weights[:, None, None], axis=0) / norm

    def evaluate(self, x, y, flux, x_0, y_0):
        """
        Evaluate the `GriddedPSFModel` for the input parameters.
        """

        # NOTE: this is needed because the PSF photometry routines input
        # length-1 values instead of scalars.  TODO: fix the photometry
        # routines.
        if not np.isscalar(x_0):
            x_0 = x_0[0]
        if not np.isscalar(y_0):
            y_0 = y_0[0]

        if (x_0 < self._xgrid_min or x_0 > self._xgrid_max or
                y_0 < self._ygrid_min or y_0 > self._ygrid_max):

            # position is outside of the grid, so simply use the
            # closest reference PSF
            self._ref_indices = np.argsort(np.hypot(self._grid_xpos - x_0,
                                                    self._grid_ypos - y_0))[0]
            self._psf_interp = self.data[self._ref_indices, :, :]
        else:
            # find the four bounding reference PSFs and interpolate
            self._ref_indices = self._find_bounding_points(x_0, y_0)
            xyref = np.array(self.grid_xypos)[self._ref_indices]
            psfs = self.data[self._ref_indices, :, :]

            self._psf_interp = self._bilinear_interp(xyref, psfs, x_0, y_0)

        # now evaluate the PSF at the (x_0, y_0) subpixel position on
        # the input (x, y) values
        psfmodel = FittableImageModel(self._psf_interp,
                                      oversampling=self.oversampling)

        return psfmodel.evaluate(x, y, flux, x_0, y_0)

class IntegratedGaussianPRF(Fittable2DModel):
    r"""
    Circular Gaussian model integrated over pixels. Because it is
    integrated, this model is considered a PRF, *not* a PSF (see
    :ref:`psf-terminology` for more about the terminology used here.)

    This model is a Gaussian *integrated* over an area of ``1`` (in
    units of the model input coordinates, e.g. 1 pixel).  This is in
    contrast to the apparently similar
    `astropy.modeling.functional_models.Gaussian2D`, which is the value
    of a 2D Gaussian *at* the input coordinates, with no integration.
    So this model is equivalent to assuming the PSF is Gaussian at a
    *sub-pixel* level.

    Parameters
    ----------
    sigma : float
        Width of the Gaussian PSF.
    flux : float (default 1)
        Total integrated flux over the entire PSF
    x_0 : float (default 0)
        Position of the peak in x direction.
    y_0 : float (default 0)
        Position of the peak in y direction.

    Notes
    -----
    This model is evaluated according to the following formula:

        .. math::

            f(x, y) =
                \frac{F}{4}
                \left[
                {\rm erf} \left(\frac{x - x_0 + 0.5}
                {\sqrt{2} \sigma} \right) -
                {\rm erf} \left(\frac{x - x_0 - 0.5}
                {\sqrt{2} \sigma} \right)
                \right]
                \left[
                {\rm erf} \left(\frac{y - y_0 + 0.5}
                {\sqrt{2} \sigma} \right) -
                {\rm erf} \left(\frac{y - y_0 - 0.5}
                {\sqrt{2} \sigma} \right)
                \right]

    where ``erf`` denotes the error function and ``F`` the total
    integrated flux.
    """

    flux = Parameter(default=1)
    x_0 = Parameter(default=0)
    y_0 = Parameter(default=0)
    sigma = Parameter(default=1, fixed=True)

    _erf = None
    fit_deriv = None

    @property
    def bounding_box(self):
        halfwidth = 4 * self.sigma
        return ((int(self.y_0 - halfwidth), int(self.y_0 + halfwidth)),
                (int(self.x_0 - halfwidth), int(self.x_0 + halfwidth)))

    def __init__(self, sigma=sigma.default,
                 x_0=x_0.default, y_0=y_0.default, flux=flux.default,
                 **kwargs):
        if self._erf is None:
            from scipy.special import erf
            self.__class__._erf = erf

        super().__init__(n_models=1, sigma=sigma, x_0=x_0, y_0=y_0, flux=flux,
                         **kwargs)

    def evaluate(self, x, y, flux, x_0, y_0, sigma):
        """Model function Gaussian PSF model."""

        return (flux / 4 *
                ((self._erf((x - x_0 + 0.5) / (np.sqrt(2) * sigma)) -
                  self._erf((x - x_0 - 0.5) / (np.sqrt(2) * sigma))) *
                 (self._erf((y - y_0 + 0.5) / (np.sqrt(2) * sigma)) -
                  self._erf((y - y_0 - 0.5) / (np.sqrt(2) * sigma)))))

class PRFAdapter(Fittable2DModel):
    """
    A model that adapts a supplied PSF model to act as a PRF. It
    integrates the PSF model over pixel "boxes".  A critical built-in
    assumption is that the PSF model scale and location parameters are
    in *pixel* units.

    Parameters
    ----------
    psfmodel : a 2D model
        The model to assume as representative of the PSF
    renormalize_psf : bool
        If True, the model will be integrated from -inf to inf and
        re-scaled so that the total integrates to 1.  Note that this
        renormalization only occurs *once*, so if the total flux of
        ``psfmodel`` depends on position, this will *not* be correct.
    xname : str or None
        The name of the ``psfmodel`` parameter that corresponds to the
        x-axis center of the PSF.  If None, the model will be assumed to
        be centered at x=0.
    yname : str or None
        The name of the ``psfmodel`` parameter that corresponds to the
        y-axis center of the PSF.  If None, the model will be assumed to
        be centered at y=0.
    fluxname : str or None
        The name of the ``psfmodel`` parameter that corresponds to the
        total flux of the star.  If None, a scaling factor will be
        applied by the ``PRFAdapter`` instead of modifying the
        ``psfmodel``.

    Notes
    -----
    This current implementation of this class (using numerical
    integration for each pixel) is extremely slow, and only suited for
    experimentation over relatively few small regions.
    """

    flux = Parameter(default=1)
    x_0 = Parameter(default=0)
    y_0 = Parameter(default=0)

    def __init__(self, psfmodel, renormalize_psf=True, flux=flux.default,
                 x_0=x_0.default, y_0=y_0.default, xname=None, yname=None,
                 fluxname=None, **kwargs):

        self.psfmodel = psfmodel.copy()

        if renormalize_psf:
            from scipy.integrate import dblquad
            self._psf_scale_factor = 1. / dblquad(self.psfmodel,
                                                  -np.inf, np.inf,
                                                  lambda x: -np.inf,
                                                  lambda x: np.inf)[0]
        else:
            self._psf_scale_factor = 1

        self.xname = xname
        self.yname = yname
        self.fluxname = fluxname

        # these can be used to adjust the integration behavior. Might be
        # used in the future to expose how the integration happens
        self._dblquadkwargs = {}

        super().__init__(n_models=1, x_0=x_0, y_0=y_0, flux=flux, **kwargs)

    def evaluate(self, x, y, flux, x_0, y_0):
        """The evaluation function for PRFAdapter."""

        if self.xname is None:
            dx = x - x_0
        else:
            dx = x
            setattr(self.psfmodel, self.xname, x_0)

        if self.xname is None:
            dy = y - y_0
        else:
            dy = y
            setattr(self.psfmodel, self.yname, y_0)

        if self.fluxname is None:
            return (flux * self._psf_scale_factor *
                    self._integrated_psfmodel(dx, dy))
        else:
            setattr(self.psfmodel, self.yname, flux * self._psf_scale_factor)
            return self._integrated_psfmodel(dx, dy)

    def _integrated_psfmodel(self, dx, dy):
        from scipy.integrate import dblquad

        # infer type/shape from the PSF model.  Seems wasteful, but the
        # integration step is a *lot* more expensive so its just peanuts
        out = np.empty_like(self.psfmodel(dx, dy))
        outravel = out.ravel()
        for i, (xi, yi) in enumerate(zip(dx.ravel(), dy.ravel())):
            outravel[i] = dblquad(self.psfmodel,
                                  xi-0.5, xi+0.5,
                                  lambda x: yi-0.5, lambda x: yi+0.5,
                                  **self._dblquadkwargs)[0]
        return out


"""
CDI.py
Kristina Davis
8/12/19

This module is used to set the CDI parameters for mini-medis.
"""
##
import numpy as np
import warnings
from scipy import linalg, interpolate
from matplotlib import pyplot as plt
from matplotlib.colors import LogNorm, SymLogNorm
import time
import pickle

from medis.params import tp, sp, ap, iop
from medis.utils import dprint
from medis.optics import extract_plane, cpx_to_intensity, extract_center
from medis.plot_tools import add_colorbar, view_timeseries

##
class Slapper:
    """Stole this idea from falco. Just passes an object you can slap stuff onto"""
    pass

class CDI_params:
    """
    contains the parameters of the CDI probes and phase sequence to apply to the DM
    """
    def __init__(self):
        # General
        self.use_cdi = False
        self.show_probe = False  # False , flag to plot phase probe or not
        self.which_DM = ''  # plane_name parameter of the DM to apply CDI probe to (must be a valid name for
        # the telescope sim you are running, eg 'tweeter'

        # Probe Dimensions (extent in pupil plane coordinates)
        self.probe_ax = 'X'  # 'Y' # direction of the probe
        self.probe_amp = 2e-6  # [m] probe amplitude, scale should be in units of actuator height limits
        self.probe_w = 10  # [actuator coordinates] width of the probe
        self.probe_h = 30  # [actuator coordinates] height of the probe
        self.probe_shift = (15, 15)  # [actuator coordinates] center position of the probe (should move off-center to
        # avoid coronagraph)
        self.probe_spacing = 10  # distance from the focal plane center to edge of the rectangular probed region

        # Phase Sequence of Probes
        self.phs_intervals = np.pi / 2  # [rad] phase interval over [0, 2pi]
        self.probe_integration_time = 1  # [s]  How long in sec to apply each probe in the sequence
        self.null_time = 5 * sp.sample_time # [s]  time between repeating probe cycles (data to be nulled using probe info)

    def __iter__(self):
        for attr, value in self.__dict__.items():
            yield attr, value

    def __name__(self):
        return self.__str__().split(' ')[0].split('.')[-1]

    def gen_phaseseries(self):
        """
        generate an array of phases per timestep for the CDI algorithm

        phase_series is used to populate cdi.phase_series, which may be longer than cdi.phase_cycle if multiple cycles
        are run, or probes may last for longer than one single timestep

        currently, I assume the timestream is not that long. Should only use this for short timestreams, and use a more
        efficient code for long simulations (scale of minutes or more)

        :return: phase_series  array of phases of CDI probes to apply to DM
        """
        self.phase_series = np.zeros(sp.numframes) * np.nan

        # Repeating Probe Phases for Integration time
        self.phase_cycle = np.arange(0, 2 * np.pi, self.phs_intervals)  # FYI not inclusive of 2pi endpoint
        self.n_probes = len(self.phase_cycle)  # number of phase probes
        if self.n_probes % 2 != 0:
            raise ValueError(f"must have even number of phase probes\n\tchange cdi.phs_intervals")
        self.DM_probe_series = np.zeros((self.n_probes, tp.act_tweeter, tp.act_tweeter))
        self.cmd_tstamps = np.zeros((sp.numframes,),  dtype='datetime64[ns]')

        if self.probe_integration_time > sp.sample_time:
            phase_hold = self.probe_integration_time / sp.sample_time
            phase_1cycle = np.repeat(self.phase_cycle, phase_hold)
        elif self.probe_integration_time == sp.sample_time:
            phase_1cycle = self.phase_cycle
        else:
            raise ValueError(f"Cannot have CDI phase probe integration time less than sp.sample_time")

        # Repeating Cycle of Phase Probes for Simulation Duration
        full_simulation_time = sp.numframes * sp.sample_time
        self.time_for_one_cycle = len(phase_1cycle) * self.probe_integration_time

        if self.time_for_one_cycle > full_simulation_time and self.probe_integration_time >= sp.sample_time:
            warnings.warn(f"\nLength of one full CDI probe cycle exceeds the "
                          f"full simulation time \n"
                          f"not all phases will be used\n"
                          f"phase reconstruction will be incomplete")
            self.phase_series = phase_1cycle[0:sp.numframes]
        elif self.time_for_one_cycle <= full_simulation_time and self.n_probes < sp.numframes:
            print(f"\nCDI Params\n\tThere will be {sp.numframes - self.n_probes} "
                  f"nulling steps after timestep {self.n_probes}")
            self.phase_series[0:self.n_probes] = phase_1cycle
        else:
            warnings.warn(f"Haven't run into  CDI phase situation like this yet")
            raise NotImplementedError

        return self.phase_series

    def save_probe(self, ix, probe):
        self.DM_probe_series[ix, :, :] = probe

    def save_tseries(self, ix, ts):
        """saves output of medis fields as 2D intensity images for CDI postprocessing"""
        self.cmd_tstamps[ix] = ts

    def save_out_to_disk(self, plot=False):
        out = Slapper()
        out.probe = Slapper()
        out.ts = Slapper()

        # Probe Info
        out.probe.direction = self.probe_ax
        out.probe.amp = self.probe_amp
        out.probe.width = self.probe_w
        out.probe.height = self.probe_h
        out.probe.shift = self.probe_shift
        out.probe.spacing = self.probe_spacing
        out.probe.DM_cmd_cycle = self.DM_probe_series
        out.probe.phs_interval = self.phs_intervals

        # Timeseries Info
        out.ts.start = 0
        out.ts.n_probes = self.n_probes
        out.ts.phase_cycle = self.phase_cycle
        out.ts.probe_integration_time = self.probe_integration_time
        out.ts.t_one_cycle = self.time_for_one_cycle
        out.ts.null_time = self.null_time
        out.ts.elapsed_time = 0
        out.ts.n_cycles = 0
        out.ts.n_cmds = sp.numframes  # TODO verify this-this was just a late night hack
        out.ts.cmd_tstamps = self.cmd_tstamps

        save_location = iop.testdir + f"/{iop.testname}_CDIparams.pkl"
        dprint(f'save_location={save_location}')
        with open(save_location, 'wb') as handle:
            pickle.dump(out, handle, protocol=pickle.HIGHEST_PROTOCOL)
        handle.close()

        # Fig
        if plot:
            if self.n_probes >= 4:
                nrows = 2
                ncols = self.n_probes//2
                figheight = 8
            else:
                nrows = 1
                ncols = self.n_probes
                figheight = 4

            fig, subplot = plt.subplots(nrows, ncols, figsize=(10, figheight))
            fig.subplots_adjust(wspace=0.5, right=0.85, left=0.05)
            fig.suptitle('Probe Series')

            for ax, ix in zip(subplot.flatten(), range(out.ts.n_probes)):
                im = ax.imshow(out.probe.DM_cmd_cycle[ix], interpolation='none', origin='lower')
                ax.set_title(f"Probe " + r'$\theta$=' + f'{out.ts.phase_cycle[ix]/np.pi:.2f}' + r'$\pi$')

            cax = fig.add_axes([0.9, 0.2, 0.03, 0.6])  # Add axes for colorbar @ position [left,bottom,width,height]
            cb = fig.colorbar(im, orientation='vertical', cax=cax)  #
            cb.set_label('Probe Height (m)')

        return out

# Sneakily Instantiating Class Objects here
cdi = CDI_params()

##
def config_probe(theta, nact, iw=0, ib=0, tstep=0):
    """
    create a probe shape to apply to the DM for CDI processing

    The probe applied to the DM to achieve CDI is that originally proposed in Giv'on et al 2011, doi: 10.1117/12.895117;
    and was used with proper in Matthews et al 2018, doi:  10.1117/1.JATIS.3.4.045001.

    The pupil coordinates used in those equations relate to the sampling of the pupil (plane of the DM). However, in
    Proper, the prop_dm uses a DM map that is the size of (n_ao_act, n_ao_act). This map was resampled from its
    original size to the actuator spacing, and the spacing in units of [m] is supplied as a keyword to the prop_dm
    function during the call. All that is to say, we apply the CDI probe using the coordinates of the DM actuators,
    and supply the probe height as an additive height to the DM map, which is passed to the prop_dm function.

    :param theta: phase of the probe
    :param nact: number of actuators in the mirror, should change if 'woofer' or 'tweeter'
    :param iw: index of wavelength number in ap.wvl_range (used for plotting only)
    :param ib: index of astronomical body eg star or companion (used for plotting only)
    :return: height of phase probes to add to the DM map in adaptive.py
    """
    x = np.linspace(-1/2-cdi.probe_shift[0]/nact, 1/2-cdi.probe_shift[0]/nact, nact)
    y = np.linspace(-1/2-cdi.probe_shift[1]/nact, 1/2-cdi.probe_shift[1]/nact, nact)
    X, Y = np.meshgrid(x, y)

    wvl_samples = np.linspace(ap.wvl_range[0], ap.wvl_range[1], ap.n_wvl_init)
    # dprint(f'iw = {iw}, lambda = {wvl_samples[iw]}')
    mag = 4 * np.pi * wvl_samples[iw] * cdi.probe_amp

    if cdi.probe_ax == 'X' or cdi.probe_ax == 'x':
        dir = X
    elif cdi.probe_ax == 'Y' or cdi.probe_ax == 'y':
        dir = Y
    else:
        raise ValueError('probe direction value not understood; must be string "X" or "Y"')

    probe = mag * np.sinc(cdi.probe_w * X) * np.sinc(cdi.probe_h * Y) \
            * np.sin(2*np.pi*cdi.probe_spacing * dir + theta)

    # Testing FF propagation
    if sp.verbose and iw == 0 and ib == 0:  # and theta == cdi.phase_series[0]
        probe_ft = (1/np.sqrt(2*np.pi)) * np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(probe)))

        fig, ax = plt.subplots(1, 3, figsize=(12, 5))
        fig.subplots_adjust(wspace=0.5)
        ax1, ax2, ax3 = ax.flatten()

        fig.suptitle(f"spacing={cdi.probe_spacing}, Dimensions {cdi.probe_w}x{cdi.probe_h} "
                     f"\nProbe Amp = {cdi.probe_amp}, " + r'$\theta$' + f"={theta/np.pi:.3f}" + r'$\pi$' + '\n')

        im1 = ax1.imshow(probe, interpolation='none')
        ax1.set_title(f"Probe on DM \n(dm coordinates)")
        # cb = fig.colorbar(im1, ax=ax1)
        add_colorbar(im1)

        im2 = ax2.imshow(np.sqrt(probe_ft.imag ** 2 + probe_ft.real ** 2), interpolation='none')
        ax2.set_title("Focal Plane Amplitude")
        add_colorbar(im2)

        im3 = ax3.imshow(np.arctan2(probe_ft.imag, probe_ft.real), interpolation='none', cmap='hsv')
        ax3.set_title("Focal Plane Phase")
        add_colorbar(im3)

        # =========================
        # Fig 2  Real & Imag
        # fig, ax = plt.subplots(1, 3, figsize=(12, 5))
        # fig.subplots_adjust(wspace=0.5)
        # ax1, ax2, ax3 = ax.flatten()
        # fig.suptitle(f'Real & Imaginary Probe Response in Focal Plane\n'
        #              f' '+r'$\theta$'+f'={theta/np.pi:.3f}'+r'$\pi$'+f', n_actuators = {nact}\n')
        #
        # im1 = ax1.imshow(probe, interpolation='none', origin='lower')
        # ax1.set_title(f"Probe on DM \n(dm coordinates)")
        # cb = fig.colorbar(im1, ax=ax1)
        #
        # im2 = ax2.imshow(probe_ft.real, interpolation='none', origin='lower')
        # ax2.set_title(f"Real FT of Probe")
        #
        # im3 = ax3.imshow(probe_ft.imag, interpolation='none', origin='lower')
        # ax3.set_title(f"Imag FT of Probe")

        plt.show()

    # Saving Probe in the series
    if iw == 0 and ib == 0:
        ip = np.argwhere(cdi.phase_series == theta)
        cdi.save_probe(ip[0,0], probe)  #
        cdi.nact = nact

    return probe

##
def cdi_postprocess(cpx_sequence, sampling, plot=False):
    """
    this is the function that accepts the timeseries of intensity images from the simulation and returns the processed
    single image. This function calculates the speckle amplitude phase, and then corrects for it to create the dark
    hole over the specified region of the image.

    From Give'on et al 2011, we have in eq 10 two different values: DeltaP-the change in the focal plane due to the
    probe, and delta, the intensity difference measurements between pairs of phase probes. (It is unfortunate that
    both these terms use delta, so I have done my best to distinguish them in the variable names)

    Here I note that in the CDI phase stream generation, for n_probes there are n_pairs = n_probes/2 pairs of probes.
    These get applied to the DM in a series such that the two probes that form the conjugate pair are separated by
    n_pairs of probes. In other words, for n_probes = 6, the 0th and 3rd probes are a pair, the 1st and 4th are a pair,
    and so on. This is a choice made when creating cdi.phase_series.

    :param cpx_sequence: #timestream of 2D images (complex) from the focal plane complex field
    :param sampling: focal plane sampling
    :return:
    """
    ##
    tic = time.time()
    focal_plane = extract_plane(cpx_sequence, 'detector')  # eliminates astro_body axis [tsteps,wvl,obj,x,y]
    fp_seq = np.sum(focal_plane, axis=(1,2))  # sum over wavelength,object

    n_pairs = cdi.n_probes//2  # number of deltas (probe differentials)
    n_nulls = sp.numframes - cdi.n_probes
    delta = np.zeros((n_pairs, sp.grid_size, sp.grid_size), dtype=float)
    # absDelta = np.zeros((n_nulls, sp.grid_size, sp.grid_size))
    # phsDelta = np.zeros((n_pairs, sp.grid_size, sp.grid_size), dtype=float)
    E_est = np.zeros((n_nulls, sp.grid_size, sp.grid_size), dtype=complex)
    I_processed = np.zeros((n_nulls, sp.grid_size, sp.grid_size))
    H = np.zeros((n_pairs, 2), dtype=float)
    b = np.zeros((n_pairs, 1))

    # Get Masked Data
    mask2D, imsk, jmsk, irng, jrng, imx, imn, jmx, jmn = get_fp_mask(cdi, thresh=1e-6)

    if sp.debug:
        fig, ax = plt.subplots(1,1)
        fig.suptitle(f'Masked FP in CDI probe Region')
        im = ax.imshow(cpx_to_intensity(fp_seq[0,:,:]*mask2D))

    for ip in range(n_pairs):
        # Compute deltas (I_ip+ - I_ip-)/2
        delta[ip] = (np.abs(fp_seq[ip])**2 - np.abs(fp_seq[ip + n_pairs])**2) / 4

    # for i,j in zip(imsk,jmsk):
    # for i,j in zip(irng,jrng):
    for i in irng:
        for j in jrng:
            for xn in range(n_nulls):
                for ip in range(n_pairs):
                    # Amplitude DeltaP
                    Ip = np.abs(fp_seq[ip, i, j]) ** 2
                    Im = np.abs(fp_seq[ip + n_pairs, i, j]) ** 2
                    Io = np.abs(fp_seq[cdi.n_probes + xn, i, j]) ** 2
                    abs = (Ip + Im) / 2 - Io
                    if abs < 0:
                        abs = 0
                    absDeltaP = np.sqrt(abs)
                    # absDeltaP = np.sqrt(np.abs((Ip + Im) / 2 - Io))

                    # phsDeltaP = phsDelta[ip, i, j]
                    # Phase DeltaP
                    # The phase of the change in the focal plane of the probe applied to the DM
                    # First subtract Eo vector from each probe phase to make new field vectors dEa, dEb,
                    # then take the angle between the two
                    dEp = fp_seq[ip, i, j] - fp_seq[cdi.n_probes + xn, i, j]
                    dEm = fp_seq[ip + n_pairs, i, j] - fp_seq[cdi.n_probes + xn, i, j]
                    phsDeltaP = np.arctan2(dEp.imag - dEm.imag, dEp.real - dEm.real)

                    cpxDeltaP = absDeltaP * np.exp(1j * phsDeltaP)
                    # cpxDeltaP = absDeltaP * np.array((np.cos(phsDeltaP), np.sin(phsDeltaP)))

                    H[ip, :] = [-cpxDeltaP.imag, cpxDeltaP.real]  # [n_pairs, 2]
                    # H[ip,:] = [cpxDeltaP[0], cpxDeltaP[1]]
                    b[ip] = delta[ip, i, j]  # [n_pairs, 1]

                a = 2 * H
                Exy = linalg.lstsq(a, b)[0]  # returns tuple, not array
                E_est[xn, i, j] = Exy[0] + (1j * Exy[1])

    toc = time.time()
    dprint(f'CDI post-processing took {(toc-tic)/60:.2} minutes\n')

    ## ===========================
    # Contrast Ratios
    # ===========================
    intensity_probe        = np.zeros(n_nulls)
    intensity_DM_FFT       = np.zeros(n_nulls)
    intensity_pre_process  = np.zeros(n_nulls)
    intensity_post_process = np.zeros(n_nulls)

    for xn in range(n_nulls):
        I_processed[xn] = np.abs(fp_seq[n_pairs+xn])**2 - np.abs(np.conj(E_est[xn])*mask2D)**2
        # I_processed[xn] = np.sqrt(np.abs(np.abs(fp_seq[n_pairs+xn])**2 - np.abs(E_est[xn]*mask2D)**2))**2
        # I_processed[xn] = np.abs(fp_seq[n_pairs+xn] - np.conj(E_est[xn]*mask2D))**2

        # Contrast
        intensity_probe[xn] = np.sum(np.abs(fp_seq[xn]*mask2D)**2)
        intensity_pre_process[xn] = np.sum(np.abs(fp_seq[n_pairs + xn]*mask2D)**2)
        intensity_post_process[xn] = np.sum(I_processed[xn]*mask2D)  #np.sum(np.abs(E_processed[xn]*mask2D)**2)

        print(f'\nIntensity in probed region for null step {xn} is '
              f'\nprobe {intensity_probe[xn]}'
              f'\npre-processed {intensity_pre_process[xn]} '
              f'\npost-processed {intensity_post_process[xn]}'
              f'\n difference = {intensity_post_process[xn] - intensity_pre_process[xn]}'
              f'\n')

 ##   if plot:
        # ==================
        # FFT of Tweeter Plane
        # ==================
        # fig, subplot = plt.subplots(1, n_pairs, figsize=(14, 5))
        # fig.subplots_adjust(wspace=0.5, right=0.85)
        # fig.suptitle('Tweeter DM Plane')
        #
        # tweet = extract_plane(cpx_sequence, 'tweeter')  # eliminates astro_body axis [tsteps,wvl,obj,x,y]
        # tweeter = np.sum(tweet, axis=(1, 2))
        # tweeter_intensity = np.abs(tweeter) ** 2
        # for ax, ix in zip(subplot.flatten(), range(n_pairs)):
        #     im = ax.imshow(extract_center(tweeter_intensity[ix], new_size=sp.grid_size*sp.beam_ratio+10),
        #                    interpolation='none', norm=LogNorm(),
        #                    vmin=5e-3, vmax=1e-2)
        #     ax.set_title(f'Probe Phase ' r'$\theta$' f'={cdi.phase_cycle[ix] / np.pi:.2f}' r'$\pi$')
        #
        # cax = fig.add_axes([0.9, 0.2, 0.03, 0.6])  # Add axes for colorbar @ position [left,bottom,width,height]
        # cb = fig.colorbar(im, orientation='vertical', cax=cax)  #
        # cb.set_label('Intensity')
        #
        # # =================================================================================
        # fig, subplot = plt.subplots(1, n_pairs, figsize=(14, 5))
        # fig.subplots_adjust(wspace=0.5, right=0.85)
        # fig.suptitle('FFT of Tweeter DM Plane')
        #
        # for ax, ix in zip(subplot.flatten(), range(n_pairs)):
        #     fft_tweeter = (1 / np.sqrt(2 * np.pi) *
        #                    np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(tweeter[ix]))))
        #     im = ax.imshow(np.abs(fft_tweeter) ** 2,
        #                    interpolation='none', norm=LogNorm(),
        #                    vmin=1e-2, vmax=10)
        #     ax.set_title(f'Probe Phase ' r'$\theta$' f'={cdi.phase_cycle[ix] / np.pi:.2f}' r'$\pi$')
        #
        # cax = fig.add_axes([0.9, 0.2, 0.03, 0.6])  # Add axes for colorbar @ position [left,bottom,width,height]
        # cb = fig.colorbar(im, orientation='vertical', cax=cax)  #
        # cb.set_label('Intensity')

        # ==================
        # Deltas
        # ==================
        # fig, subplot = plt.subplots(1, n_pairs, figsize=(14,5))
        # fig.subplots_adjust(wspace=0.5, right=0.85)
        # fig.suptitle('Deltas for CDI Probes')
        #
        # for ax, ix in zip(subplot.flatten(), range(n_pairs)):
        #     im = ax.imshow(delta[ix]*1e6*mask2D, interpolation='none',
        #                    norm=SymLogNorm(linthresh=1),
        #                    # vmin=-1, vmax=1
        #                    ) #, norm=SymLogNorm(linthresh=1e-5))
        #     ax.set_title(f"Diff Probe\n" + r'$\theta$' + f'={cdi.phase_series[ix]/np.pi:.3f}' +
        #                  r'$\pi$ -$\theta$' + f'={cdi.phase_series[ix+n_pairs]/np.pi:.3f}' + r'$\pi$')
        #
        # cax = fig.add_axes([0.9, 0.2, 0.03, 0.6])  # Add axes for colorbar @ position [left,bottom,width,height]
        # cb = fig.colorbar(im, orientation='vertical', cax=cax)  #
        # cb.set_label('Intensity')

        # ==================
        # Null Probe E-Field
        # ==================
        fig, subplot = plt.subplots(1, n_nulls, figsize=(14, 5))
        fig.subplots_adjust(wspace=0.5, right=0.85)
        fig.suptitle('Original (Null-Probe) Image Plane Intensity')

        for ax, ix in zip(subplot.flatten(), range(n_nulls)):
            im = ax.imshow(np.abs(fp_seq[cdi.n_probes + ix, irng[0]:irng[-1], jrng[0]:jrng[-1]]) ** 2,  # , 250:270, 150:170  *mask2D , irng[0]:irng[-1], jrng[0]:jrng[-1]
                           interpolation='none', norm=LogNorm(),
                           vmin=1e-7, vmax=1e-4)
            ax.set_title(f'Null Step {ix}')

        cax = fig.add_axes([0.9, 0.2, 0.03, 0.6])  # Add axes for colorbar @ position [left,bottom,width,height]
        cb = fig.colorbar(im, orientation='vertical', cax=cax)  #
        cb.set_label('Intensity')

        # ===============================
        # E-filed Estimate
        # ===============================
        fig, subplot = plt.subplots(1, n_nulls, figsize=(14, 5))
        fig.subplots_adjust(wspace=0.5, right=0.85)
        fig.suptitle('Pairwise Probe Estimated Image Plane intensity')

        for ax, ix in zip(subplot.flatten(), range(n_nulls)):
            im = ax.imshow(np.abs(E_est[ix, irng[0]:irng[-1], jrng[0]:jrng[-1]])**2,  # , 250:270, 150:170  *mask2D , irng[0]:irng[-1], jrng[0]:jrng[-1]
                           interpolation='none',
                           norm=LogNorm(),
                           vmin=1e-7, vmax=1e-4
                           )  # , norm=SymLogNorm(linthresh=1e-5))
            ax.set_title(f'Null Step {ix}')

        cax = fig.add_axes([0.9, 0.2, 0.03, 0.6])  # Add axes for colorbar @ position [left,bottom,width,height]
        cb = fig.colorbar(im, orientation='vertical', cax=cax)  #
        cb.set_label('Intensity')

        # ==================================
        # Processed E-field (CDI Subtracted)
        # ==================================
        fig, subplot = plt.subplots(1, n_nulls, figsize=(14, 5))
        fig.subplots_adjust(wspace=0.5, right=0.85)
        fig.suptitle('CDI Subtracted Image Plane Intensity')

        for ax, ix in zip(subplot.flatten(), range(n_nulls)):
            im = ax.imshow(I_processed[ix],  # , 250:270, 150:170  *mask2D , irng[0]:irng[-1], jrng[0]:jrng[-1]
                           interpolation='none',
                           norm=LogNorm(),
                           vmin=1e-7, vmax=1e-4
                           )  # , norm=SymLogNorm(linthresh=1e-5))
            ax.set_title(f'Null Step {ix}')

        cax = fig.add_axes([0.9, 0.2, 0.03, 0.6])  # Add axes for colorbar @ position [left,bottom,width,height]
        cb = fig.colorbar(im, orientation='vertical', cax=cax)  #
        cb.set_label('Intensity')

        # =======================================
        # Compare Unprobed to Estimated E-fields
        # ======================================
        fig, subplot = plt.subplots(2, n_nulls, figsize=(16, 12))
        fig.subplots_adjust(wspace=0.5, right=0.85)
        fig.suptitle('Compare Null Probed to Pairwise Probe Estimated Image Plane Intensity')
        ax1, ax2, ax3, ax4 = subplot.flatten()

        ax1.imshow(np.abs(fp_seq[-2, irng[0]:irng[-1], jrng[0]:jrng[-1]]) ** 2,  # , 250:270, 150:170  *mask2D
                   interpolation='none',
                   norm=LogNorm(),
                   # vmin=1e-8, vmax=1e-2
                   )
        ax1.set_title(f'Original Null Step 0')

        ax2.imshow(np.abs(fp_seq[-1, irng[0]:irng[-1], jrng[0]:jrng[-1]]) ** 2,  # , 250:270, 150:170  *mask2D
                   interpolation='none',
                   norm=LogNorm(),
                   # vmin=1e-8, vmax=1e-2
                   )
        ax2.set_title(f'Original Null Step 1')

        ax3.imshow(np.abs(E_est[0, irng[0]:irng[-1], jrng[0]:jrng[-1]]) ** 2,  # , 250:270, 150:170  *mask2D
                   interpolation='none',
                   norm=LogNorm(),
                   # vmin=1e-8, vmax=1e-2
                   )
        ax3.set_title(f'PP Estimated Null Step 0')

        ax4.imshow(np.abs(E_est[1, irng[0]:irng[-1], jrng[0]:jrng[-1]]) ** 2,  # , 250:270, 150:170  *mask2D
                   interpolation='none',
                   norm=LogNorm(),
                   # vmin=1e-8, vmax=1e-2
                   )
        ax4.set_title(f'PP Estimated Null Step 1')

        cax = fig.add_axes([0.9, 0.2, 0.03, 0.6])  # Add axes for colorbar @ position [left,bottom,width,height]
        cb = fig.colorbar(im, orientation='vertical', cax=cax)  #
        cb.set_label('Intensity')

        # =======================================
        # Compare Unprobed to Estimated E-fields
        # ======================================
        fig, subplot = plt.subplots(2, n_nulls, figsize=(16, 12))
        fig.subplots_adjust(wspace=0.5, right=0.85)
        fig.suptitle('Compare Null Probed to Pairwise Probe Estimated Image Plane Intensity')
        ax1, ax2, ax3, ax4 = subplot.flatten()

        ax1.imshow(np.abs(fp_seq[-2]) ** 2,  # , 250:270, 150:170  *mask2D
                   interpolation='none',
                   norm=LogNorm(),
                   # vmin=1e-8, vmax=1e-2
                   )
        ax1.set_title(f'Original Null Step 0')

        ax2.imshow(np.abs(fp_seq[-1]) ** 2,  # , 250:270, 150:170  *mask2D
                   interpolation='none',
                   norm=LogNorm(),
                   # vmin=1e-8, vmax=1e-2
                   )
        ax2.set_title(f'Original Null Step 1')

        ax3.imshow(I_processed[0],  # , 250:270, 150:170  *mask2D
                   interpolation='none',
                   norm=LogNorm(),
                   # vmin=1e-8, vmax=1e-2
                   )
        ax3.set_title(f'CDI Processed Null Step 0')

        ax4.imshow(I_processed[1],  # , 250:270, 150:170  *mask2D
                   interpolation='none',
                   norm=LogNorm(),
                   # vmin=1e-8, vmax=1e-2
                   )
        ax4.set_title(f'CDI Processed Null Step 1')

        cax = fig.add_axes([0.9, 0.2, 0.03, 0.6])  # Add axes for colorbar @ position [left,bottom,width,height]
        cb = fig.colorbar(im, orientation='vertical', cax=cax)  #
        cb.set_label('Intensity')

        # ==================
        # View Time Series
        # ==================
        view_timeseries(cpx_to_intensity(fp_seq[:, 100:300, 300:500]), cdi, title=f"White Light Timeseries",
                        subplt_cols=sp.tseries_cols,
                        logZ=False,
                        vlim=(1e-7, 1e-4),
                        )

        plt.show()

##
def get_fp_mask(cdi, thresh=1e-7):
    """
    returns a mask of the CDI probe pattern in focal plane coordinates

    :param cdi: structure containing all CDI probe parameters
    :param thresh: intensity threshold for determining probed coordinates
    :return: fp_mask: boolean array where True marks the probed coordinates
             imsk, jmsk:
             irng, jrng:

    """
    nx = sp.grid_size
    ny = sp.grid_size
    dm_act = cdi.nact

    fftA = (1 / np.sqrt(2 * np.pi) *
            np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(cdi.DM_probe_series[0]))))

    Ar = interpolate.interp2d(range(dm_act), range(dm_act), fftA.real, kind='cubic')
    Ai = interpolate.interp2d(range(dm_act), range(dm_act), fftA.imag, kind='cubic')
    ArI = Ar(np.linspace(0, dm_act, ny), np.linspace(0, dm_act, nx))
    AiI = Ai(np.linspace(0, dm_act, ny), np.linspace(0, dm_act, nx))

    fp_probe = np.sqrt(ArI**2 + AiI**2)
    # fp_mask = (fp_probe > 1e-7)
    # (imsk, jmsk) = (fp_probe > 1e-7).nonzero()
    fp_mask = (fp_probe > thresh)
    (imsk, jmsk) = (fp_probe > thresh).nonzero()

    irng = range(min(imsk), max(imsk), 1)
    jrng = range(min(jmsk), max(jmsk), 1)

    imx = max(irng)-1  # -1 is to get index values for plotting purposes
    imn = min(irng)-1
    jmx = max(jrng)-1
    jmn = min(jrng)-1

    return fp_mask, imsk, jmsk, irng, jrng, imx, imn, jmx, jmn

##
if __name__ == '__main__':
    dprint(f"Testing CDI probe")
    cdi.use_cdi = True; cdishow_probe = True

    cdi.probe_amp = 2e-6  # [m] probe amplitude, scale should be in units of actuator height limits
    cdi.probe_w = 10  # [actuator coordinates] width of the probe
    cdi.probe_h = 30  # [actuator coordinates] height of the probe
    cdi.probe_shift = [5, 5]  # [actuator coordinates] center position of the probe
    cdi.probe_spacing = 15

    tp.act_tweeter = 49

    sp.numframes = 10
    cdi.gen_phaseseries()
    cdi.init_probes(tp.act_tweeter)
    # cdiphase_series = [-1*np.pi/4]
    config_probe(cdi.phase_series[0], tp.act_tweeter)  #