"""
I'm copying this from some of my much older work. It's probably going to be
buggy. Yep, it was. I've done some refactoring and whatnot but it could probably
use some more. Maybe combine with `GeoDFUtils`? It would be nice to be able to
easily operate on GDF subsets and selections.

`buffer` and `rasterize` are working pretty well. Still need tests and whatnot
but those methods are the main use for this module at the moment. The two
`error_matrix` functions are useful and working too.
"""

#from error_matrix import *
from RasterDS import RasterDS
from ErrorMatrix import ErrorMatrix
from scipy.stats.stats import mode
import numpy as np
import pandas as pd
import geopandas as gpd
import os
from osgeo import ogr, gdal, osr
import shapely as shpl
from scipy.stats import mode as scipymode
from tempfile import mkdtemp
import shutil

class GroundTruthGDF(gpd.GeoDataFrame):
    def __init__(self, *args, **kwargs):
        hf = kwargs.pop('habfield', 'habitat')
        hc = kwargs.pop('habcodefield', 'hab_num')
        super(GroundTruthGDF, self).__init__(*args, **kwargs)
        self.habfield = hf
        self.habcodefld = hc

    @classmethod
    def new(cls,*args,**kwargs):
        return cls(*args,**kwargs)

    @classmethod
    def from_file(cls, filename, **kwargs):
        hf = kwargs.pop('habfield', 'habitat')
        hc = kwargs.pop('habcodefield', 'hab_num')
        gdf = gpd.io.file.read_file(filename, **kwargs)
        return cls(gdf, habfield=hf, habcodefield=hc)

    @property
    def codes_habitat(self):
        """
        Return a dictionary just like habitat_codes only backwards.
        """
        hf = self.habfield
        hcf = self.habcodefld
        hcd = dict()
        for cl in self[hcf].unique():
            if cl > 0:
                sers = self[self[hcf]==cl][hf]
                if sers.count() > 1:
                    hcd[cl] = sers.mode().item()
                elif sers.count() > 0:
                    hcd[cl] = sers.item()
        return hcd

    def __getitem__(self, key):
        result = super(GroundTruthGDF, self).__getitem__(key)
        if isinstance(result, gpd.GeoDataFrame):
            result.__class__ = GroundTruthGDF
            result.habfield = self.habfield
            result.habcodefld = self.habcodefld
        return result

    def query(self, expr, inplace=False, **kwargs):
        result = super(GroundTruthGDF, self).query(expr, inplace=False, **kwargs)
        if isinstance(result, gpd.GeoDataFrame):
            result.__class__ = GroundTruthGDF
            result.habfield = self.habfield
            result.habcodefld = self.habcodefld
        return result

    def comparison_df(self, rds, radius=0, generous=False, band_index=0,
                       out_of_bounds=np.nan, with_unclassed=False):
        """
        There can be problems if there are codes in the raster that do not exist
        in the geodataframe. I should probably check for this condition and
        raise an exception. No time right now.
        """
        pred = self.compare_raster(rds, radius=radius, generous=generous,
                                   band_index=band_index,
                                   out_of_bounds=out_of_bounds)
        truth = self.__getitem__(self.habcodefld)
        truth.name = 'truth'
        pred.name = 'pred'
        preddf = pd.concat((truth, pred), axis=1)
        if not with_unclassed:
            # Get rid of any row that has a zero in it
            preddf = preddf[(preddf!=0).all(1)]
        return preddf

    def error_matrix(self, rds, radius=0, generous=False, band_index=0,
                       out_of_bounds=np.nan, with_unclassed=False):
        from sklearn.metrics import confusion_matrix
        compdf = self.comparison_df(rds, radius=radius, generous=generous,
                                    band_index=band_index,
                                    out_of_bounds=out_of_bounds,
                                    with_unclassed=with_unclassed).dropna()
        # scikit-learn returns pred on x and true on y. I want it the other
        # way around so .T
        em = confusion_matrix(compdf.truth, compdf.pred).T.view(ErrorMatrix)
        codes = np.sort(np.unique(compdf.dropna()))
        em.categories = map(lambda s: self.codes_habitat.get(s, "Unclassified"),
                            codes)
        return em

    def compare_raster(self, rds, radius=0, generous=False, band_index=0,
                       out_of_bounds=np.nan):
        """
        Compare habitat codes in `gdf` with codes in corresponding locations of
        a raster habitat map (`rds`). This can be an exact point to point
        comparison (when `radius`=0) or can be more forgiving. When `radius`>0
        and `generous` is `False`, the mode (most common) value within `radius`
        of each point will be returned. When `radius`>0 and `generous` is True,
        ground truth habitat codes will be returned if found within `radius` of
        each point, and the mode will be returned if not.

        Parameters
        ----------
        rds : OpticalRS.RasterDS
            The habitat map (or whatever raster) you want to compare to the
            `GroundTruthShapefile` (self). The projection of this raster must
            match the projection of the `GroundTruthShapefile`. If it doesn't
            match, you might get results but they'll be wrong.
        radius : float
            The radius with which to buffer `point`. The units of this value
            depend on the projection being used.
        generous : boolean
            If False (default), mode will be returned. If True, habitat code will be
            returned if within `radius`. See function description for more info.
        band_index : int
            Index of the image band to sample. Zero indexed (band 1 = 0). For
            single band rasters, this should be left at the default value (0).
        out_of_bounds : float, int, or nan (default)
            If `point` is not within `self.raster_extent`, `out_of_bounds` will
            be returned.

        Returns
        -------
        pandas Series
            The values from `rds` that correspond to each point in `gdf`.
        """
        column = self.habcodefld
        if generous:
            rcheck = lambda row: rds.radiused_point_check(row.geometry,
                                                          radius=radius,
                                                          search_value=row[column],
                                                          band_index=band_index,
                                                          out_of_bounds=out_of_bounds)
        else:
            rcheck = lambda row: rds.radiused_point_check(row.geometry,
                                                          radius=radius,
                                                          search_value=None,
                                                          band_index=band_index,
                                                          out_of_bounds=out_of_bounds)
        return self.apply(rcheck, axis=1)

class GroundTruthShapefile(object):
    """
    This class contains code for relating point ground truth shapefiles (such as
    the ones generated by Benthic Photo Survey) to raster maps. The default
    values (for `habfield` and `habcodefield`) assume that there's a field
    called habitat that contains a text description of the habitat class for
    each point.

    """
    def __init__(self, file_path, habfield='habitat', habcodefield='hab_num'):
        self.habfield = habfield
        self.habcodefld = habcodefield
        self.file_path = file_path
        self.ds = open_shapefile(self.file_path)
        self.hab_dict = self.__setup_hab_dict()
        self.legit_habs = sorted( [ h for h in self.habitats if h ] ) # Exclude None as a habitat value
        self.habitat_codes = self.__setup_hab_codes() # dict( zip( legit_habs, range( 1, len(legit_habs) + 1 ) ) )

    def __setup_hab_dict(self):
        """
        The hab_dict is a dictionary that contains a list of ogr features for
        each habitat key.
        """
        hab_dict = {}
        for hab in self.habitats:
            hab_dict[hab] = [f for f in self.features if f.__getattr__(self.habfield)==hab]
        return hab_dict

    def __setup_hab_codes(self):
        """
        There should be habitat codes in the shapefile in a field called
        hab_num. We need to get them and set up the matching names. This only
        works for BPS shapefiles with a hab_num field set up to match the
        habitat field. If `self.habfield` is set to something else, we'll just
        generate integer codes.
        """
        # Exclude None from list of habitats
        hcd = {}
        if self.habcodefld is not None:
            for hab in self.legit_habs:
                feat = self.hab_dict[hab][0]
                hcd[hab] = feat.__getattr__(self.habcodefld)
        else:
            for i, hab in enumerate(self.legit_habs):
                hcd[hab] = i+1 # +1 to make it not zero indexed
        return hcd

    @property
    def features(self):
        fts = [f for f in self.ds.GetLayer()]
        self.ds.GetLayer().ResetReading()
        return fts

    @property
    def habitats(self):
        habs = sorted( set([f.__getattr__(self.habfield) for f in self.features]))
        return habs

    @property
    def legit_habs_code_sorted(self):
        """
        Return the legit habitats sorted by order of their numeric codes.
        """
        return [v for k,v in sorted(self.codes_habitat.items())]

    @property
    def geo_data_frame(self):
        """
        Return a GeoPandas GeoDataFrame object.
        """
        gtgdf = GroundTruthGDF.from_file(self.file_path, habfield=self.habfield,
                                         habcodefield=self.habcodefld)
        # gtgdf = gpd.GeoDataFrame.from_file(self.file_path)
        return gtgdf

    def geopandas_subset(self, query, file_name=None):
        """
        Create a `GroundTruthShapefile` based on a geopandas subset of
        `self.geo_data_frame`. If `file_name` is `None` (default), then the file
        will only be temporarily saved. It will be deleted before this function
        returns. This seems to work fine for generating error matrices from
        subsets but it could have unintended consequences elsewhere. If you
        provide a `file_name`, a shapefile will be saved from the output.

        Parameters
        ----------
        query : string or pandas Series
            If `query` is a string, `pandas.DataFrame.query` will be used to
            generate the subset. Otherwise, query is assumed to be a series that
            can be used to index `self.geo_data_frame`.
        file_name : string file path or None
            If `None`, a temporary shapefile will be created and immediately
            deleted. Otherwise, the subset will be saved as a shapefile.

        Returns
        -------
        GroundTruthShapefile
            A `GroundTruthShapefile` object containing only the selected subset
            of features.
        """
        if file_name is None:
            tdir = mkdtemp()
            tfn = os.path.join(tdir, 'temp.shp')
        else:
            tfn = file_name
        if type(query) is str:
            gdf = self.geo_data_frame.query(query)
        else:
            gdf = self.geo_data_frame[query]
        # save the subset to a file
        gdf.to_file(tfn)
        # make a new GroundTruthShapefile
        gts = GroundTruthShapefile(tfn, self.habfield, self.habcodefld)
        if file_name is None:
            shutil.rmtree(tdir)
        return gts

    @property
    def spatial_reference(self):
        """
        Return the OGR spatial reference object for the shapefile.
        """
        return self.ds.GetLayer().GetSpatialRef()

    @property
    def projection_wkt(self):
        """
        Return the well known text (WKT) representation of the shapefile's projection.
        """
        return self.spatial_reference.ExportToWkt()

    @property
    def projcs(self):
        """
        Return the PROJCS value from the shapefile's spatial reference. This is
        basically the name of the projection. ...I think.
        """
        return self.spatial_reference.GetAttrValue('PROJCS')

    @property
    def geometry_type(self):
        """
        Just return whether it's a type of point, line, or polygon.
        """
        type_name = ogr.GeometryTypeToName( self.ds.GetLayer().GetGeomType() ).lower()
        if type_name.find('point') <> -1:
            return 'point'
        elif type_name.find('line') <> -1:
            return 'line'
        elif type_name.find('polygon') <> -1:
            return 'polygon'
        else:
            return None

    @property
    def hab_colors(self):
        """
        return a dictionary with hab codes as keys and hab colors as values.
        """
        legit_habs = sorted( [ h for h in self.habitats if h ] )
        hcd = {}
        for hab in legit_habs:
            feat = self.hab_dict[hab][0]
            hcd[hab] = feat.hab_color
        return hcd

    @property
    def codes_habitat(self):
        """
        Return a dictionary just like habitat_codes only backwards.
        """
        chd = {}
        for k,v in self.habitat_codes.items():
            chd[v] = k
        return chd

    @property
    def qgis_vector(self):
        qvl = QgsVectorLayer(self.file_path,'grnd_truth','ogr')
        if qvl.isValid():
            return qvl
        else:
            raise Exception("Failed to create a QGis Vector Layer. QGis provider path problems, perhaps?")

    def buffer(self, radius=1.0, file_path=None):
        """
        Buffer the geometries in `self` and return a new `ogr` datasource. If
        `file_path` is `None`, just create the datasource in memory. If a file
        path is given, write out a shapefile. All fields and values (aside from
        geometry) are cloned.
        """
        if file_path == None:
            drvname = 'Memory'
        else:
            drvname = 'ESRI Shapefile'
        srcds = self.ds
        # get projection
        lyr = srcds.GetLayer(0)
        sptrf = lyr.GetSpatialRef()
        proj = osr.SpatialReference()
        proj.ImportFromWkt(sptrf.ExportToWkt())
        drv = ogr.GetDriverByName(drvname)
        if file_path == None:
            dst_ds = drv.CreateDataSource('out')
        elif os.path.exists(file_path):
            raise Exception("{} already exists!".format(file_path))
        else:
            dst_ds = drv.CreateDataSource(file_path)
        dst_lyr = dst_ds.CreateLayer('', srs=proj, geom_type=ogr.wkbPolygon)
        # copy all the fields to the destination ds
        featr = lyr.GetFeature(0)
        nfields = featr.GetFieldCount()
        for i in range(nfields):
            fld = featr.GetFieldDefnRef(i)
            dst_lyr.CreateField(fld)
        feat_defn = dst_lyr.GetLayerDefn()
        # reset the feature counter
        lyr.ResetReading()
        # buffer the geometries and copy the fields
        for i in range(lyr.GetFeatureCount()):
            # get the feature and geometry
            feat = lyr.GetFeature(i)
            geom = feat.GetGeometryRef()
            # create a new feature
            newfeat = feat.Clone()
            # get the buffered geometry
            bufgeom = geom.Buffer(radius)
            # set the new geometry to the buffered geom
            newfeat.SetGeometry(bufgeom)
            # add the new feature to the destination layer
            dst_lyr.CreateFeature(newfeat)
            # clean up
            newfeat.Destroy()
            feat.Destroy()
        # ensure the new features are written
        dst_lyr.SyncToDisk()
        return dst_ds

    def rasterize(self, buffer_radius=None, raster_template=None,
                  pixel_size=1.99976, value_field='hab_num', float_values=False,
                  array_only=False, out_file_path=None):
        """
        Return a raster that can be used for classification training.

        buffer_radius: A float value in projection units to buffer the
        geometries by. If buffer_radius is left None then only pixels right
        under points will be classified.

        raster_template: A RasterDS object. If supplied, the resulting
        rasterized image will have the same extent and geotransform as the
        template. Also, if a raster_template is provided, the pixel_size keyword
        value will be ignored and pixel size will come from the template.

        pixel_size: A float value representing pixel size in projection units.
        This value will be ignored if a raster_template is supplied.

        value_field: A string representing the name of the field in the
        shapefile that holds the numeric code that will be burned into the
        raster output as the pixel value.

        float_values: Boolean. If `True`, the output raster will contain floats.
        If `False`, the output will be integers. Default is `False`.

        array_only: A boolean. If true we'll try to just write the raster to
        memory and not to disk. If you don't need to keep the raster, this will
        just keep you from having to clean up useless files later. Then we'll
        just return an array instead of GroundTruthRaster object.

        out_file_path: String. Path to the raster file output. If `None`
        (default) and `array_only=False`, a file name based on the
        `GroundTruthShapefile` file name will be created. If `array_only=True`,
        `out_file_path` is ignored.
        """
        if float_values:
            datatype = gdal.GDT_Float32
        else:
            datatype = gdal.GDT_Byte
        # Make a copy of the layer's data source because we'll need to
        # modify its attributes table
        if buffer_radius:
            source_ds = ogr.GetDriverByName("Memory").CopyDataSource( self.buffer(radius=buffer_radius), "" )
        else:
            source_ds = ogr.GetDriverByName("Memory").CopyDataSource( self.ds, "")
        source_layer = source_ds.GetLayer(0)
        source_srs = source_layer.GetSpatialRef()

        if raster_template:
            gTrans = raster_template.gdal_ds.GetGeoTransform()
            pixsizeX = gTrans[1]
            pixsizeY = gTrans[5]
            x_res = raster_template.gdal_ds.RasterXSize
            y_res = raster_template.gdal_ds.RasterYSize
            rdsarr = raster_template.band_array
            # if np.ma.is_masked(rdsarr):
            #     mask = rdsarr[...,0].mask
            # else:
            #     mask = None
        else:
            x_min, x_max, y_min, y_max = source_layer.GetExtent()
            # Create the destination data source
            x_res = int((x_max - x_min) / pixel_size)
            y_res = int((y_max - y_min) / pixel_size)

        if out_file_path:
            targ_fn = out_file_path
        else:
            # make a target ds with filename based on source filename
            targ_fn = self.file_path.rsplit(os.path.extsep, 1)[0] + '_rast' + os.path.extsep + 'tif'
        # print "x_res: %i, y_res: %i" % (x_res,y_res)
        target_ds = gdal.GetDriverByName('GTiff').Create(targ_fn, x_res, y_res, 1, datatype)

        if raster_template:
            # Use the raster template supplied so that we get the same extent as the raster
            # we're trying to classify
            target_ds.SetGeoTransform( gTrans )
        else:
            # None supplied so use the pixel_size value and the extent of the shapefile
            target_ds.SetGeoTransform(( x_min, pixel_size, 0, y_max, 0, -pixel_size, ))
        if raster_template:
            target_ds.SetProjection( raster_template.gdal_ds.GetProjection() )
        elif source_srs:
            # Make the target raster have the same projection as the source
            target_ds.SetProjection(source_srs.ExportToWkt())
        else:
            # Source has no projection (needs GDAL >= 1.7.0 to work)
            target_ds.SetProjection('LOCAL_CS["arbitrary"]')
        # Rasterize
        err = gdal.RasterizeLayer(target_ds, [1], source_layer,
                burn_values=[0],
                options=["ATTRIBUTE=%s" % value_field])
        if err != 0:
            raise Exception("error rasterizing layer: %s" % err)
        # clean up
        source_layer = None
        source_srs = None
        source_ds = None

        if array_only:
            out_array = target_ds.ReadAsArray()
            target_ds = None
            os.remove( targ_fn )
            return out_array
        else:
            target_ds = None
            return RasterDS(targ_fn)

    def error_matrix(self, classification_ds, with_unclassed=False):
        """
        Take a RasterDS (classification_ds) and create a user / producer
        accuracy table. Return as an array so it can be displayed in multiple
        ways. See the `ErrorMatrix` module for more information on the returned
        object.

        Parameters
        ----------
        classification_ds : OpticalRS.RasterDS
            The habitat map (or whatever raster) you want to compare to the
            `GroundTruthShapefile` (self). The projection of this raster must
            match the projection of the `GroundTruthShapefile`. If it doesn't
            match, you might get results but they'll be wrong.

        Returns
        -------
        ErrorMatrix
            See the `ErrorMatrix` module for more information on the returned
            object.

        Notes
        -----
        This function should be merged in some way with `error_matrix_buffered`.
        There's a bunch of redundancy between the two. I don't have time to do
        it right now.
        """
        maxcode = max(self.habitat_codes.values())
        if with_unclassed:
            maxcode += 1
        errmat = np.zeros((maxcode, maxcode), int)
        cats = list()
        rext = classification_ds.raster_extent
        for hab,code in self.habitat_codes.items():
            for feature in self.hab_dict[hab]:
                ref_val = code
                geom = feature.geometry()
                pnt = shpl.geometry.base.geom_from_wkb(geom.ExportToWkb())
                if pnt.within(rext):
                    cls_val = classification_ds.value_at_point( geom )
                else:
                    # this means that the point is not within the raster
                    # I think that means we don't want to count this point at
                    # all in the accuracy assessment.
                    continue
                if with_unclassed:
                    errmat[ cls_val ][ ref_val ] += 1
                elif cls_val == 0:
                    # If we're not including unclassified values
                    # we don't want this showing up in the totals.
                    continue
                else:
                    errmat[ cls_val - 1 ][ ref_val - 1 ] += 1
        # Get rid of all zero rows and columns. This can happen if hab codes
        # skip an integer.
        em = errmat.view( ErrorMatrix ).clean_zeros(with_unclassed)
        # Rows and Columns of errmat end up sorted by hab code. This next line
        # will give the habitat names sorted by hab code number.
        if with_unclassed:
            em.categories = ['Unclassified'] + sorted(self.habitat_codes, key=self.habitat_codes.get)
        else:
            em.categories = sorted(self.habitat_codes, key=self.habitat_codes.get)
        return em

    def error_matrix_buffered(self, classification_ds, radius=2.0, with_unclassed=False):
        """
        Take a RasterDS (classification_ds) and create a user / producer
        accuracy table. Ground Truth points will be buffered and matching
        habitat codes within `radius` of a point will be considered success.
        Return as an array so it can be displayed in multiple ways. See the
        `ErrorMatrix` module for more information on the returned object.

        Parameters
        ----------
        classification_ds : OpticalRS.RasterDS
            The habitat map (or whatever raster) you want to compare to the
            `GroundTruthShapefile` (self). The projection of this raster must
            match the projection of the `GroundTruthShapefile`. If it doesn't
            match, you might get results but they'll be wrong.
        radius : float
            The radius with which to buffer points. The units of this value
            depend on the projection being used. You can use
            `GroundTruthShapefile.projection_wkt` to examine the projection and
            find the units.

        Returns
        -------
        ErrorMatrix
            See the `ErrorMatrix` module for more information on the returned
            object.
        """
        maxcode = max(self.habitat_codes.values())
        if with_unclassed:
            maxcode += 1
        errmat = np.zeros((maxcode, maxcode), int)
        cats = list()
        rext = classification_ds.raster_extent
        for hab,code in self.habitat_codes.items():
            for feature in self.hab_dict[hab]:
                ref_val = code
                geom = feature.geometry()
                pnt = shpl.geometry.base.geom_from_wkb(geom.ExportToWkb())
                if pnt.within(rext):
                    clsarr = classification_ds.geometry_subset(pnt.buffer(radius),
                                                               all_touched=True)
                else:
                    # this means that the point is not within the raster
                    # I think that means we don't want to count this point at
                    # all in the accuracy assessment.
                    continue

                if ref_val in clsarr.compressed():
                    cls_val = ref_val # this counts as success
                elif not pnt.within(rext):
                    # this means that the point is not within the raster
                    # I think that means we don't want to count this point at
                    # all in the accuracy assessment.
                    continue
                else:
                    # our reference value was not found within radius of point
                    # so we'll report it as the most common class within radius
                    if len(clsarr.compressed()) == 0:
                        cls_val = 0 # Assuming zero is code for unclassified
                    else:
                        cls_val = scipymode(clsarr.compressed()).mode.item()
                if with_unclassed:
                    errmat[ cls_val ][ ref_val ] += 1
                elif cls_val == 0:
                    # If we're not including unclassified values
                    # we don't want this showing up in the totals.
                    continue
                else:
                    errmat[ cls_val - 1 ][ ref_val - 1 ] += 1
        # Get rid of all zero rows and columns. This can happen if hab codes
        # skip an integer.
        em = errmat.view( ErrorMatrix ).clean_zeros(with_unclassed)
        # Rows and Columns of errmat end up sorted by hab code. This next line
        # will give the habitat names sorted by hab code number.
        if with_unclassed:
            em.categories = ['Unclassified'] + sorted(self.habitat_codes, key=self.habitat_codes.get)
        else:
            em.categories = sorted(self.habitat_codes, key=self.habitat_codes.get)
        return em

    @property
    def hab_dict_counts(self):
        ret_dict = {}
        for hab in self.habitats:
            ret_dict[hab] = len( self.hab_dict[hab] )
        return ret_dict

    def add_raster_values(self, raster_ds):
        """
        The raster data source here is assumed to be a classified image. The raster
        values should correspond to classes.
        """
        trans = transform_dict(raster_ds)
        band = raster_ds.GetRasterBand(1)
        self.features = [ add_raster_value(f,trans,band) for f in self.ds.GetLayer() ]
        self.ds.GetLayer().ResetReading()
        self.hab_dict = self.__setup_hab_dict()

    @property
    def unsupervised_habitat_class_dict(self):
        """
        For each habitat, give a list of raster values that correspond to the ground truth
        points of that habitat type. This will be used with unsupervised classifications to
        figure out which, if any, of the classes correspond to particular habitat types.
        """
        try:
            hcd = {}
            for hab in self.habitats:
                hcd[hab] = [ f.raster_value for f in self.hab_dict[hab] ]
        except AttributeError:
            raise AttributeError("Features need to be assigned raster values before you can create a habitat class dictionary.")
        return hcd

    @property
    def unsupervised_habitat_class_modes(self):
        hcm = {}
        for hab in self.habitats:
            md, cn = mode( self.unsupervised_habitat_class_dict[hab] )
            if len( md )==1:
                hcm[hab] = md[0]
            else:
                hcm[hab] = None
        return hcm

    def __output_training_LAN(self,img,buffer_radius=3.5,driver_str='LAN'):
        """
        DEPRICATED! -> This only works for points. I think I can use the
        rasterize method instead. I need to verify and then get rid of this
        method. This method also has the habitat field hard coded (search for
        feat.habitat). That would need to be changed to
        feat.__getattr__(self.habfield) to make this work correctly.

        Create a raster input for supervised classifications. img is the image
        that we want to classify (in the form of a gdal datasource). Spectral
        can't use tifs so we will create LAN file.

        A buffer radius of 3.5 meters gives us 3 x 3 sets of pixels with our
        point feature in the center. This, of course, assumes that we're dealing
        with WV2 imagery and a projection with meters as the units. This works
        for me on my project but might now work for others.
        """
        if driver_str=='LAN':
            f_ext = 'lan'
        elif driver_str=='GTiff':
            f_ext = 'tif'
        else:
            raise ValueError("At this point, the output_training_LAN method only knows how to deal with LAN and GTiff file types. Sorry.")

        lyr = self.ds.GetLayer()
        lyr.ResetReading()
        trans = transform_dict(img)
        driver = gdal.GetDriverByName(driver_str)
        rows = img.RasterYSize
        cols = img.RasterXSize
        fname = img.GetDescription().rsplit(os.path.extsep)[0] + '_train' + os.path.extsep + f_ext
        add_num = 0
        while os.path.exists(fname):
            add_num += 1
            if add_num==1:
                fname = fname.replace( os.path.extsep + f_ext, '_%i' % add_num + os.path.extsep + f_ext )
            else:
                old = '_%i.%s' % ( add_num - 1, f_ext )
                new = '_%i.%s' % ( add_num, f_ext )
                fname = fname.replace( old, new )
        outDs = driver.Create(fname, cols, rows, 1, GDT_Int16)
        if outDs is None:
            print 'Could not create %s' % fname
            sys.exit(1)

        outBand = outDs.GetRasterBand(1)

        pixel_count = 0
        hab_pix_count = dict( zip( [h for h in self.habitats if h], np.zeros( len([h for h in self.habitats if h]), dtype=np.int ) ) )
        for feat in lyr:
            if not feat.habitat:
                continue
            if self.hab_dict_counts[feat.habitat] < 24:
                continue
            if buffer_radius:
                geom = feat.geometry().Buffer(buffer_radius)
                elp = envelope_dict(geom)
                xtop = elp['xLeft']
                ytop = elp['yTop']
                xOffset = int( (xtop - trans['originX']) / trans['pixWidth'] )
                yOffset = int( (ytop - trans['originY']) / trans['pixHeight'] )
                xdist = elp['xRight'] - elp['xLeft']
                ydist = elp['yBottom'] - elp['yTop']
                cols = int( xdist / trans['pixWidth'] )
                rows = int( ydist / trans['pixHeight'] )
                pixarr = int( self.habitat_codes[feat.habitat] ) * np.ones((rows,cols), dtype=np.int16)
            else:
                geom = feat.geometry()
                xOffset = int( (geom.GetX() - trans['originX']) / trans['pixWidth'] )
                yOffset = int( (geom.GetY() - trans['originY']) / trans['pixHeight'] )
                pixarr = np.array( [[ self.habitat_codes[feat.habitat] ]] )

            outBand.WriteArray(pixarr,xOffset,yOffset)
            pixel_count += pixarr.size
            hab_pix_count[feat.habitat] += pixarr.size

        outBand.FlushCache()
        outBand.SetNoDataValue(0)
        # georeference the image and set the projection
        outDs.SetGeoTransform(img.GetGeoTransform())
        outDs.SetProjection(img.GetProjection())

        # build pyramids
        gdal.SetConfigOption('HFA_USE_RRD', 'YES')
        outDs.BuildOverviews(overviewlist=[2,4,8,16,32,64,128])

        print "%i pixels total" % pixel_count
        for hab in self.habitats:
            if hab:
                print "%i pixels for %s" % ( hab_pix_count[hab], hab )

        return GroundTruthRaster( outDs.GetDescription() )

    def training_classes(self, rds, buffer_radius=None,calc_stats=0):
        """
        I think I should move some of this functionality over to the GroundTruthRaster class
        in common.py.  I'm generating classes okay from what I can tell but I get a singular
        matrix error when I try to run the Gaussian Classifier. I have no idea why. Baffled,
        I am.
        """
        grnd_truth = self.rasterize(buffer_radius=buffer_radius,raster_template=rds,array_only=True)
        sp_img = rds.spy_image.load()
        return sp.create_training_classes(sp_img, grnd_truth,calc_stats=calc_stats)

def add_raster_value(feature, trans, band ):
    geom = feature.geometry()
    x = geom.GetX()
    y = geom.GetY()

    xOffset = int( (x - trans['originX']) / trans['pixWidth'] )
    yOffset = int( (y - trans['originY']) / trans['pixHeight'] )

    data = band.ReadAsArray(xOffset, yOffset, 1, 1)
    feature.raster_value = data[0,0]
    return feature

def open_shapefile(filename):
    """Take a file path string and return an ogr shape"""
    # open the shapefile and get the layer
    driver = ogr.GetDriverByName('ESRI Shapefile')
    shp = driver.Open(filename)
    if shp is None:
        print 'Could not open %s' % filename
        sys.exit(1)
    return shp


"""
Copyright 2013 Steven Diamond

This file is part of CVXPY.

CVXPY is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

CVXPY is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with CVXPY.  If not, see <http://www.gnu.org/licenses/>.
"""

import cvxpy.settings as s
from cvxpy.problems.solvers.ecos_intf import ECOS
import ecos

class ECOS_BB(ECOS):
    """An interface for the ECOS BB solver.
    """

    # Solver capabilities.
    LP_CAPABLE = True
    SOCP_CAPABLE = True
    SDP_CAPABLE = False
    EXP_CAPABLE = False
    MIP_CAPABLE = True

    def name(self):
        """The name of the solver.
        """
        return s.ECOS_BB

    @staticmethod
    def _noncvx_id_to_idx(dims, var_offsets, var_sizes):
        """Converts the nonconvex constraint variable ids in dims into indices.

        Parameters
        ----------
        dims : dict
            The dimensions of the cones.
        var_offsets : dict
            A dict of variable id to horizontal offset.
        var_sizes : dict
            A dict of variable id to variable dimensions.

        Returns
        -------
        tuple
            A list of indices for the boolean variables and integer variables.
        """
        bool_idx = []
        int_idx = []
        for indices, constr_type in zip([bool_idx, int_idx],
                                        [s.BOOL_IDS, s.INT_IDS]):
            for var_id in dims[constr_type]:
                offset = var_offsets[var_id]
                size = var_sizes[var_id]
                for i in range(size[0]*size[1]):
                    indices.append(offset + i)
            del dims[constr_type]

        return bool_idx, int_idx

    def get_problem_data(self, objective, constraints, cached_data):
        """Returns the argument for the call to the solver.

        Parameters
        ----------
        objective : LinOp
            The canonicalized objective.
        constraints : list
            The list of canonicalized cosntraints.
        cached_data : dict
            A map of solver name to cached problem data.

        Returns
        -------
        dict
            The arguments needed for the solver.
        """
        data = super(ECOS_BB, self).get_problem_data(objective, constraints,
                                                     cached_data)
        sym_data = self.get_sym_data(objective, constraints, cached_data)
        bool_idx, int_idx = self._noncvx_id_to_idx(data[s.DIMS],
                                                   sym_data.var_offsets,
                                                   sym_data.var_sizes)
        data[s.BOOL_IDX] = bool_idx
        data[s.INT_IDX] = int_idx
        return data

    def solve(self, objective, constraints, cached_data,
              warm_start, verbose, solver_opts):
        """Returns the result of the call to the solver.

        Parameters
        ----------
        objective : LinOp
            The canonicalized objective.
        constraints : list
            The list of canonicalized cosntraints.
        cached_data : dict
            A map of solver name to cached problem data.
        warm_start : bool
            Not used.
        verbose : bool
            Should the solver print output?
        solver_opts : dict
            Additional arguments for the solver.

        Returns
        -------
        tuple
            (status, optimal value, primal, equality dual, inequality dual)
        """
        data = self.get_problem_data(objective, constraints, cached_data)
        # Default verbose to false for BB wrapper.
        mi_verbose = solver_opts.get('mi_verbose', False)
        results_dict = ecos.solve(data[s.C], data[s.G], data[s.H],
                                  data[s.DIMS], data[s.A], data[s.B],
                                  verbose=verbose,
                                  mi_verbose=mi_verbose,
                                  bool_vars_idx=data[s.BOOL_IDX],
                                  int_vars_idx=data[s.INT_IDX],
                                  **solver_opts)
        return self.format_results(results_dict, None,
                                   data[s.OFFSET], cached_data)


"""Estimation functions for WFSC."""

import numpy as np
import multiprocessing
# from astropy.io import fits
import matplotlib.pyplot as plt
import falco
from . import check

def perfect(mp):
    """
    Return the perfect-knowledge E-field from the full model.

    Optionally add Zernikes at the input pupil.

    Parameters
    ----------
    mp : ModelParameters
        Structure containing optical model parameters

    Returns
    -------
    Emat : numpy ndarray
        2-D array with the vectorized, complex E-field of the dark hole pixels
        for each mode included in the control Jacobian.
    """
    if type(mp) is not falco.config.ModelParameters:
        raise TypeError('Input "mp" must be of type ModelParameters')

    if mp.flagParallel:

        Emat = np.zeros((mp.Fend.corr.Npix, mp.jac.Nmode), dtype=complex)

        # Loop over all modes and wavelengths
        inds_list = [(x, y) for x in range(mp.jac.Nmode) for y in range(mp.Nwpsbp)]
        Nvals = mp.jac.Nmode*mp.Nwpsbp

        pool = multiprocessing.Pool(processes=mp.Nthreads)
        resultsRaw = [pool.apply_async(_est_perfect_Efield_with_Zernikes_in_parallel,
                                       args=(mp, ilist, inds_list)) for ilist in range(Nvals)]
        results = [p.get() for p in resultsRaw]  # All the images in a list
        pool.close()
        pool.join()

        # Re-order for easier indexing
        Ecube = np.zeros((mp.Fend.corr.Npix, mp.jac.Nmode, mp.Nwpsbp), dtype=complex)
        for iv in range(Nvals):
            im = inds_list[iv][0]  # Index of the Jacobian mode
            wi = inds_list[iv][1]   # Index of the wavelength in the sub-bandpass
            Ecube[:, im, wi] = results[iv]
        Emat = np.mean(Ecube, axis=2)  # Average over wavelengths in the subband

#        EmatAll = np.zeros((mp.Fend.corr.Npix, Nval))
#        for iv in range(Nval):
#            EmatAll[:, iv] = results[iv]
#
#        counter = 0;
#        for im=1:mp.jac.Nmode
#            EsbpMean = 0;
#            for wi=1:mp.Nwpsbp
#                counter = counter + 1;
#                EsbpMean = EsbpMean + EmatAll(:,counter) * \
#                     mp.full.lambda_weights(wi);
#            end
#            Emat(:,im) = EsbpMean;
#        end

    else:

        Emat = np.zeros((mp.Fend.corr.Npix, mp.jac.Nmode), dtype=complex)
        modvar = falco.config.Object()

        for im in range(mp.jac.Nmode):
            modvar.sbpIndex = mp.jac.sbp_inds[im]
            modvar.zernIndex = mp.jac.zern_inds[im]
            modvar.whichSource = 'star'

            # Take the mean over the wavelengths within the sub-bandpass
            EmatSbp = np.zeros((mp.Fend.corr.Npix, mp.Nwpsbp), dtype=complex)
            for wi in range(mp.Nwpsbp):
                modvar.wpsbpIndex = wi
                E2D = falco.model.full(mp, modvar)
                # Actual field in estimation area. Apply spectral weight
                # within the sub-bandpass
                EmatSbp[:, wi] = mp.full.lambda_weights[wi] * \
                    E2D[mp.Fend.corr.maskBool]
            Emat[:, im] = np.sum(EmatSbp, axis=1)

    return Emat

# Extra function needed to use parfor (because parfor can have only a
# single changing input argument).
def _est_perfect_Efield_with_Zernikes_in_parallel(mp, ilist, inds_list):

    im = inds_list[ilist][0]  # Index of the Jacobian mode
    wi = inds_list[ilist][1]   # Index of the wavelength in the sub-bandpass

    modvar = falco.config.Object()
    modvar.sbpIndex = mp.jac.sbp_inds[im]
    modvar.zernIndex = mp.jac.zern_inds[im]
    modvar.wpsbpIndex = wi
    modvar.whichSource = 'star'

    E2D = falco.model.full(mp, modvar)

    # Actual field in estimation area. Don't apply spectral weight here.
    return E2D[mp.Fend.corr.maskBool]

def pairwise_probing(mp, ev, jacStruct=np.array([])):
    """
    Estimate the dark hole E-field with pair-wise probing.

    Parameters
    ----------
    mp : falco.config.ModelParameter
        Object containing all model parameters.
    ev : falco.config.Object()
    jacStruct : array_like, optional
        Array containing the control Jacobian. Default is an empty array.

    Returns
    -------
    None
        Outputs are included in the object ev.

    """
    # "ev" is passed in only for the Kalman filter. Reset it for the batch
    # process to avoid accidentally using old data.
    # if 'pwp-bp' == mp.estimator.lower():
    #     ev = falco.config.Object()

    # Select number of actuators across based on chosen DM for the probing
    if mp.est.probe.whichDM == 1:
        Nact = mp.dm1.Nact
    elif mp.est.probe.whichDM == 2:
        Nact = mp.dm2.Nact
    else:
        raise ValueError('mp.est.probe.whichDM must equal 1 or 2.')

    # Store the initial DM commands
    if np.any(mp.dm_ind == 1):
        DM1Vnom = mp.dm1.V

    if np.any(mp.dm_ind == 2):
        DM2Vnom = mp.dm2.V
    else:
        DM2Vnom = np.zeros_like(mp.dm1.V)

    # Definitions:
    Npairs = mp.est.probe.Npairs  # Number of image PAIRS
    ev.Icube = np.zeros((mp.Fend.Neta, mp.Fend.Nxi, 1+2*Npairs))
    if np.any(mp.dm_ind == 1):
        ev.Vcube1 = np.zeros((mp.dm1.Nact, mp.dm1.Nact, 1+2*Npairs))
    if np.any(mp.dm_ind == 2):
        ev.Vcube2 = np.zeros((mp.dm2.Nact, mp.dm2.Nact, 1+2*Npairs))

    # Generate evenly spaced probes along the complex unit circle
    # NOTE: Nprobes=Npairs*2
    probePhaseVec = np.array([0, Npairs])
    for k in range(Npairs-1):
        probePhaseVec = np.append(probePhaseVec, probePhaseVec[-1]-(Npairs-1))
        probePhaseVec = np.append(probePhaseVec, probePhaseVec[-1]+Npairs)
    probePhaseVec = probePhaseVec*np.pi/(Npairs)

    badAxisVec = ''
    if mp.est.probe.axis.lower() == 'y':
        for _iter in range(2*Npairs):
            badAxisVec += 'y'
    elif mp.est.probe.axis.lower() == 'x':
        for _iter in range(2*Npairs):
            badAxisVec += 'x'
    elif mp.est.probe.axis.lower() in ('alt', 'xy', 'alternate'):
        for iPair in range(2*Npairs):
            if (iPair+1) % 4 == 1 or (iPair+1) % 4 == 2:
                badAxisVec += 'x'
            elif (iPair+1) % 4 == 3 or (iPair+1) % 4 == 0:
                badAxisVec += 'y'
    elif mp.est.probe.axis.lower() in ('m', 'multi'):
        for _iter in range(2*Npairs):
            badAxisVec += 'm'
    else:
        raise ValueError('Incorrect value for mp.est.probe.axis')

    # Initialize output arrays
    ev.Eest = np.zeros((mp.Fend.corr.Npix, mp.Nsbp), dtype=complex)
    ev.IincoEst = np.zeros((mp.Fend.corr.Npix, mp.Nsbp))
    ev.I0mean = 0
    ev.IprobedMean = 0

    # Get images and perform estimates in each sub-bandpass
    print('Estimating electric field with batch process estimation ...')

    for si in range(mp.Nsbp):
        print('Wavelength: %u/%u ... ' % (si, mp.Nsbp-1))

        # Valid for all calls to model_compact.m:
        modvar = falco.config.Object()  # Initialize the new structure
        modvar.sbpIndex = si
        modvar.whichSource = 'star'

        # Measure current contrast level average
        # Reset DM commands to the unprobed state:
        mp.dm1.V = DM1Vnom
        mp.dm2.V = DM2Vnom
        # Separate out image values at DH pixels and delta DM voltage settings
        Iplus = np.zeros((mp.Fend.corr.Npix, Npairs))
        Iminus = np.zeros((mp.Fend.corr.Npix, Npairs))
        DM1Vplus = np.zeros((Nact, Nact, Npairs))
        DM1Vminus = np.zeros((Nact, Nact, Npairs))
        DM2Vplus = np.zeros((Nact, Nact, Npairs))
        DM2Vminus = np.zeros((Nact, Nact, Npairs))

        # Compute probe shapes and take probed images:

        # Take initial, unprobed image (for unprobed DM settings).
        whichImg = 1
        I0 = falco.imaging.get_sbp_image(mp, si)
        I0vec = I0[mp.Fend.corr.maskBool]  # Vectorize the correction region
        ev.I0mean = ev.I0mean+I0/mp.Nsbp  # Getting Inorm for whole bandpass

        # Store values for first image and its DM commands
        ev.Icube[:, :, whichImg] = I0
        if np.any(mp.dm_ind == 1):
            ev.Vcube1[:, :, whichImg] = mp.dm1.V
        if np.any(mp.dm_ind == 2):
            ev.Vcube2[:, :, whichImg] = mp.dm2.V

        # Compute the average Inorm in the scoring and correction regions
        ev.InormScore = np.mean(I0[mp.Fend.score.maskBool])
        ev.InormCorr = np.mean(I0[mp.Fend.corr.maskBool])
        print('Measured unprobed Inorm (Corr / Score): %.2e \t%.2e \n' %
              (ev.InormCorr, ev.InormScore))

        # Set (approximate) probe intensity based on current measured Inorm
        if mp.flagFiber:
            ev.InormProbeMax = 1e-5
            InormProbe = np.min([np.sqrt(np.max(I0)*1e-8), ev.InormProbeMax])
        else:
            ev.InormProbeMax = 1e-4
            InormProbe = np.min([np.sqrt(np.max(I0vec)*1e-5),
                                 ev.InormProbeMax])
            # Change this to a high percentile value (e.g., 90%) instead of the
            # max to avoid being tricked by noise
        print('Chosen probe intensity: %.2e' % InormProbe)

        # Perform the probing
        iOdd = 0  # Initialize index counters
        iEven = 0
        for iProbe in range(2*Npairs):

            # Generate the command map for the probe
            probeCmd = gen_pairwise_probe(mp, InormProbe,
                                probePhaseVec[iProbe], badAxisVec[iProbe])

            # Select which DM to use for probing. Allocate probe to that DM
            if mp.est.probe.whichDM == 1:
                dDM1Vprobe = probeCmd/mp.dm1.VtoH  # Now in volts
                dDM2Vprobe = 0
            elif mp.est.probe.whichDM == 2:
                dDM1Vprobe = 0
                dDM2Vprobe = probeCmd/mp.dm1.VtoH  # Now in volts
            else:
                raise ValueError('DM for probing must be 1 or 2.')

            if np.any(mp.dm_ind == 1):
                mp.dm1.V = DM1Vnom + dDM1Vprobe
            if np.any(mp.dm_ind == 2):
                mp.dm2.V = DM2Vnom + dDM2Vprobe

            # Take probed image
            if mp.flagFiber:
                Im = falco.imaging.get_sbp_image_fiber(mp, si)
            else:
                Im = falco.imaging.get_sbp_image(mp, si)

            # plt.imshow(np.log10(Im)); plt.title('Probed Image %d' % iProbe); plt.colorbar(); plt.pause(1e-2);

            ImNonneg = Im
            ImNonneg[Im < 0] = 0
            whichImg = 1+iProbe  # Increment image counter
            # Inorm averaged over all the probed images
            ev.IprobedMean = ev.IprobedMean + \
                np.mean(Im[mp.Fend.corr.maskBool]) / (2*Npairs)

            # Store probed image and its DM settings
            ev.Icube[:, :, whichImg] = Im
            if np.any(mp.dm_ind == 1):
                ev.Vcube1[:, :, whichImg] = mp.dm1.V
            if np.any(mp.dm_ind == 2):
                ev.Vcube2[:, :, whichImg] = mp.dm2.V

            # Report results
            probeSign = '-+'
            print('Actual Probe %d%s Contrast is: %.2e' % (
                    np.floor(iProbe/2), probeSign[(iProbe+1) % 2],
            np.mean(Im[mp.Fend.corr.maskBool])))

            # Assign image to positive or negative probe collection:
            if (iProbe+1) % 2 == 1:  # Odd; for plus probes
                if np.any(mp.dm_ind == 1):
                    DM1Vplus[:, :, iOdd] = dDM1Vprobe + DM1Vnom
                if np.any(mp.dm_ind == 2):
                    DM2Vplus[:, :, iOdd] = dDM2Vprobe + DM2Vnom
                Iplus[:, iOdd] = Im[mp.Fend.corr.maskBool]
                iOdd += 1
            elif (iProbe+1) % 2 == 0:  # Even; for minus probes
                if np.any(mp.dm_ind == 1):
                    DM1Vminus[:, :, iEven] = dDM1Vprobe + DM1Vnom
                if np.any(mp.dm_ind == 2):
                    DM2Vminus[:, :, iEven] = dDM2Vprobe + DM2Vnom
                Iminus[:, iEven] = Im[mp.Fend.corr.maskBool]
                iEven += 1

        # Calculate probe amplitudes and measurement vector.
        # (Refer again to Give'on+ SPIE 2011 to undersand why.)
        ampSq = (Iplus+Iminus)/2 - np.tile(I0vec.reshape((-1, 1)), (1, Npairs))  # square of probe E-field amplitudes
        ampSq[ampSq < 0] = 0  # If probe amplitude is zero, set amp = 0
        amp = np.sqrt(ampSq)  # E-field amplitudes, dimensions: [mp.Fend.corr.Npix, Npairs]
        isnonzero = np.all(amp, 1)
        zAll = ((Iplus-Iminus)/4).T  # Measurement vector, dimensions: [Npairs,mp.Fend.corr.Npix]
        ampSq2Dcube = np.zeros((mp.Fend.Neta, mp.Fend.Nxi, mp.est.probe.Npairs))
        for iProbe in range(Npairs):  # Display the actual probe intensity
            ampSq2D = np.zeros((mp.Fend.Neta, mp.Fend.Nxi))
            ampSq2D[mp.Fend.corr.maskBool] = ampSq[:, iProbe]
            ampSq2Dcube[:, :, iProbe] = ampSq2D
            print('*** Mean measured Inorm for probe #%d  =\t%.3e' %
                  (iProbe, np.mean(ampSq2D[mp.Fend.corr.maskBool])))

        # Plot relevant data for all the probes
        # falco_plot_pairwi1se_probes(mp, ev,
        #         DM1Vplus-repmat(DM1Vnom, [1,1,size(DM1Vplus,3)]),
        #         ampSq2Dcube)

        # ################# Perform the estimation ############################

        # Use Jacobian for estimation. This is fully model-based if the
        # Jacobian is purely model-based, or it is better if the Jacobian is
        # adaptive based on empirical data.
        if mp.est.flagUseJac:

            dEplus = np.zeros_like(Iplus, dtype=complex)
            for iProbe in range(Npairs):
                if mp.est.probe.whichDM == 1:
                    dV = DM1Vplus[:, :, iProbe] - DM1Vnom
                    dEplus[:, iProbe] = np.squeeze(jacStruct.G1[:, :, si]) * \
                        dV[mp.dm1.act_ele]
                elif mp.est.probe.whichDM == 2:
                    dV = DM2Vplus[:, :, iProbe] - DM2Vnom
                    dEplus[:, iProbe] = np.squeeze(jacStruct.G2[:, :, si]) * \
                        dV[mp.dm2.act_ele]

        # Get the probe phase from the model and measure the probe amplitude
        else:

            # For unprobed field based on model:
            if np.any(mp.dm_ind == 1):
                mp.dm1.V = DM1Vnom
            if np.any(mp.dm_ind == 2):
                mp.dm2.V = DM2Vnom
            if mp.flagFiber:
                pass
                # [~, E0] = model_compact(mp, modvar)
            else:
                E0 = falco.model.compact(mp, modvar)

            E0vec = E0[mp.Fend.corr.maskBool]

            # For probed fields based on model:
            Eplus = np.zeros_like(Iplus, dtype=complex)
            Eminus = np.zeros_like(Iminus, dtype=complex)
            for iProbe in range(Npairs):
                # For plus probes:
                if np.any(mp.dm_ind == 1):
                    mp.dm1.V = np.squeeze(DM1Vplus[:, :, iProbe])
                if np.any(mp.dm_ind == 2):
                    mp.dm2.V = np.squeeze(DM2Vplus[:, :, iProbe])
                if(mp.flagFiber):
                    pass
                    # [~, Etemp] = model_compact(mp, modvar);
                else:
                    Etemp = falco.model.compact(mp, modvar)
                Eplus[:, iProbe] = Etemp[mp.Fend.corr.maskBool]

                # For minus probes:
                if np.any(mp.dm_ind == 1):
                    mp.dm1.V = np.squeeze(DM1Vminus[:, :, iProbe])
                if np.any(mp.dm_ind == 2):
                    mp.dm2.V = np.squeeze(DM2Vminus[:, :, iProbe])
                if mp.flagFiber:
                    pass
                    # [~, Etemp] = model_compact(mp, modvar)
                else:
                    Etemp = falco.model.compact(mp, modvar)
                Eminus[:, iProbe] = Etemp[mp.Fend.corr.maskBool]

            # Create delta E-fields for each probe image.
            # Then create Npairs phase angles.
            dEplus = Eplus - np.tile(E0vec.reshape((-1, 1)), (1, Npairs))
            dEminus = Eminus - np.tile(E0vec.reshape((-1, 1)), (1, Npairs))
            dphdm = np.zeros((mp.Fend.corr.Npix, Npairs))  # phases
            for iProbe in range(Npairs):
                dphdm[:, iProbe] = np.arctan2(
                    np.imag(dEplus[:, iProbe]) - np.imag(dEminus[:, iProbe]),
                    np.real(dEplus[:, iProbe]) - np.real(dEminus[:, iProbe]))

        # Batch process the measurements to estimate the electric field in the
        # dark hole. Done pixel by pixel.

        if (mp.estimator.lower() == 'pwp-bp') or \
            (mp.estimator.lower() == 'pwp-kf' and ev.Itr < mp.est.ItrStartKF):

            Eest = np.zeros((mp.Fend.corr.Npix,), dtype=complex)
            zerosCounter = 0  # number of zeroed-out dark hole pixels
            for ipix in range(mp.Fend.corr.Npix):

                if mp.est.flagUseJac:
                    dE = dEplus[ipix, :].T
                    H = np.array([np.real(dE), np.imag(dE)])
                else:
                    H = np.zeros([Npairs, 2])  # Observation matrix
                    # Leave Eest for a pixel as zero if any probe amp is 0
                    if isnonzero[ipix] == 1:
                        for iProbe in range(Npairs):
                            H[iProbe, :] = amp[ipix, iProbe] * \
                                np.array([np.cos(dphdm[ipix, iProbe]),
                                          np.sin(dphdm[ipix, iProbe])])
                    else:
                        zerosCounter += 1

                Epix = np.linalg.pinv(H) @ zAll[:, ipix]  # Batch processing
                Eest[ipix] = Epix[0] + 1j*Epix[1]

            # If estimate is too bright, the estimate was probably bad.
            # !!!!!!!!!!!!!!BE VERY CAREFUL WITH THIS HARD-CODED VALUE!!!!!!!!!
            Eest[np.abs(Eest)**2 > 1e-2] = 0.0

            print('%d of %d pixels were given zero probe amplitude.' %
                    (zerosCounter, mp.Fend.corr.Npix))

            # Initialize the state and state covariance estimates for Kalman
            # filter. The state is the real and imag parts of the E-field.
            if mp.estimator.lower() == 'pwp-kf':
                # Re-organize the batch-processed E-field estimate into the 1st
                # state estimate for the Kalman filter
                xOld = np.zeros((2*mp.Fend.corr.Npix, 1))
                for ii in range(mp.Fend.corr.Npix):
                    xOld[2*(ii-1)+0:2*(ii-1)+1] = np.array([np.real(Eest[ii]),
                                                            np.imag(Eest[ii])])
                ev.xOld = xOld  # Save out for returning later

                # Initialize the state covariance matrix
                # (2x2 for each dark hole pixel)
                ev.Pold_KF_array = np.tile(mp.est.Pcoef0*np.eye(2),
                                           (mp.Fend.corr.Npix, 1, mp.Nsbp))

        # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        # Begin Kalman Filter Update
        # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        #
        # To be completed later...
        #
        # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        # End Kalman Filter Update
        # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

        # Save out the estimates
        ev.Eest[:, si] = Eest
        Iest = np.abs(Eest)**2
        ev.IincoEst[:, si] = I0vec - Iest  # incoherent light

    # Other data to save out
    ev.ampSqMean = np.mean(ampSq)  # Mean probe intensity
    ev.ampNorm = amp/np.sqrt(InormProbe)  # Normalized probe amplitude maps

    # Calculate the mean normalized intensity over the whole dark hole at all
    # wavelengths.
    ev.InormEst = np.mean(Iest)

    # Reset DM commands to their values before probing
    if np.any(mp.dm_ind == 1):
        mp.dm1.V = DM1Vnom

    if np.any(mp.dm_ind == 2):
        mp.dm2.V = DM2Vnom

    print('Completed pairwise probing estimation.')

    pass
    # return None

def gen_pairwise_probe(mp, InormDes, psi, badAxis):
    """
    Generate delta DM commands that probe the dark hole.

    Parameters
    ----------
    mp : falco.config.ModelParameter
        Object containing all model parameters.
    InormDes : float
        Desired normalized intensity of the probes in the image.
    psi : float
        phase angle of the sinusoidal part of the probe. Units of radians.
    badAxis : TYPE
        DESCRIPTION.

    Returns
    -------
    probeCmd : array_like
         Nact x Nact array of delta DM actuator commands to make a probe.
    """
    check.real_positive_scalar(InormDes, 'InormDes', ValueError)
    check.real_scalar(psi, 'psi', TypeError)
    if badAxis.lower() not in ('x', 'y', 'm'):
        raise ValueError('Invalid value for badAxis.')

    # Number of actuators across DM surface
    # (independent of beam diameter for time being)
    if mp.est.probe.whichDM == 1:
        Nact = mp.dm1.Nact
        dm = mp.dm1
    elif mp.est.probe.whichDM == 2:
        Nact = mp.dm2.Nact
        dm = mp.dm2

    # Coordinates in actuator space
    xs = np.arange(-(Nact-1)/2, (Nact+1)/2)/Nact - \
        np.round(mp.est.probe.offsetX)/Nact
    ys = np.arange(-(Nact-1)/2, (Nact+1)/2)/Nact - \
        np.round(mp.est.probe.offsetY)/Nact
    [XS, YS] = np.meshgrid(xs, ys)

    # Restrict the probing region if it is not possible to achieve
    if mp.est.probe.radius > Nact/2.0:
        mp.est.probe.radius = Nact/2.0

    # Generate the DM command for the probe
    magn = 4*np.pi*mp.lambda0*np.sqrt(InormDes)  # surface height to get desired intensity [meters]
    if badAxis.lower() == 'y':
        mX = mp.est.probe.radius
        mY = 2*mp.est.probe.radius
        omegaX = mp.est.probe.radius/2
        probeSurf = magn*np.sinc(mX*XS)*np.sinc(mY*YS)*np.cos(2*np.pi*omegaX*XS + psi)

    elif badAxis.lower() == 'x':
        mX = 2*mp.est.probe.radius
        mY = mp.est.probe.radius
        omegaY = mp.est.probe.radius/2
        probeSurf = magn*np.sinc(mX*XS)*np.sinc(mY*YS)*np.cos(2*np.pi*omegaY*YS + psi)

    elif badAxis.lower() == 'm':
        omegaX = mp.est.probe.Xloc/2
        omegaY = mp.est.probe.Yloc/2
        probeSurf = np.zeros_like(XS)
        for i in range(mp.Fend.Nfiber):
            probeSurf = probeSurf + \
                magn*np.sin(2*np.pi*omegaX(i)*XS + 2*np.pi*omegaY(i)*YS + psi)

    # Option to use just the sincs for a zero phase shift. This avoids the
    # phase discontinuity along one axis (for this probe only!).
    if psi == 0:
        m = 2*mp.est.probe.radius
        probeSurf = magn*np.sinc(m*XS)*np.sinc(m*YS)

    probeCmd = falco.dm.fit_surf_to_act(dm, probeSurf)

    # Scale the probe amplitude empirically if needed
    probeCmd = mp.est.probe.gainFudge*probeCmd

    return probeCmd