#!/usr/bin/env python
"""
This module contains functions and class definitions for running forward
models of models based on logistic regression.
"""

# stdlib imports
import numpy as np
import os.path
import re
import collections
import copy
# from scipy import sparse
import shutil
import tempfile
from timeit import default_timer as timer

# third party imports
from mapio.shake import ShakeGrid
from mapio.shake import getHeaderData
from mapio.gmt import GMTGrid
from mapio.gdal import GDALGrid
from mapio.grid2d import Grid2D
from mapio.geodict import GeoDict

from gfail.temphdf import TempHdf
from gfail.spatial import quickcut, trim_ocean

# temporary until mapio is updated
import warnings
warnings.filterwarnings('ignore')

PARAM_PATTERN = 'b[0-9]+'
LAYER_PATTERN = '_layer'
TERM_PATTERN = 'term'

SM_TERMS = ['MW', 'YEAR', 'MONTH', 'DAY', 'HOUR', 'pga', 'pgv', 'mmi']
SM_GRID_TERMS = ['pga', 'pgv', 'mmi']
# these will get np. prepended
OPERATORS = ['log', 'log10', 'arctan', 'power', 'sqrt', 'minimum', 'pi']
FLOATPAT = '[+-]?(?=\d*[.eE])(?=\.?\d)\d*\.?\d*(?:[eE][+-]?\d+)?'
INTPAT = '[0-9]+'
OPERATORPAT = '[\+\-\*\/]*'
MONTHS = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct',
          'Nov', 'Dec']

class LogisticModel(object):
    def __init__(self, shakefile, config, uncertfile=None, saveinputs=False,
                 slopefile=None, bounds=None, numstd=1, slopemod=None,
                 trimfile=None):
        """
        Sets up the logistic model

        Args:
            shakefile (str): Path to shakemap grid.xml file for the event.
            config: configobj object defining the model and its inputs. Only
                one model should be described in each config file.
            uncertfile (str): Path to uncertainty.xml file.
            saveinputs (bool): Save input layers as Grid2D objects in addition
                to the model? If false (the default), it will just output the
                model.
            slopefile (str): Optional path to slopefile that will be resampled
                to the other input files for applying thresholds. OVERWRITES
                VALUE IN CONFIG.
            bounds (dict): Default of None uses ShakeMap boundaries, otherwise
                a dictionary of boundaries to cut to like

                .. code-block:: python

                    bounds = {
                        'xmin': lonmin, 'xmax': lonmax,
                        'ymin': latmin, 'ymax': latmax
                    }

            numstd (float): Number of +/- standard deviations to use if
                uncertainty is computed.
            slopemod (str): How slope input should be modified to be in
                degrees: e.g., ``np.arctan(slope) * 180. / np.pi`` or
                ``slope/100.`` (note that this may be in the config file
                already).
            trimfile (str): shapefile of earth's landmasses to use to cut
                offshore areas.
        """
        mnames = getLogisticModelNames(config)
        if len(mnames) == 0:
            raise Exception('No config file found or problem with config '
                            'file format')
        if len(mnames) > 1:
            raise Exception('Config file contains more than one model which '
                            'is no longer allowed, update your config file '
                            'to the newer format')

        self.model = mnames[0]
        self.config = config
        cmodel = config[self.model]
        self.modeltype = cmodel['gfetype']
        self.coeffs = validateCoefficients(cmodel)
        # key = layer name, value = file name
        self.layers = validateLayers(cmodel)
        self.terms, timeField = validateTerms(cmodel, self.coeffs, self.layers)
        self.interpolations = validateInterpolations(cmodel, self.layers)
        self.units = validateUnits(cmodel, self.layers)
        self.gmused = [value for term, value in cmodel['terms'].items()
                       if 'pga' in value.lower() or 'pgv' in
                       value.lower() or 'mmi' in value.lower()]
        self.modelrefs, self.longrefs, self.shortrefs = validateRefs(cmodel)
        self.numstd = numstd
        self.clips = validateClips(cmodel, self.layers, self.gmused)

        if cmodel['baselayer'] not in list(self.layers.keys()):
            raise Exception('You must specify a base layer corresponding to '
                            'one of the files in the layer section.')
        self.saveinputs = saveinputs
        if slopefile is None:
            try:
                self.slopefile = cmodel['slopefile']
            except:
                # print('Slopefile not specified in config, no slope '
                #      'thresholds will be applied\n')
                self.slopefile = None
        else:
            self.slopefile = slopefile
        if slopemod is None:
            try:
                self.slopemod = cmodel['slopemod']
            except:
                self.slopemod = None

        # See if trimfile exists
        if trimfile is not None:
            if not os.path.exists(trimfile):
                print(
                    'trimfile defined does not exist: %s\nOcean will not be trimmed' % trimfile)
                self.trimfile = None
            elif os.path.splitext(trimfile)[1] != '.shp':
                print('trimfile must be a shapefile, ocean will not be trimmed')
                self.trimfile = None
            else:
                self.trimfile = trimfile
        else:
            self.trimfile = None

        # Get month of event
        griddict, eventdict, specdict, fields, uncertainties = \
            getHeaderData(shakefile)
        MONTH = MONTHS[(eventdict['event_timestamp'].month) - 1]

        # Figure out how/if need to cut anything
        geodict = ShakeGrid.getFileGeoDict(shakefile, adjust='res')
        if bounds is not None:  # Make sure bounds are within ShakeMap Grid
            if (geodict.xmin > bounds['xmin'] or
                    geodict.xmax < bounds['xmax'] or
                    geodict.ymin > bounds['ymin'] or
                    geodict.ymax < bounds['ymax']):
                print('Specified bounds are outside shakemap area, using '
                      'ShakeMap bounds instead.')
                bounds = None
        if bounds is not None:
            tempgdict = GeoDict.createDictFromBox(
                bounds['xmin'], bounds['xmax'],
                bounds['ymin'], bounds['ymax'],
                geodict.dx, geodict.dy, inside=False)
            gdict = geodict.getBoundsWithin(tempgdict)
        else:
            gdict = geodict

        # Now find the layer that is our base layer and get the largest bounds
        # we can guarantee not to exceed shakemap bounds
        basefile = self.layers[cmodel['baselayer']]
        ftype = getFileType(basefile)
        if ftype == 'esri':
            basegeodict, firstcol = GDALGrid.getFileGeoDict(basefile)
            if basegeodict == gdict:
                sampledict = gdict
            else:
                sampledict = basegeodict.getBoundsWithin(gdict)
        elif ftype == 'gmt':
            basegeodict, firstcol = GMTGrid.getFileGeoDict(basefile)
            if basegeodict == gdict:
                sampledict = gdict
            else:
                sampledict = basegeodict.getBoundsWithin(gdict)
        else:
            raise Exception('All predictor variable grids must be a valid '
                            'GMT or ESRI file type.')

        # Do we need to subdivide baselayer?
        if 'divfactor' in self.config[self.model].keys():
            divfactor = float(self.config[self.model]['divfactor'])
            if divfactor != 1.:
                # adjust sampledict so everything will be resampled
                newxmin = sampledict.xmin - sampledict.dx / \
                    2. + sampledict.dx/(2.*divfactor)
                newymin = sampledict.ymin - sampledict.dy / \
                    2. + sampledict.dy/(2.*divfactor)
                newxmax = sampledict.xmax + sampledict.dx / \
                    2. - sampledict.dx/(2.*divfactor)
                newymax = sampledict.ymax + sampledict.dy / \
                    2. - sampledict.dy/(2.*divfactor)
                newdx = sampledict.dx/divfactor
                newdy = sampledict.dy/divfactor

                sampledict = GeoDict.createDictFromBox(
                    newxmin, newxmax, newymin,
                    newymax, newdx, newdy, inside=True)

        # Find slope thresholds, if applicable
        self.slopemin = 'none'
        self.slopemax = 'none'
        if slopefile is not None:
            try:
                self.slopemin = float(config[self.model]['slopemin'])
                self.slopemax = float(config[self.model]['slopemax'])
            except:
                print('Could not find slopemin and/or slopemax in config, '
                      'limits. No slope thresholds will be applied.')
                self.slopemin = 'none'
                self.slopemax = 'none'

        # Make temporary directory for hdf5 pytables file storage
        self.tempdir = tempfile.mkdtemp()

        # now load the shakemap, resampling and padding if necessary
        temp = ShakeGrid.load(shakefile)  # , adjust='res')
        self.shakedict = temp.getShakeDict()
        self.eventdict = temp.getEventDict()
        self.shakemap = {}

        # Read both PGA and PGV in, may need them for thresholds
        for gm in ['pga', 'pgv']:
            junkfile = os.path.join(self.tempdir, 'temp.bil')
            GDALGrid.copyFromGrid(temp.getLayer(gm)).save(junkfile)
            if gm in self.interpolations.keys():
                intermeth = self.interpolations[gm]
            else:
                intermeth = 'bilinear'
            junkgrid = quickcut(junkfile, sampledict, precise=True,
                                method=intermeth)
            if gm in self.clips:
                junkgrid.setData(np.clip(junkgrid.getData(),
                                         self.clips[gm][0], self.clips[gm][1]))
            self.shakemap[gm] = TempHdf(
                junkgrid, os.path.join(self.tempdir, '%s.hdf5' % gm))
            os.remove(junkfile)
        del(temp)

        # get updated geodict
        sampledict = junkgrid.getGeoDict()

        # take uncertainties into account, if available
        if uncertfile is not None:
            self.uncert = {}
            try:
                # Only read in the ones that will be needed
                temp = ShakeGrid.load(uncertfile)
                for gm in self.gmused:
                    if 'pgv' in gm:
                        gmsimp = 'pgv'
                    elif 'pga' in gm:
                        gmsimp = 'pga'
                    elif 'mmi' in gm:
                        gmsimp = 'mmi'
                    junkfile = os.path.join(self.tempdir, 'temp.bil')
                    GDALGrid.copyFromGrid(temp.getLayer(
                        'std%s' % gmsimp)).save(junkfile)
                    if gmsimp in self.interpolations.keys():
                        intermeth = self.interpolations[gmsimp]
                    else:
                        intermeth = 'bilinear'
                    junkgrid = quickcut(junkfile, sampledict, precise=True,
                                        method=intermeth)
                    if gmsimp in self.clips:
                        junkgrid.setData(
                            np.clip(junkgrid.getData(), self.clips[gmsimp][0],
                                    self.clips[gmsimp][1]))
                    self.uncert['std' + gmsimp] = TempHdf(
                        junkgrid, os.path.join(self.tempdir,
                                               'std%s.hdf5' % gmsimp))
                    os.remove(junkfile)
                del(temp)
            except:
                print('Could not read uncertainty file, ignoring '
                      'uncertainties')
                self.uncert = None
        else:
            self.uncert = None

        # Load the predictor layers, save as hdf5 temporary files, put file
        # locations into a dictionary.

        # Will be replaced in the next section if a slopefile was defined
        self.nonzero = None

        # key = layer name, value = grid object
        self.layerdict = {}

        didslope = False
        for layername, layerfile in self.layers.items():
            start = timer()
            if isinstance(layerfile, list):
                for lfile in layerfile:
                    if timeField == 'MONTH':
                        if lfile.find(MONTH) > -1:
                            layerfile = lfile
                            ftype = getFileType(layerfile)
                            interp = self.interpolations[layername]
                            temp = quickcut(layerfile, sampledict,
                                            precise=True, method=interp)
                            if layername in self.clips:
                                temp.setData(
                                    np.clip(temp.getData(),
                                            self.clips[layername][0],
                                            self.clips[layername][1]))
                            self.layerdict[layername] = TempHdf(
                                temp, os.path.join(self.tempdir,
                                                   '%s.hdf5' % layername))
                            del(temp)
            else:
                interp = self.interpolations[layername]
                temp = quickcut(layerfile, sampledict,
                                precise=True, method=interp)
                if layername in self.clips:
                    temp.setData(
                        np.clip(temp.getData(),
                                self.clips[layername][0],
                                self.clips[layername][1]))
                self.layerdict[layername] = TempHdf(
                    temp, os.path.join(self.tempdir, '%s.hdf5' % layername))
                td = temp.getGeoDict()
                if td != sampledict:
                    raise Exception(
                        'Geodictionaries of resampled files do not match')

                if layerfile == self.slopefile:
                    flag = 0
                    if self.slopemin == 'none' and self.slopemax == 'none':
                        flag = 1
                    if self.slopemod is None:
                        slope1 = temp.getData().astype(float)
                        slope = 0
                    else:
                        try:
                            slope = temp.getData().astype(float)
                            slope1 = eval(self.slopemod)
                        except:
                            print('slopemod provided not valid, continuing '
                                  'without slope thresholds.')
                            flag = 1
                    if flag == 0:
                        nonzero = np.array(
                            [(slope1 > self.slopemin) &
                             (slope1 <= self.slopemax)])
                        self.nonzero = nonzero[0, :, :]
                        del(slope1)
                        del(slope)
                    else:
                        # Still remove areas where the slope equals exactly
                        # 0.0 to remove offshore liq areas.
                        nonzero = np.array([slope1 != 0.0])
                        self.nonzero = nonzero[0, :, :]
                        del(slope1)
                    didslope = True
                del(temp)

            print('Loading %s layer: %1.1f sec'
                  % (layername, timer() - start))

        if didslope is False and self.slopefile is not None:
            # Slope didn't get read in yet
            temp = quickcut(self.slopefile, sampledict, precise=True,
                            method='bilinear')
            flag = 0
            if self.slopemin == 'none' and self.slopemax == 'none':
                flag = 1
            if self.slopemod is None:
                slope1 = temp.getData().astype(float)
                slope = 0
            else:
                try:
                    slope = temp.getData().astype(float)
                    slope1 = eval(self.slopemod)
                except:
                    print('slopemod provided not valid, continuing without '
                          'slope thresholds')
                    flag = 1
            if flag == 0:
                nonzero = np.array([(slope1 > self.slopemin) &
                                    (slope1 <= self.slopemax)])
                self.nonzero = nonzero[0, :, :]
                del(slope1)
                del(slope)
            else:
                # Still remove areas where the slope equals exactly
                # 0.0 to remove offshore liq areas.
                nonzero = np.array([slope1 != 0.0])
                self.nonzero = nonzero[0, :, :]
                del(slope1)

        self.nuggets = [str(self.coeffs['b0'])]

        ckeys = list(self.terms.keys())
        ckeys.sort()
        for key in ckeys:
            term = self.terms[key]
            coeff = self.coeffs[key]
            self.nuggets.append('(%g * %s)' % (coeff, term))

        self.equation = ' + '.join(self.nuggets)

        if self.uncert is not None:
            self.nugmin = copy.copy(self.nuggets)
            self.nugmax = copy.copy(self.nuggets)

            # Find the term with the shakemap input and replace for these
            # nuggets.
            for gm in ['pga', 'mmi', 'pgv']:
                for k, nug in enumerate(self.nuggets):
                    tempnug = ("self.shakemap['%s'].getSlice(rowstart, "
                               "rowend, colstart, colend, name='%s')"
                               % (gm, gm))
                    if tempnug in nug:
                        newnug = ("np.exp(np.log(%s) - self.numstd * "
                                  "self.uncert['std%s'].getSlice(rowstart, "
                                  "rowend, colstart, colend, name='std%s'))"
                                  % (tempnug, gm, gm))
                        self.nugmin[k] = self.nugmin[k].replace(
                            tempnug, newnug)
                        newnug = ("np.exp(np.log(%s) + self.numstd * "
                                  "self.uncert['std%s'].getSlice(rowstart, "
                                  "rowend, colstart, colend, name='std%s'))"
                                  % (tempnug, gm, gm))
                        self.nugmax[k] = self.nugmax[k].replace(
                            tempnug, newnug)

            self.equationmin = ' + '.join(self.nugmin)
            self.equationmax = ' + '.join(self.nugmax)
        else:
            self.equationmin = None
            self.equationmax = None

        self.geodict = sampledict

    def getEquations(self):
        """
        Method for LogisticModel class to extract strings defining the
        equations for the model for median ground motions and +/- one standard
        deviation (3 total).

        Returns:
            tuple: (equation, equationmin, equationmax) where:
                * equation: the equation for median ground motions,
                * equationmin: the equation for the same model but using
                  median ground motions minus 1 standard deviation
                * equationmax: same as above but for plus 1 standard deviation.
        """
        return self.equation, self.equationmin, self.equationmax

    def getGeoDict(self):
        """
        Returns the geodictionary of the LogisticModel class defining bounds
        and resolution of model inputs and outputs.

        Returns:
            geodict: mapio geodict object
        """
        return self.geodict

    def calculate(self, cleanup=True, rowmax=300, colmax=None):
        """
        Calculate the model.

        Args:
            cleanup (bool): If True, delete temporary hdf5 files
            rowmax (int): Number of rows to compute at once; If None, all rows
                will be computed at once.
            colmax (int): Number of columns to compute at once; If None, all
                columns will be computed at once.
        Returns:
            dict: Dictionary containing the model results (and model inputs if
            saveinputs was set to True). See
            `the description <https://github.com/usgs/groundfailure#api-for-model-output>`_
            of the structure.
        """
        tk = list(self.shakemap.keys())[0]
        # Figure out what slices to do
        rowstarts, rowends, colstarts, colends = \
            self.shakemap[tk].getSliceDiv(rowmax, colmax)

        # Make empty matrix to fill
        X = np.empty([self.geodict.ny, self.geodict.nx])

        # Loop through slices, appending output each time
        for rowstart, rowend, colstart, colend in \
                zip(rowstarts, rowends, colstarts, colends):
            X[rowstart:rowend, colstart:colend] = eval(self.equation)

        P = 1/(1 + np.exp(-X))

        if 'vs30max' in self.config[self.model].keys():
            vs30 = self.layerdict['vs30'].getSlice(
                None, None, None, None, name='vs30')
            P[vs30 > float(self.config[self.model]['vs30max'])] = 0.0

        if 'minpgv' in self.config[self.model].keys():
            pgv = self.shakemap['pgv'].getSlice(
                None, None, None, None, name='pgv')
            P[pgv < float(self.config[self.model]['minpgv'])] = 0.0

        if 'minpga' in self.config[self.model].keys():
            pga = self.shakemap['pga'].getSlice(
                None, None, None, None, name='pga')
            P[pga < float(self.config[self.model]['minpga'])] = 0.0

        if 'coverage' in self.config[self.model].keys():
            eqn = self.config[self.model]['coverage']['eqn']
            P = eval(eqn)

        if self.uncert is not None:
            # Make empty matrix to fill
            Xmin = np.empty([self.geodict.ny, self.geodict.nx])
            Xmax = Xmin.copy()
            # Loop through slices, appending output each time
            for rowstart, rowend, colstart, colend in \
                    zip(rowstarts, rowends, colstarts, colends):
                Xmin[rowstart:rowend, colstart:colend] = eval(self.equationmin)
                Xmax[rowstart:rowend, colstart:colend] = eval(self.equationmax)

            Pmin = 1/(1 + np.exp(-Xmin))
            Pmax = 1/(1 + np.exp(-Xmax))

            if 'vs30max' in self.config[self.model].keys():
                vs30 = self.layerdict['vs30'].getSlice(
                    None, None, None, None, name='vs30')
                Pmin[vs30 > float(self.config[self.model]['vs30max'])] = 0.0
                Pmax[vs30 > float(self.config[self.model]['vs30max'])] = 0.0

            if 'minpgv' in self.config[self.model].keys():
                pgv = self.shakemap['pgv'].getSlice(
                    None, None, None, None, name='pgv')
                Pmin[pgv < float(self.config[self.model]['minpgv'])] = 0.0
                Pmax[pgv < float(self.config[self.model]['minpgv'])] = 0.0

            if 'minpga' in self.config[self.model].keys():
                pga = self.shakemap['pgv'].getSlice(
                    None, None, None, None, name='pga')
                Pmin[pga < float(self.config[self.model]['minpga'])] = 0.0
                Pmax[pga < float(self.config[self.model]['minpga'])] = 0.0

            if 'coverage' in self.config[self.model].keys():
                eqnmin = eqn.replace('P', 'Pmin')
                eqnmax = eqn.replace('P', 'Pmax')
                Pmin = eval(eqnmin)
                Pmax = eval(eqnmax)

            #Pmin[np.isnan(Pmin)] = 0.0
            #Pmax[np.isnan(Pmax)] = 0.0

        #P[np.isnan(P)] = 0.0

        if self.slopefile is not None and self.nonzero is not None:
            # Apply slope min/max limits
            print('applying slope thresholds')
            P = P * self.nonzero
            #P[P==0.0] = float('nan')
            #P[np.isnan(P)] = 0.0
            if self.uncert is not None:
                Pmin = Pmin * self.nonzero
                Pmax = Pmax * self.nonzero
                #Pmin[Pmin==0.0] = float('nan')
                #Pmax[Pmax==0.0] = float('nan')
                #Pmin[np.isnan(Pmin)] = 0.0
                #Pmax[np.isnan(Pmax)] = 0.0

        # Stuff into Grid2D object
        if 'Jessee' in self.modelrefs['shortref']:
            if 'coverage' not in self.config[self.model].keys():
                units5 = 'relative hazard'
            else:
                units5 = 'areal coverage'
        elif 'Zhu' in self.modelrefs['shortref']:
            units5 = 'areal coverage'
        else:
            units5 = 'probability'

        shakedetail = (
            '%s_ver%s'
            % (self.shakedict['shakemap_id'],
               self.shakedict['shakemap_version']))
        description = {
            'name': self.modelrefs['shortref'],
            'longref': self.modelrefs['longref'],
            'units': units5,
            'shakemap': shakedetail,
            'event_id': self.eventdict['event_id'],
            'parameters': {'slopemin': self.slopemin,
                           'slopemax': self.slopemax,
                           'modeltype': self.modeltype}}
        if 'vs30max' in self.config[self.model].keys():
            description['vs30max'] = float(self.config[self.model]['vs30max'])
        if 'minpgv' in self.config[self.model].keys():
            description['minpgv'] = float(self.config[self.model]['minpgv'])

        Pgrid = Grid2D(P, self.geodict)
        if self.trimfile is not None:
            # Turn all offshore cells to nan
            Pgrid = trim_ocean(Pgrid, self.trimfile, nodata=float('nan'))
        rdict = collections.OrderedDict()
        rdict['model'] = {
            'grid': Pgrid,
            'label': ('%s %s') % (self.modeltype.capitalize(),
                                  units5.title()),
            'type': 'output',
            'description': description
        }
        if self.uncert is not None:
            Pmingrid = Grid2D(Pmin, self.geodict)
            Pmaxgrid = Grid2D(Pmax, self.geodict)
            if self.trimfile is not None:
                Pmingrid = trim_ocean(
                    Pmingrid, self.trimfile, nodata=float('nan'))
                Pmaxgrid = trim_ocean(
                    Pmaxgrid, self.trimfile, nodata=float('nan'))
            rdict['modelmin'] = {
                'grid': Pmingrid,
                'label': ('%s %s (-%0.1f std ground motion)'
                          % (self.modeltype.capitalize(),
                             units5.title(),
                             self.numstd)),
                'type': 'output',
                'description': description
            }
            rdict['modelmax'] = {
                'grid': Pmaxgrid,
                'label': ('%s %s (+%0.1f std ground motion)'
                          % (self.modeltype.capitalize(),
                             units5.title(),
                             self.numstd)),
                'type': 'output',
                'description': description
            }

        # This step might swamp memory for higher resolution runs
        if self.saveinputs is True:
            for layername, layergrid in list(self.layerdict.items()):
                units = self.units[layername]
                if units is None:
                    units = ''
                rdict[layername] = {
                    'grid': Grid2D(
                        layergrid.getSlice(
                            None, None, None, None, name=layername),
                        self.geodict
                    ),
                    'label': '%s (%s)' % (layername, units),
                    'type': 'input',
                    'description': {
                        'units': units,
                        'name': self.shortrefs[layername],
                        'longref': self.longrefs[layername]
                    }
                }
            for gmused in self.gmused:
                if 'pga' in gmused:
                    units = '%g'
                    getkey = 'pga'
                elif 'pgv' in gmused:
                    units = 'cm/s'
                    getkey = 'pgv'
                elif 'mmi' in gmused:
                    units = 'intensity'
                    getkey = 'mmi'
                else:
                    continue
                    # Layer is derived from several input layers, skip
                    # outputting this layer

                if getkey in rdict:
                    continue

                layer = self.shakemap[getkey].getSlice(
                    None, None, None, None, name=getkey)
                rdict[getkey] = {
                    'grid': Grid2D(layer, self.geodict),
                    'label': '%s (%s)' % (getkey.upper(), units),
                    'type': 'input',
                    'description': {
                        'units': units,
                        'shakemap': shakedetail
                    }
                }
                if self.uncert is not None:
                    uncertlayer = self.uncert[getkey].getSlice(
                        None, None, None, None, name='std'+getkey)
                    layer1 = np.exp(np.log(layer) - uncertlayer)
                    rdict[getkey + 'modelmin'] = {
                        'grid': Grid2D(layer1, self.geodict),
                        'label': ('%s - %0.1f std (%s)'
                                  % (getkey.upper(),
                                     self.numstd, units)),
                        'type': 'input',
                        'description': {'units': units,
                                        'shakemap': shakedetail}
                    }
                    layer2 = np.exp(np.log(layer) + uncertlayer)
                    rdict[getkey + 'modelmax'] = {
                        'grid': Grid2D(layer2, self.geodict),
                        'label': ('%s + %0.1f std (%s)'
                                  % (getkey.upper(),
                                     self.numstd, units)),
                        'type': 'input',
                        'description': {'units': units,
                                        'shakemap': shakedetail}
                    }
        if cleanup:
            shutil.rmtree(self.tempdir)
        return rdict

def getLogisticModelNames(config):
    """
    Get the names of the models present in the configobj

    Args:
        config: configobj object defining the model and its inputs.

    Returns:
        list: list of model names.
    """
    names = []
    lmodel_space = config
    for key, value in lmodel_space.items():
        if isinstance(value, str):
            continue
        else:  # this is a model
            names.append(key)
    return names

def getFileType(filename):
    """
    Determine whether input file is a shapefile or a grid (ESRI or GMT).

    Args:
        filename (str): Path to candidate filename.

    Returns:
        str: 'shapefile', 'grid', or 'unknown'.
    """
    #TODO MOVE TO MAPIO.
    if os.path.isdir(filename):
        return 'dir'
    ftype = GMTGrid.getFileType(filename)
    if ftype != 'unknown':
        return 'gmt'
    # Skip over ESRI header files
    if filename.endswith('.hdr'):
        return 'unknown'
    try:
        GDALGrid.getFileGeoDict(filename)
        return 'esri'
    except:
        pass
    return 'unknown'

def getAllGridFiles(indir):
    """
    Get list of all gmt or esri (.grd, .bil) files in a directory.

    Args:
        indir (str): Directory to search.
    Returns:
        list: List of file names.
    """
    #TODO MOVE TO MAPIO
    tflist = os.listdir(indir)
    flist = []
    for tf in tflist:
        fullfile = os.path.join(indir, tf)
        ftype = getFileType(fullfile)
        if ftype in ['gmt', 'esri']:
            flist.append(fullfile)
    return flist

def validateCoefficients(cmodel):
    """
    Ensures coefficients provided in model description are valid and outputs
    a dictionary of the coefficients.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model,
            for example:

            .. code-block:: python

                cmodel = config['test_model']

    Returns:
        dict: a dictionary of model coefficients named b0, b1, b2...
    """
    coeffs = {}
    for key, value in cmodel['coefficients'].items():
        if re.search('b[0-9]*', key) is None:
            raise Exception('coefficients must be named b0, b1, ...')
        coeffs[key] = float(value)
    if 'b0' not in list(coeffs.keys()):
        raise Exception('coefficients must include an intercept '
                        'coefficient named b0.')
    return coeffs

def validateClips(cmodel, layers, gmused):
    """
    Ensures coefficients provided in model description are valid and outputs
    a dictionary of the coefficients.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model,
            for example:

            .. code-block:: python

                cmodel = config['test_model']

    Returns:
        dict: a dictionary of clip values for each layer (if exists)
    """
    clips = {}
    if 'clip' in cmodel:
        for key, value in cmodel['clip'].items():
            if key not in layers:
                if key not in gmused:
                    x1 = [par for par in gmused if key in par]
                    if len(x1) == 0:
                        raise Exception(
                            'Clipping key %s does not match any names of layers'
                            % key)
            clips[key] = (float(value[0]), float(value[1]))
    return clips

def validateLayers(cmodel):
    """
    Ensures all input files required to run the model exist and are valid
    file types.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model,
            for example,

            .. code-block:: python

                cmodel = config['test_model']

    Returns:
        dict: a dictionary of file names, e.g.

        .. code-block:: python

            {
                'slope': 'slopefile.bil',
                'vs30': 'vs30.grd'
            }

    """
    layers = {}
    longrefs = {}
    shortrefs = {}
    for key in cmodel['layers'].keys():
        for item, value in cmodel['layers'][key].items():
            if item == 'file':
                ftype = getFileType(value)
                if ftype == 'unknown':
                    raise Exception('layer file %s is not a valid GMT or '
                                    'ESRI file.' % value)
                if ftype == 'dir':
                    value = getAllGridFiles(value)
                layers[key] = value
            elif item == 'shortref':
                shortrefs[key] = value
            elif item == 'longref':
                longrefs[key] = value
    return layers

def validateTerms(cmodel, coeffs, layers):
    """
    Reformats model inputs from config file, replacing functions with numpy
    functions, inserting code for extracting data from each layer (required
    to run eval in the calculate step), addressing any time variables, and
    checks that term names match coefficient names.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model,
            e.g.

            .. code-block:: python

                cmodel = config['test_model']

        coeffs (dict): Dictionary of model coefficients, e.g.

            .. code-block:: python

                {'b0': 3.5, 'b1': -0.01}

        layers (dict): Dictionary of file names for all input layers, e.g.

            .. code-block:: python

                {'slope': 'slopefile.bil', 'vs30': 'vs30.grd'}

    Returns:
        tuple: (terms, timeField), where
            - 'terms' is a dictionary of terms that form the model equation,
              e.g.

            .. code-block:: python

                {
                    'b1': "self.layerdict['friction'].getData()",
                    'b2': "self.layerdict['slope'].getData()/100."
                }

            - 'timeField' indicates the time that is used to know which input
              file to read in, e.g. for monthly average precipitation, 'MONTH'.
    """
    #TODO:
    #    - Return a time field for every term, not just one global one.

    terms = {}
    timeField = None
    for key, value in cmodel['terms'].items():
        if key not in list(coeffs.keys()):
            raise Exception('Term names must match names of coefficients')
        # replace log with np.log, make sure variables are all in layers list,
        # etc.
        term, rem, tTimeField = checkTerm(value, layers)
        if tTimeField is not None:
            timeField = tTimeField
        if len(rem):
            msg = ('Term "%s" contains the unknown text fragment "%s". '
                   'This may cause the expression to fail.')
            tpl = (term, rem)
            raise Exception(msg % tpl)
        terms[key] = term
    return (terms, timeField)

def validateInterpolations(cmodel, layers):
    """Validate logistic model interpolation.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model.
        layers (dict): Dictionary of file names for all input layers.

    Returns:
        dict: Model interpolation methods.
    """
    interpolations = {}
    for key, value in cmodel['interpolations'].items():
        if key not in list(layers.keys()):
            raise Exception(
                'Interpolation key %s does not match any names of layers'
                % key)
        methods = ['linear', 'nearest', 'cubic', 'bilinear']
        if value not in methods:
            raise Exception(
                'Interpolation method %s not in approved list of methods: %s'
                % (key, str(methods)))
        interpolations[key] = value
    for key in list(layers.keys()):
        if key not in list(interpolations.keys()):
            raise Exception(
                'No interpolation method configured for layer %s' % key)
    return interpolations

def validateUnits(cmodel, layers):
    """Validate model units.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model.
        layers (dict): Dictionary of file names for all input layers.

    Returns:
        dict: Model units.
    """
    units = {}
    for key in cmodel['layers'].keys():
        if 'units' in cmodel['layers'][key]:
            units[key] = cmodel['layers'][key]['units']
        else:
            raise Exception('No unit string configured for layer %s' % key)
    return units

def validateLogisticModels(config):
    """Validate model names.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model.
        layers (dict): Dictionary of file names for all input layers.

    Returns:
        bool: True if the model names are valid
    """
    mnames = getLogisticModelNames(config)
    if len(mnames) > 1:
        raise Exception('Config file contains more than one model which is '
                        'no longer allowed, update your config file to the '
                        'newer format')
    for cmodelname in mnames:
        try:
            cmodel = config[cmodelname]
            coeffs = validateCoefficients(cmodel)
            # key = layer name, value = file name
            layers = validateLayers(cmodel)
            terms, timeField = validateTerms(cmodel, coeffs, layers)
            if timeField is not None:
                for (layer, layerfile) in list(layers.items()):
                    if isinstance(layerfile, list):
                        for lfile in layerfile:
                            if timeField == 'MONTH':
                                pass
            validateInterpolations(cmodel, layers)
            if cmodel['baselayer'] not in layers:
                raise Exception(
                    'Model %s missing baselayer parameter.' % cmodelname)
        except Exception as e:
            raise Exception('Validation failed with error: "%s" on model %s'
                            % (str(e), cmodelname))

    return True

def validateRefs(cmodel):
    """Validate references for models and layers.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model.

    Returns:
        tuple: (modelrefs, longrefs, shortrefs) where:
            * modelrefs: dictionary of citation information for model
                keys='longref', 'shortref'
            * shortrefs: dictionary containing short reference for each
                input layer
            * longrefs: dictionary containing full references for each
                input layer

    """
    longrefs = {}
    shortrefs = {}
    modelrefs = {}
    for key in cmodel['layers'].keys():
        if 'longref' in cmodel['layers'][key]:
            longrefs[key] = cmodel['layers'][key]['longref']
        else:
            print('No longref provided for layer %s' % key)
            longrefs[key] = 'unknown'
        if 'shortref' in cmodel['layers'][key]:
            shortrefs[key] = cmodel['layers'][key]['shortref']
        else:
            print('No shortref provided for layer %s' % key)
            shortrefs[key] = 'unknown'
    try:
        modelrefs['longref'] = cmodel['longref']
    except:
        print('No model longref provided')
        modelrefs['longref'] = 'unknown'
    try:
        modelrefs['shortref'] = cmodel['shortref']
    except:
        print('No model shortref provided')
        modelrefs['shortref'] = 'unknown'
    return modelrefs, longrefs, shortrefs

def checkTerm(term, layers):
    """Checks terms of equation and replaces text with machine readable operators

    Args:
        term: term from model configuration file
        layers: dictionary of file names for all input layers

    Returns:
        tuple: (term, tterm, timeField) where:
            * term: dictionary of verified terms for equation with keys corresponding
                to each layer name
            * tterm: any unconverted and unverified text that may cause expression to fail
            * timeField: if any inputs are time dependent, output is unit of time (e.g., 'YEAR'),
                otherwise, None.
    """
    # startterm = term
    # Strip out everything that isn't: 0-9.() operators, +-/* or layer names.
    # Anything left is an unknown symbol.
    tterm = term
    # remove log, sqrt, etc.
    for op in OPERATORS:
        tterm = tterm.replace(op, '')
    # remove ShakeMap variables
    for sm_term in SM_TERMS:
        tterm = tterm.replace(sm_term, '')
    # remove layer names
    for layer in layers:
        tterm = tterm.replace(layer, '')
    # remove arithmetic operators
    tterm = re.sub(OPERATORPAT, '', tterm)
    # remove floating point numbers
    tterm = re.sub(FLOATPAT, '', tterm)
    # remove integer numbers
    tterm = re.sub(INTPAT, '', tterm)
    # remove parentheses
    tterm = re.sub('[()]*', '', tterm)
    # remove any blank spaces
    tterm = tterm.strip()
    # remove commas
    tterm = tterm.strip(',')
    # anything left *might* cause an error
    for op in OPERATORS:
        if term.find(op) > -1:
            term = term.replace(op, 'np.'+op)

    for sm_term in SM_GRID_TERMS:
        term = term.replace(
            sm_term,
            "self.shakemap['%s'].getSlice(rowstart, rowend, "
            "colstart, colend, name='%s')" % (sm_term, sm_term))

    # replace the macro MW with the magnitude value from the shakemap
    term = term.replace('MW', "self.eventdict['magnitude']")

    # term.replace('YEAR',"self.shakemap.getEventDict()['event_time'].year")
    # hasTime = False
    timeField = None
    for unit in ['YEAR', 'MONTH', 'DAY', 'HOUR']:
        if term.find(unit) > -1:
            term = term.replace(unit, '')
            timeField = unit

    for layer in layers:
        if layer == 'friction':
            term = term.replace(
                layer,
                "np.nan_to_num(self.layerdict['%s'].getSlice(rowstart, "
                "rowend, colstart, colend, name='%s'))" % (layer, layer))
        else:
            term = term.replace(
                layer,
                "self.layerdict['%s'].getSlice(rowstart, rowend, colstart, "
                "colend, name='%s')" % (layer, layer))
    return (term, tterm, timeField)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

'''

    ImagePlayer: denoising HAADF images.

    Copyright (C) 2021  Feng Wang

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

    This program 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 Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.

'''

import os
os.environ["QT_LOGGING_RULES"]= '*.debug=false;qt.qpa.*=false'
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]=""
#os.environ["CUDA_VISIBLE_DEVICES"]="0"

from PySide6.QtPrintSupport import QPrintDialog, QPrinter
from PySide6.QtWidgets import (QApplication, QDialog, QFileDialog, QLabel, QMainWindow, QMenuBar, QMessageBox, QScrollArea, QScrollBar, QSizePolicy, QStatusBar)
from PySide6.QtGui import (QAction, QClipboard, QColorSpace, QGuiApplication, QImage, QImageReader, QImageWriter, QKeySequence, QPalette, QPainter, QPixmap, QScreen, QPainter, QIcon)
from PySide6.QtCore import QDir, QMimeData, QStandardPaths, Qt, Slot, QSize, QPoint
from argparse import ArgumentParser, RawTextHelpFormatter
from PySide6.QtWidgets import (QApplication)
from qt_material import apply_stylesheet

import tempfile
import glob
import importlib, pathlib, sys
import imageio
import numpy as np
import shutil
import random
import string
import sys
from requests import get
import os
import string
import sys
import random
import pathlib
import importlib
import tempfile
import imageio
import imageio.plugins.pillow
import tifffile
from scipy.ndimage import zoom

from tensorflow.keras.models import model_from_json
from tensorflow.keras.layers import Layer, InputSpec
from tensorflow.keras import initializers, regularizers, constraints
from tensorflow.python.keras.utils.generic_utils import get_custom_objects
from tensorflow.keras import backend as K

class InstanceNormalization(Layer):
    def __init__(self,
                 axis=None,
                 epsilon=1e-3,
                 center=True,
                 scale=True,
                 beta_initializer='zeros',
                 gamma_initializer='ones',
                 beta_regularizer=None,
                 gamma_regularizer=None,
                 beta_constraint=None,
                 gamma_constraint=None,
                 **kwargs):
        super(InstanceNormalization, self).__init__(**kwargs)
        self.supports_masking = True
        self.axis = axis
        self.epsilon = epsilon
        self.center = center
        self.scale = scale
        self.beta_initializer = initializers.get(beta_initializer)
        self.gamma_initializer = initializers.get(gamma_initializer)
        self.beta_regularizer = regularizers.get(beta_regularizer)
        self.gamma_regularizer = regularizers.get(gamma_regularizer)
        self.beta_constraint = constraints.get(beta_constraint)
        self.gamma_constraint = constraints.get(gamma_constraint)

    def build(self, input_shape):
        ndim = len(input_shape)
        if self.axis == 0:
            raise ValueError('Axis cannot be zero')

        if (self.axis is not None) and (ndim == 2):
            raise ValueError('Cannot specify axis for rank 1 tensor')

        self.input_spec = InputSpec(ndim=ndim)

        if self.axis is None:
            shape = (1,)
        else:
            shape = (input_shape[self.axis],)

        if self.scale:
            self.gamma = self.add_weight(shape=shape,
                                         name='gamma',
                                         initializer=self.gamma_initializer,
                                         regularizer=self.gamma_regularizer,
                                         constraint=self.gamma_constraint)
        else:
            self.gamma = None
        if self.center:
            self.beta = self.add_weight(shape=shape,
                                        name='beta',
                                        initializer=self.beta_initializer,
                                        regularizer=self.beta_regularizer,
                                        constraint=self.beta_constraint)
        else:
            self.beta = None
        self.built = True

    def call(self, inputs, training=None):
        input_shape = K.int_shape(inputs)
        reduction_axes = list(range(0, len(input_shape)))

        if self.axis is not None:
            del reduction_axes[self.axis]

        del reduction_axes[0]

        mean = K.mean(inputs, reduction_axes, keepdims=True)
        stddev = K.std(inputs, reduction_axes, keepdims=True) + self.epsilon
        normed = (inputs - mean) / stddev

        broadcast_shape = [1] * len(input_shape)
        if self.axis is not None:
            broadcast_shape[self.axis] = input_shape[self.axis]

        if self.scale:
            broadcast_gamma = K.reshape(self.gamma, broadcast_shape)
            normed = normed * broadcast_gamma
        if self.center:
            broadcast_beta = K.reshape(self.beta, broadcast_shape)
            normed = normed + broadcast_beta
        return normed

    def get_config(self):
        config = {
            'axis': self.axis,
            'epsilon': self.epsilon,
            'center': self.center,
            'scale': self.scale,
            'beta_initializer': initializers.serialize(self.beta_initializer),
            'gamma_initializer': initializers.serialize(self.gamma_initializer),
            'beta_regularizer': regularizers.serialize(self.beta_regularizer),
            'gamma_regularizer': regularizers.serialize(self.gamma_regularizer),
            'beta_constraint': constraints.serialize(self.beta_constraint),
            'gamma_constraint': constraints.serialize(self.gamma_constraint)
        }
        base_config = super(InstanceNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

get_custom_objects().update({'InstanceNormalization': InstanceNormalization})

def read_model(directory):
    #weights_path = f'{directory}/weights.h5'
    weights_path = os.path.join( directory, 'weights.h5' )
    if not os.path.isfile(weights_path):
        print( f'Failed to find weights from file {weights_path}' )
        return None

    #json_path = f'{directory}/js.json'
    json_path = os.path.join( directory, 'js.json' )
    if not os.path.isfile(json_path):
        print( f'Failed to find model from file {json_path}' )
        return None

    js_file = open( json_path, 'r' )
    model_json = js_file.read()
    js_file.close()
    model = model_from_json( model_json, custom_objects={"InstanceNormalization": InstanceNormalization} )
    model.load_weights( weights_path )
    #print( model.summary() )
    return model

def s9_denoising( model, input_image_path, output_image_path ):
    print( f'Trying to denoising image from {input_image_path}' )
    x = imageio.imread( input_image_path )
    x = np.squeeze( x )

    if not len(x.shape) == 2:
        print( f'Input image with shape {x.shape} is not a gray image' )
        return False

    x = np.reshape(x, (1,) + x.shape + (1,) )
    x = np.asarray( x, dtype='float32' )
    x /= np.amax(x) + 1.0e-10
    y = model.predict( x )
    y = np.squeeze( y )
    y /= np.amax(y) + 1.0e-10
    imageio.imwrite( output_image_path, np.asarray( y*255.0, dtype='uint8' ) )
    print( f'Writing output image to {output_image_path}' )
    return True

uncultivated_widget = None
s9_denoising_model = None

def update_widget_image( image_widget, output_image_path ):
    global uncultivated_widget
    if uncultivated_widget is None:
        uncultivated_widget = image_widget
    uncultivated_widget.show()
    uncultivated_widget.update_content_file( output_image_path, rescaling_flag=False, in_a_new_window=False );
    QApplication.processEvents()

def s9_implementation( image_widget ):

    global s9_denoising_model
    if s9_denoising_model is None: # load module dynamically
        if getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS'):
            current_directory = os.path.dirname( sys.executable ) # <-- updated
        else:
            current_directory = os.path.dirname(os.path.realpath(__file__))
        #current_directory = os.path.dirname( sys.executable ) # <-- updated
        local_s9_denoising_model_path = os.path.join(current_directory, 'models', 's9_model' )
        s9_denoising_model = read_model( local_s9_denoising_model_path )

    '''
    input_image_path = image_widget.get_snapshot_file()
    random_file_prefix = ''.join(random.choices(string.ascii_uppercase + string.digits, k = 13))
    output_image_path = os.path.join( tempfile.gettempdir(), f'{random_file_prefix}_s9_denoising_cache.png' )

    if s9_denoising( s9_denoising_model, input_image_path, output_image_path ):
        image_widget.update_content_file( output_image_path )
    '''

    if not image_widget.is_tiff: # case of single image
        input_image_path = image_widget.get_snapshot_file()
        random_file_prefix = ''.join(random.choices(string.ascii_uppercase + string.digits, k = 13))
        output_image_path = os.path.join( tempfile.gettempdir(), f'{random_file_prefix}_s9_denoising_cache.png' )

        if s9_denoising( s9_denoising_model, input_image_path, output_image_path ):
            image_widget.update_content_file( output_image_path )
    else: # case of tif file
        image_widget.tiff_data_denoised = []
        n, _, _ = image_widget.tiff_data.shape
        for idx in range( n ):
            tmp_file_path = image_widget.get_new_snapshot_file()
            imageio.imwrite( tmp_file_path, image_widget.tiff_data[idx] )

            input_image_path = image_widget.get_snapshot_file()
            random_file_prefix = ''.join(random.choices(string.ascii_uppercase + string.digits, k = 13))
            output_image_path = os.path.join( tempfile.gettempdir(), f'{random_file_prefix}_s9_denoising_cache.png' )

            if s9_denoising( s9_denoising_model, tmp_file_path, output_image_path ):
                image_widget.tiff_data_denoised.append( imageio.imread( output_image_path ) )
                update_widget_image( image_widget, output_image_path )

            else:
                break

        if len( image_widget.tiff_data_denoised ):
            image_widget.tiff_data_denoised = np.asarray( image_widget.tiff_data_denoised )

def s9_interface():
    def detailed_implementation( image_widget ):
        def fun():
            return s9_implementation( image_widget )
        return fun

    return 'S9Denoising', detailed_implementation

def sd_denoising( model, input_image_path, output_image_path ):
    x = imageio.imread( input_image_path )
    x = np.squeeze( x )

    if not len(x.shape) == 2:
        print( f'Input image with shape {x.shape} is not a gray image' )
        return False

    x = zoom( x, 2, order=3 )
    x = np.reshape(x, (1,) + x.shape + (1,) )
    x = np.asarray( x, dtype='float32' )
    x /= np.amax(x) + 1.0e-10
    y = model.predict( x )
    #y = make_zoom_prediction( model, x )

    y = np.squeeze( y )
    y /= np.amax(y) + 1.0e-10
    imageio.imwrite( output_image_path, np.asarray( y*255.0, dtype='uint8' ) )
    return True

sd_denoising_model = None
def sd_implementation( image_widget ):

    global sd_denoising_model
    if sd_denoising_model is None: # load module dynamically
        #current_directory = os.path.dirname(os.path.realpath(__file__))
        if getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS'):
            current_directory = os.path.dirname( sys.executable ) # <-- updated
        else:
            current_directory = os.path.dirname(os.path.realpath(__file__))
        local_sd_denoising_model_path = os.path.join(current_directory, 'models', 'debora_model' )
        sd_denoising_model = read_model( local_sd_denoising_model_path )

    if not image_widget.is_tiff: # case of single image
        input_image_path = image_widget.get_snapshot_file()
        random_file_prefix = ''.join(random.choices(string.ascii_uppercase + string.digits, k = 13))
        output_image_path = os.path.join( tempfile.gettempdir(), f'{random_file_prefix}_debora_denoising_cache.png' )

        if sd_denoising( sd_denoising_model, input_image_path, output_image_path ):
            image_widget.update_content_file( output_image_path )
    else: # case of tif file
        image_widget.tiff_data_denoised = []
        n, _, _ = image_widget.tiff_data.shape
        for idx in range( n ):
            tmp_file_path = image_widget.get_new_snapshot_file()
            imageio.imwrite( tmp_file_path, image_widget.tiff_data[idx] )

            input_image_path = image_widget.get_snapshot_file()
            random_file_prefix = ''.join(random.choices(string.ascii_uppercase + string.digits, k = 13))
            output_image_path = os.path.join( tempfile.gettempdir(), f'{random_file_prefix}_debora_denoising_cache.png' )

            if sd_denoising( sd_denoising_model, tmp_file_path, output_image_path ):
                image_widget.tiff_data_denoised.append( imageio.imread( output_image_path ) )
                update_widget_image( image_widget, output_image_path )

            else:
                break

        if len( image_widget.tiff_data_denoised ):
            image_widget.tiff_data_denoised = np.asarray( image_widget.tiff_data_denoised )

ABOUT = "Deep Image Denoising."

class ImagePlayer(QMainWindow):
    def __init__(self, parent=None):
        super().__init__(parent)
        self._scale_factor = 1.0
        self._first_file_dialog = True
        self._image_label = QLabel()
        self._image_label.setBackgroundRole(QPalette.Base)
        self._image_label.setSizePolicy(QSizePolicy.Ignored, QSizePolicy.Ignored)
        self._image_label.setScaledContents(True)

        self._scroll_area = QScrollArea()
        self._scroll_area.setBackgroundRole(QPalette.Dark)
        self._scroll_area.setWidget(self._image_label)
        self._scroll_area.setVisible(False)
        self.setCentralWidget(self._scroll_area)

        self._create_actions()

        self.resize(QGuiApplication.primaryScreen().availableSize() ) # PySide6.QtCore.QSize(1920, 1200)

        # data zone
        self.tmp_image_counter = 0
        self.current_cached_image_path = None
        self.current_image_presented = None
        self.random_file_prefix = ''.join(random.choices(string.ascii_uppercase + string.digits, k = 13))

        dir_path = os.path.dirname(os.path.realpath(__file__))
        self._plugin_list = []
        self.load_plugins(os.path.join(dir_path, 'plugins'))

        self._model_list = []
        self.load_models(os.path.join(dir_path, 'models'))

        # drop image file to open, accept
        self.setAcceptDrops(True)

        icon_path = self.load_plugins(os.path.join(dir_path, 'src', 'logo.png'))
        self.setWindowIcon(QIcon(icon_path))

        # home model path
        self.user_model_path = os.path.join( os.path.expanduser('~'), '.deepoffice', 'imageplayer', 'model' )
        #if not os.path.exists(self.user_model_path):
        #    os.makedirs( self.user_model_path )

        # denoising stacked tiff file
        self.is_tiff = False
        self.tiff_path = None
        self.tiff_data = None
        self.tiff_data_denoised = None
        self.tiff_data_denoised_presentation_index = None

    # download model files from github release
    def download_remote_model(self, model_name, model_url):
        model_path = os.path.join( self.user_model_path, model_name )

        if not os.path.exists(model_path):
            os.makedirs(model_path)

        file_name = model_url.rsplit('/', 1)[1]
        local_model_path = os.path.join( model_path, file_name )
        if not os.path.isfile(local_model_path):
            print( f'downloading model file {local_model_path} from {model_url}' )
            with open(local_model_path, "wb") as file:
                response = get(model_url)
                file.write(response.content)
            print( f'downloaded model file {local_model_path} from {model_url}' )

        return local_model_path

    # interface to plugins 1
    def get_snapshot_file( self, index=None ):
        return self.get_new_snapshot_file( self.tmp_image_counter-1 )

    # interface to plugins 2: a single file
    def update_content_file(self, fileName, rescaling_flag=False, in_a_new_window=True):
        self.load_file(fileName, rescaling_flag)
        '''
        if not in_a_new_window:
            self.load_file( fileName, rescaling_flag )
        else:
            image_player = ImagePlayer()
            image_player.show()
            image_player.update_content_file( fileName, rescaling_flag=rescaling_flag, in_a_new_window=False );
        '''

    # interface to plugins 3: many files
    def update_content_files(self, fileNames, rescaling_flag=False, in_a_new_window=True):
        for fileName in fileNames:
            self.update_content_file( fileName, rescaling_flag, in_a_new_window )

    def get_new_snapshot_file( self, index=None ):
        if index is None:
            index = self.tmp_image_counter
            self.tmp_image_counter += 1
        tmp_dir = tempfile.gettempdir()
        tmp_png_file = os.path.join( tmp_dir, f'{self.random_file_prefix}_{index}.png' )
        return tmp_png_file

    def dragEnterEvent(self, e):
        if e.mimeData().hasUrls():
            e.accept()
        else:
            e.ignore()
            print( f'ignoring {e.mimeData().text()}' )

    def dropEvent(self, e):
        if e.mimeData().hasUrls():
            for url in  e.mimeData().urls():
                self.load_file( url.path() )
                break

    def wheelEvent(self, event):
        numDegrees = event.angleDelta() / 8.0

        if not numDegrees.isNull():
            x = numDegrees.x() / 150.0
            y = numDegrees.y() / 150.0

            new_scale = x + 1.0
            if abs(y) > abs(x):
                new_scale = y + 1.0

            new_scale = min( 2.0, max( 0.5, new_scale) )
            self._scale_image(new_scale)

        event.accept()

    def _show_tiff( self, index=0 ):
        img_to_show = self.tiff_data
        if len(self.tiff_data.shape) == 3:
            n, _, _ = self.tiff_data.shape
            index = index % n
            img_to_show = self.tiff_data[index]

        tmp_file_path = self.get_new_snapshot_file()
        imageio.imwrite( tmp_file_path, img_to_show )
        self.load_file( tmp_file_path )

    def load_file(self, fileName, rescaling_flag=True):
        _, file_extension = os.path.splitext( fileName )

        if file_extension == '.tif' or file_extension == '.tiff':
            print( f'reading tifffile from {fileName}' )
            self.is_tiff = True
            self.tiff_path = fileName
            self.tiff_data = np.squeeze( tifffile.imread( fileName ) )
            self.tiff_data_denoised = None

            if self.tiff_data.size == 0:
                return False

            self._show_tiff(0)
            return True

        reader = QImageReader(fileName)
        reader.setAutoTransform(True)
        new_image = reader.read()
        native_filename = QDir.toNativeSeparators(fileName)
        if new_image.isNull():
            error = reader.errorString()
            #QMessageBox.information(self, QGuiApplication.applicationDisplayName(), f"Cannot load {native_filename}: {error}")
            error_message = f'cannot open file {native_filename} -->  Error: {error}'
            self.statusBar().showMessage(error_message)
            return False
        self.current_image_presented = fileName
        self._set_image(new_image)
        self.setWindowFilePath(fileName)

        w = max( self._image.width() * 1.15, 300 )
        h = max( self._image.height() * 1.15, 300 )
        d = self._image.depth()
        max_size = QGuiApplication.primaryScreen().availableSize()
        self._current_size = QSize(min(w, max_size.width()), min( h, max_size.height()) )
        self.resize( self._current_size )
        if rescaling_flag:
            self._scale_image( self._scale_factor )
        message = f'Opened "{native_filename}": {self._image.width()}x{self._image.height()}'
        self.statusBar().showMessage(message)

        self._save_tmp_file()
        return True

    def _set_image(self, new_image):
        self._image = new_image
        if self._image.colorSpace().isValid():
            self._image.convertToColorSpace(QColorSpace.SRgb)
        self._image_label.setPixmap(QPixmap.fromImage(self._image))
        self._scroll_area.setAlignment(Qt.AlignCenter)
        self._scale_factor = 1.0

        self._scroll_area.setVisible(True)
        self._print_act.setEnabled(True)
        self._update_actions()

        if not self._fit_to_window_act.isChecked():
            self._image_label.adjustSize()

    def _save_file(self, fileName):
        _, new_extension = os.path.splitext( fileName )
        if new_extension == '.tif' or new_extension == '.tiff':
            if self.tiff_data_denoised is not None:
                tifffile.imwrite( fileName, self.tiff_data_denoised )
                print( f'writting denoised tiff data to {fileName}' )
                return True

        _, current_extension = os.path.splitext( self.current_image_presented )
        if (current_extension == new_extension ):
            if self.current_image_presented != fileName:
                shutil.copy( self.current_image_presented, fileName )
            return True

        writer = QImageWriter(fileName)

        native_filename = QDir.toNativeSeparators(fileName)
        if not writer.write(self._image):
            error = writer.errorString()
            message = f"Cannot write {native_filename}: {error}"
            QMessageBox.information(self, QGuiApplication.applicationDisplayName(), message)
            self.statusBar().showMessage( message );
            return False
        return True

    def _save_tmp_file( self ):
        tmp_png_file = self.get_new_snapshot_file()
        if not self._save_file(tmp_png_file):
            print( f'Failed saving tmp file to {tmp_png_file}' )
            return None
        print( f'saving tmp file: {tmp_png_file}' )
        self.tmp_image_counter += 1
        self.current_cached_image_path = tmp_png_file
        return tmp_png_file

    @Slot()
    def _undo( self ):
        prev_cache_file = self.get_new_snapshot_file( self.tmp_image_counter-2 )
        if os.path.isfile( prev_cache_file ):
            self.load_file( prev_cache_file )
            self.tmp_image_counter -=  2
            print( f'Undo: updating image counter to {self.tmp_image_counter}')
        else:
            print( f'cannot load {prev_cache_file=}' )

    @Slot()
    def _redo( self ):
        next_cache_file = self.get_new_snapshot_file( self.tmp_image_counter )
        if os.path.isfile( next_cache_file ):
            self.load_file( next_cache_file )
            print( f'Redo: updating image counter to {self.tmp_image_counter}')
        else:
            print( f'cannot load {next_cache_file=}' )

    @Slot()
    def _clean_tmp_file( self ):
        for idx in range( self.tmp_image_counter ):
            tmp_dir = tempfile.gettempdir()
            tmp_png_file = os.path.join( tmp_dir, f'{self.random_file_prefix}_{idx}.png' )
            if os.path.isfile( tmp_png_file ):
                os.remove( tmp_png_file )
                print( f'removing file {tmp_png_file}' )
        self.tmp_image_counter = 0
        self.current_cached_image_path = None

    @Slot()
    def _open(self):
        dialog = QFileDialog(self, "Open File")
        self._initialize_image_filedialog(dialog, QFileDialog.AcceptOpen)
        while (dialog.exec() == QDialog.Accepted
               and not self.load_file(dialog.selectedFiles()[0])):
            pass

    @Slot()
    def _save_as(self):
        dialog = QFileDialog(self, "Save File As")
        self._initialize_image_filedialog(dialog, QFileDialog.AcceptSave)
        while (dialog.exec() == QDialog.Accepted and not self._save_file(dialog.selectedFiles()[0])):
            pass

    @Slot()
    def _print_(self):
        printer = QPrinter()
        dialog = QPrintDialog(printer, self)
        if dialog.exec() == QDialog.Accepted:
            painter = QPainter(printer)
            pixmap = self._image_label.pixmap()
            rect = painter.viewport()
            size = pixmap.size()
            size.scale(rect.size(), Qt.KeepAspectRatio)
            painter.setViewport(rect.x(), rect.y(), size.width(), size.height())
            painter.setWindow(pixmap.rect())
            painter.drawPixmap(0, 0, pixmap)
            painter.end()

    @Slot()
    def _copy(self):
        QGuiApplication.clipboard().setImage(self._image)

    @Slot()
    def _paste(self):
        new_image = QGuiApplication.clipboard().image()
        if new_image.isNull():
            self.statusBar().showMessage("No image in clipboard")
        else:
            self._set_image(new_image)
            self.setWindowFilePath('')
            w = new_image.width()
            h = new_image.height()
            d = new_image.depth()
            message = f"Obtained image from clipboard, {w}x{h}, Depth: {d}"
            self.statusBar().showMessage(message)

    @Slot()
    def _zoom_in(self):
        if self.current_image_presented is not None:
            self._scale_image(1.25)

    @Slot()
    def _zoom_out(self):
        if self.current_image_presented is not None:
            self._scale_image(0.8)

    @Slot()
    def _normal_size(self):
        if self.current_image_presented is not None:
            self._image_label.adjustSize()
            self._scale_factor = 1.0

    @Slot()
    def _fit_to_window(self):
        if self.current_image_presented is not None:
            fit_to_window = self._fit_to_window_act.isChecked()
            self._scroll_area.setWidgetResizable(fit_to_window)
            if not fit_to_window:
                self._normal_size()
            self._update_actions()

    @Slot()
    def _about(self):
        QMessageBox.about(self, "About Image Viewer", ABOUT)

    def _create_actions(self):
        file_menu = self.menuBar().addMenu("&File")

        self._open_act = file_menu.addAction("&Open...")
        self._open_act.triggered.connect(self._open)
        self._open_act.setShortcut(QKeySequence.Open)

        self._save_as_act = file_menu.addAction("&Save As...")
        self._save_as_act.triggered.connect(self._save_as)
        self._save_as_act.setEnabled(False)

        self._print_act = file_menu.addAction("&Print...")
        self._print_act.triggered.connect(self._print_)
        self._print_act.setShortcut(QKeySequence.Print)
        self._print_act.setEnabled(False)

        file_menu.addSeparator()

        self._exit_act = file_menu.addAction("E&xit")
        self._exit_act.triggered.connect(self.close)
        self._exit_act.triggered.connect(self._clean_tmp_file)
        self._exit_act.setShortcut("Ctrl+Q")

        edit_menu = self.menuBar().addMenu("&Edit")

        self._copy_act = edit_menu.addAction("Undo")
        self._copy_act.triggered.connect(self._undo)

        self._copy_act = edit_menu.addAction("Redo")
        self._copy_act.triggered.connect(self._redo)

        self._copy_act = edit_menu.addAction("&Copy")
        self._copy_act.triggered.connect(self._copy)
        self._copy_act.setShortcut(QKeySequence.Copy)
        self._copy_act.setEnabled(False)

        self._paste_act = edit_menu.addAction("&Paste")
        self._paste_act.triggered.connect(self._paste)
        self._paste_act.setShortcut(QKeySequence.Paste)

        view_menu = self.menuBar().addMenu("&View")

        self._zoom_in_act = view_menu.addAction("Zoom &In (25%)")
        self._zoom_in_act.setShortcut(QKeySequence.ZoomIn)
        self._zoom_in_act.triggered.connect(self._zoom_in)
        self._zoom_in_act.setEnabled(False)

        self._zoom_out_act = view_menu.addAction("Zoom &Out (25%)")
        self._zoom_out_act.triggered.connect(self._zoom_out)
        self._zoom_out_act.setShortcut(QKeySequence.ZoomOut)
        self._zoom_out_act.setEnabled(False)

        self._normal_size_act = view_menu.addAction("&Normal Size")
        self._normal_size_act.triggered.connect(self._normal_size)
        self._normal_size_act.setShortcut("Ctrl+S")
        self._normal_size_act.setEnabled(False)

        view_menu.addSeparator()

        self._fit_to_window_act = view_menu.addAction("&Fit to Window")
        self._fit_to_window_act.triggered.connect(self._fit_to_window)
        self._fit_to_window_act.setEnabled(False)
        self._fit_to_window_act.setCheckable(True)
        self._fit_to_window_act.setShortcut("Ctrl+F")

        self.plugin_menu = self.menuBar().addMenu("&Plugins")

        self.model_menu = self.menuBar().addMenu("&Models")

        self.s9_act = self.model_menu.addAction("Denoising S9")
        self.s9_act.triggered.connect(self._s9_denosing)
        self.s9_act.setEnabled(True)

        self.sd_act = self.model_menu.addAction("Denoising SD")
        self.sd_act.triggered.connect(self._sd_denosing)
        self.sd_act.setEnabled(True)

        help_menu = self.menuBar().addMenu("&Help")

        about_act = help_menu.addAction("&About")
        about_act.triggered.connect(self._about)
        about_qt_act = help_menu.addAction("About &Qt")
        about_qt_act.triggered.connect(QApplication.aboutQt)

    def load_plugin( self, plugin_path ):
        module_name = pathlib.Path(plugin_path).stem
        module = importlib.import_module(module_name)
        plugin_name, event = module.interface()
        _act = self.plugin_menu.addAction( plugin_name )
        _act.triggered.connect( event(self) )
        self._plugin_list.append( plugin_name )

    def load_plugins( self, folder ):
        pass
        '''
        sys.path.append( folder )
        plugin_paths = glob.glob( f'{folder}/*.py' ) +  glob.glob( f'{str(pathlib.Path.home())}/.deepoffice/plugins/*.py' )
        for plugin_path in plugin_paths:
            self.load_plugin( plugin_path )
        '''

    def load_model( self, model_path ):
        module_name = pathlib.Path(model_path).stem
        module = importlib.import_module(module_name)
        model_name, event = module.interface()
        _act = self.model_menu.addAction( model_name )
        _act.triggered.connect( event(self) )
        self._model_list.append( model_name )

    def load_models( self, folder ):
        pass
        '''
        sys.path.append( folder )
        model_paths = glob.glob( f'{folder}/*.py' ) +  glob.glob( f'{str(pathlib.Path.home())}/.deepoffice/models/*.py' )
        for model_path in model_paths:
            self.load_model( model_path )
        '''

    def _update_actions(self):
        has_image = not self._image.isNull()
        self._save_as_act.setEnabled(has_image)
        self._copy_act.setEnabled(has_image)
        enable_zoom = not self._fit_to_window_act.isChecked()
        self._zoom_in_act.setEnabled(enable_zoom)
        self._zoom_out_act.setEnabled(enable_zoom)
        self._normal_size_act.setEnabled(enable_zoom)

    def _scale_image(self, factor):
        if self.current_image_presented is not None:
            self._scale_factor *= factor
            new_size = self._scale_factor * self._image_label.pixmap().size()
            self._image_label.resize(new_size)

            self._adjust_scrollbar(self._scroll_area.horizontalScrollBar(), factor)
            self._adjust_scrollbar(self._scroll_area.verticalScrollBar(), factor)

            self._zoom_in_act.setEnabled(self._scale_factor < 3.0)
            self._zoom_out_act.setEnabled(self._scale_factor > 0.333)

            self._current_size *= factor

    def _adjust_scrollbar(self, scrollBar, factor):
        pos = int(factor * scrollBar.value() + ((factor - 1) * scrollBar.pageStep() / 2))
        print( f'adjusting scrollbar to {pos=}' )
        scrollBar.setValue(pos)

    def _initialize_image_filedialog(self, dialog, acceptMode):
        if self._first_file_dialog:
            self._first_file_dialog = False
            locations = QStandardPaths.standardLocations(QStandardPaths.PicturesLocation)
            directory = locations[-1] if locations else QDir.currentPath()
            dialog.setDirectory(directory)

        mime_types = [m.data().decode('utf-8') for m in QImageWriter.supportedMimeTypes()]
        mime_types.sort()

        dialog.setMimeTypeFilters(mime_types)
        dialog.setAcceptMode(acceptMode)
        if acceptMode == QFileDialog.AcceptSave:
            dialog.setDefaultSuffix("png")

    def _s9_denosing( self ):
        s9_implementation( self )

    def _sd_denosing( self ):
        sd_implementation( self )

    # TODO: left/right arrow to navigate

if __name__ == '__main__':
    arg_parser = ArgumentParser(description="Image Viewer", formatter_class=RawTextHelpFormatter)
    arg_parser.add_argument('file', type=str, nargs='?', help='Image file')
    args = arg_parser.parse_args()

    app = QApplication(sys.argv)
    image_player = ImagePlayer()
    extra = { 'danger': '#dc3545', 'warning': '#ffc107', 'success': '#17a2b8', 'font-family': 'Roboto', }
    apply_stylesheet(app, 'light_cyan_500.xml', invert_secondary=True, extra=extra)

    if args.file and not image_player.load_file(args.file):
        sys.exit(-1)

    image_player.show()
    sys.exit(app.exec())


import os
import shutil
import pandas as pd
import numpy as np
import time
import multiprocessing as mp
import astropy.coordinates as astroCoords
import astropy.units as u
import csv
import trajectoryFiltering as tf
from kbmodpy import kbmod as kb
from astropy.io import fits
from astropy.wcs import WCS
from sklearn.cluster import DBSCAN
from skimage import measure
from analysis_utils import analysis_utils, \
    return_indices, stamp_filter_parallel
from collections import OrderedDict

class run_search(analysis_utils):

    def __init__(self, v_list, ang_list, num_obs):

        """
        Input
        --------

        v_list : list

            [min_velocity, max_velocity, velocity_steps]

        ang_list: list

            [radians below ecliptic,
             radians above ecliptic,
             steps]

        num_obs : integer

            Number of images a trajectory must be unmasked.
        """

        self.v_arr = np.array(v_list)
        self.ang_arr = np.array(ang_list)
        self.num_obs = num_obs

        return

    def run_search(self, im_filepath, res_filepath, out_suffix, time_file,
                   likelihood_level=10., mjd_lims=None, num_fakes=25,
                   rand_seed=42):

        visit_nums, visit_times = np.genfromtxt(time_file, unpack=True)
        image_time_dict = OrderedDict()
        for visit_num, visit_time in zip(visit_nums, visit_times):
            image_time_dict[str(int(visit_num))] = visit_time

        chunk_size = 100000

        start = time.time()

        patch_visits = sorted(os.listdir(im_filepath))
        patch_visit_ids = np.array([int(visit_name[1:7]) for visit_name in patch_visits])
        patch_visit_times = np.array([image_time_dict[str(visit_id)] for visit_id in patch_visit_ids])

        if mjd_lims is None:
            use_images = patch_visit_ids
        else:
            visit_only = np.where(((patch_visit_times > mjd_lims[0])
                                   & (patch_visit_times < mjd_lims[1])))[0]
            print(visit_only)
            use_images = patch_visit_ids[visit_only]

        image_mjd = np.array([image_time_dict[str(visit_id)] for visit_id in use_images])
        times = image_mjd - image_mjd[0]

        flags = ~0 # mask pixels with any flags
        flag_exceptions = [32,39] # unless it has one of these special combinations of flags
        master_flags = int('100111', 2) # mask any pixels which have any of
        # these flags in more than two images

        hdulist = fits.open('%s/v%i-fg.fits' % (im_filepath, use_images[0]))
        f0 = hdulist[0].header['FLUXMAG0']
        w = WCS(hdulist[1].header)
        ec_angle = self.calc_ecliptic_angle(w)
        del(hdulist)

        images = [kb.layered_image('%s/v%i-fg.fits' % (im_filepath, f)) for f in np.sort(use_images)]
        print('Images Loaded')

        p = kb.psf(1.4)

        # Add fakes steps
        print('Adding fake objects')
        x_fake_range = (5, 3650)
        y_fake_range = (5, 3650)
        angle_range = (ec_angle-(np.pi/15.), ec_angle+(np.pi/15.))
        velocity_range = (100, 500)
        mag_range = (20, 26)

        fake_results = []
        fake_output = []

        np.random.seed(rand_seed)
        for val in range(num_fakes):
            traj = kb.trajectory()
            traj.x = int(np.random.uniform(*x_fake_range))
            traj.y = int(np.random.uniform(*y_fake_range))
            ang = np.random.uniform(*angle_range)
            vel = np.random.uniform(*velocity_range)
            traj.x_v = vel*np.cos(ang)
            traj.y_v = vel*np.sin(ang)
            mag_val = np.random.uniform(*mag_range)
            traj.flux = f0*np.power(10, -0.4*mag_val)
            fake_results.append(traj)
            fake_output.append([traj.x, traj.y, traj.x_v, traj.y_v, traj.flux, mag_val])

        for fake_obj in fake_results:
            tf.add_trajectory(images, fake_obj, p, times)

        stack = kb.image_stack(images)
        del(images)
        stack.apply_mask_flags(flags, flag_exceptions)
        stack.apply_master_mask(master_flags, 2)

        stack.grow_mask()
        stack.grow_mask()

        # stack.apply_mask_threshold(120.)

        stack.set_times(times)
        print("Times set")
        x_size = stack.get_width()
        y_size = stack.get_width()

        search = kb.stack_search(stack, p)
        del(stack)
        ang_min = ec_angle - self.ang_arr[0]
        ang_max = ec_angle + self.ang_arr[1]
        vel_min = self.v_arr[0]
        vel_max = self.v_arr[1]
        print("Starting Search")
        print('---------------------------------------')
        param_headers = ("Ecliptic Angle", "Min. Search Angle", "Max Search Angle",
                         "Min Velocity", "Max Velocity")
        param_values = (ec_angle, ang_min, ang_max, vel_min, vel_max)
        for header, val in zip(param_headers, param_values):
            print('%s = %.4f' % (header, val))
        search.gpu(int(self.ang_arr[2]),int(self.v_arr[2]),ang_min,ang_max,
                   vel_min,vel_max,int(self.num_obs))

        keep_stamps = []
        keep_snr = []
        keep_new_lh = []
        keep_results = []
        keep_times = []
        memory_error = False
        keep_lc = []
        filter_stats = np.zeros(4)

        likelihood_limit = False
        res_num = 0
        chunk_size = 500000
        print('---------------------------------------')
        print("Processing Results")
        print('---------------------------------------')
        while likelihood_limit is False:
            pool = mp.Pool(processes=16)
            results = search.get_results(res_num,chunk_size)
            chunk_headers = ("Chunk Start", "Chunk Size", "Chunk Max Likelihood",
                             "Chunk Min. Likelihood")
            chunk_values = (res_num, len(keep_results), results[0].lh, results[-1].lh)
            for header, val, in zip(chunk_headers, chunk_values):
                if type(val) == np.int:
                    print('%s = %i' % (header, val))
                else:
                    print('%s = %.2f' % (header, val))
            print('---------------------------------------')
            psi_curves = []
            phi_curves = []
            for line in results:
                psi_curve, phi_curve = search.lightcurve(line)
                psi_curves.append(np.array(psi_curve).flatten())
                phi_curve = np.array(phi_curve).flatten()
                phi_curve[phi_curve == 0.] = 99999999.
                phi_curves.append(phi_curve)
                if line.lh < likelihood_level:
                    likelihood_limit = True
                    break
            keep_idx_results = pool.starmap_async(return_indices,
                                                  zip(psi_curves, phi_curves,
                                                      [j for j in range(len(psi_curves))]))
            pool.close()
            pool.join()
            keep_idx_results = keep_idx_results.get()

            filter_stats[0] += len(psi_curves)
            if len(keep_idx_results[0]) < 3:
                keep_idx_results = [(0, [-1], 0.)]

            for result_on in range(len(psi_curves)):
                if keep_idx_results[result_on][1][0] == -1:
                    continue
                elif len(keep_idx_results[result_on][1]) < 3:
                    continue
                elif keep_idx_results[result_on][2] < likelihood_level:
                    continue
                else:
                    keep_idx = keep_idx_results[result_on][1]
                    new_likelihood = keep_idx_results[result_on][2]
                    keep_results.append(results[result_on])
                    keep_new_lh.append(new_likelihood)
                    stamps = search.sci_stamps(results[result_on], 10)
                    stamp_arr = np.array([np.array(stamps[s_idx]) for s_idx in keep_idx])
                    keep_stamps.append(np.sum(stamp_arr, axis=0))
                    keep_lc.append((psi_curves[result_on]/phi_curves[result_on])[keep_idx])
                    keep_snr.append((psi_curves[result_on]/np.sqrt(phi_curves[result_on]))[keep_idx])
                    #keep_times.append(image_mjd[keep_idx])
                    keep_times.append(keep_idx)

            # if len(keep_results) > 800000:
            #     with open('%s/memory_error_tr_%s.txt' %
            #               (res_filepath, out_suffix), 'w') as f:
            #         f.write('In %i total results, %i were kept. Needs manual look.' %
            #                 (res_num + chunk_size, len(keep_results)))
            #     memory_error = True
            #     likelihood_limit = True

            # if res_num+chunk_size >= 8000000:
            #     likelihood_level = 20.
            #     with open('%s/overload_error_tr_%s.txt' %
            #               (res_filepath, out_suffix), 'w') as f:
            #         f.write('In %i total results, %i were kept. Likelihood level down to %f.' %
            #                 (res_num + chunk_size, len(keep_results), line.lh))

            res_num += chunk_size

        del(search)

        lh_sorted_idx = np.argsort(np.array(keep_new_lh))[::-1]
        filter_stats[1] = len(lh_sorted_idx)

        if len(lh_sorted_idx) > 0:
            print("Stamp filtering %i results" % len(lh_sorted_idx))
            pool = mp.Pool(processes=16)
            stamp_filt_pool = pool.map_async(stamp_filter_parallel,
                                             np.array(keep_stamps)[lh_sorted_idx])
            pool.close()
            pool.join()
            stamp_filt_results = stamp_filt_pool.get()
            stamp_filt_idx = lh_sorted_idx[np.where(np.array(stamp_filt_results) == 1)]
            filter_stats[2] = len(stamp_filt_idx)
            if len(stamp_filt_idx) > 0:
                print("Clustering %i results" % len(stamp_filt_idx))
                cluster_idx = self.cluster_results(np.array(keep_results)[stamp_filt_idx],
                                                   x_size, y_size, [vel_min, vel_max],
                                                   [ang_min, ang_max])
                final_results = stamp_filt_idx[cluster_idx]
            else:
                cluster_idx = []
                final_results = []
            del(cluster_idx)
            del(stamp_filt_results)
            del(stamp_filt_idx)
            del(stamp_filt_pool)
        else:
            final_results = lh_sorted_idx

        print('Keeping %i results' % len(final_results))
        filter_stats[3] = len(final_results)

        np.savetxt('%s/results_%s.txt' % (res_filepath, out_suffix),
                   np.array(keep_results)[final_results], fmt='%s')
        np.savetxt('%s/results_fakes_%s.txt' % (res_filepath, out_suffix),
                   np.array(fake_output), header='x,y,xv,yv,flux,mag')
        # np.savetxt('%s/lc_%s.txt' % (res_filepath, out_suffix),
        #            np.array(keep_lc)[final_results], fmt='%s')
        with open('%s/lc_%s.txt' % (res_filepath, out_suffix), 'w') as f:
            writer = csv.writer(f)
            writer.writerows(np.array(keep_lc)[final_results])
        with open('%s/snr_%s.txt' % (res_filepath, out_suffix), 'w') as f:
            writer = csv.writer(f)
            writer.writerows(np.array(keep_snr)[final_results])
        # np.savetxt('%s/times_%s.txt' % (res_filepath, out_suffix),
        #            np.array(keep_times)[final_results], fmt='%s')
        with open('%s/times_%s.txt' % (res_filepath, out_suffix), 'w') as f:
            writer = csv.writer(f)
            writer.writerows(np.array(keep_times)[final_results])
        np.savetxt('%s/filtered_likes_%s.txt' % (res_filepath, out_suffix),
                   np.array(keep_new_lh)[final_results], fmt='%.4f')
        np.savetxt('%s/ps_%s.txt' % (res_filepath, out_suffix),
                   np.array(keep_stamps).reshape(len(keep_stamps), 441)[final_results], fmt='%.4f')
        np.savetxt('%s/filt_stats_%s.txt' % (res_filepath, out_suffix), np.array(filter_stats), fmt='%i')

        end = time.time()

        del(keep_stamps)
        del(keep_times)
        del(keep_results)
        del(keep_new_lh)
        del(keep_lc)

        print("Time taken for patch: ", end-start)

import numpy,import astropy,from astropy