"""
Utilities for fitting stacked (drizzled) spectra
"""
from collections import OrderedDict
from imp import reload

import astropy.io.fits as pyfits
import astropy.units as u

import numpy as np

from . import utils
from .utils import GRISM_COLORS, GRISM_MAJOR, GRISM_LIMITS, DEFAULT_LINE_LIST

from .fitting import GroupFitter

def make_templates(grism='G141', return_lists=False, fsps_templates=False,
                   line_list=DEFAULT_LINE_LIST):
    """Generate template savefile

    This script generates the template sets with the emission line
    complexes and with individual lines.

    Parameters
    ----------
    grism : str
        Grism of interest, which defines what FWHM to use for the line
        templates.

    return_lists : bool
        Return the templates rather than saving them to a file

    Returns
    -------
    t_complexes, t_lines : list
        If `return` then return two lists of templates.  Otherwise,
        store them to a `~numpy` save file "templates_{fwhm}.npy".

    """

    from .multifit import MultiBeam

    if grism == 'G141':    # WFC3/IR
        fwhm = 1100
    elif grism == 'G800L': # ACS/UVIS
        fwhm = 1400
    elif grism == 'G280':  # WFC3/UVIS
        fwhm = 1500
    elif grism == 'GRISM': # WFIRST
        fwhm = 350
    else:
        fwhm = 700 # G102

    # Line complex templates
    t_complexes = utils.load_templates(fwhm=fwhm, line_complexes=True,
                                           fsps_templates=fsps_templates)

    # Individual lines
    # line_list = ['SIII', 'SII', 'Ha', 'OI-6302', 'OIII', 'Hb',
    #              'OIII-4363', 'Hg', 'Hd', 'NeIII', 'OII', 'MgII']

    t_lines = utils.load_templates(fwhm=fwhm, line_complexes=False,
                                       full_line_list=line_list,
                                       fsps_templates=fsps_templates)

    if return_lists:
        return t_complexes, t_lines
    else:
        # Save them to a file
        np.save('templates_{0}.npy'.format(fwhm), [t_complexes, t_lines])
        print('Wrote `templates_{0}.npy`'.format(fwhm))

class StackFitter(GroupFitter):
    def __init__(self, files='gnt_18197.stack.fits', group_name=None, sys_err=0.02, mask_min=0.1, fit_stacks=True, fcontam=1, PAs=None, extensions=None, min_ivar=0.01, overlap_threshold=3, verbose=True, eazyp=None, eazy_ix=0, MW_EBV=0., chi2_threshold=1.5, min_DoF=200):
        """Object for fitting stacked spectra.

        Parameters
        ----------
        files : str or list of str
            Stack FITS filename.  If a list is supplied, e.g., the product
            of a `~glob` command, then append all specified files.

        group_name : str
            Rootname to associate with the object.  If none, then default to
            `files`.

        sys_err : float
            Minimum systematic error, interpreted as a fractional error.
            The adjusted variance is taken to be

                >>> var = var0 + (sys_err*flux)**2

        mask_min : float
            Only fit 2D pixels where the flat-flambda model has pixel values
            greater than `mask_min` times the maximum of the model.

        fit_stacks : bool
            Fit the stacks of each grism combined from all available PAs.  If
            False, then fit the PAs individually.

        fcontam : float
            Parameter to control weighting of contaminated pixels for
            `fit_stacks=False`.

        """
        if isinstance(files, list):
            file=files[0]
        else:
            file=files

        self.files = [file]
        if group_name is not None:
            self.group_name = group_name
        else:
            self.group_name = file

        if verbose:
            print('Load file {0}'.format(file))

        self.file = file
        self.hdulist = pyfits.open(file)
        self.min_ivar = min_ivar
        self.sys_err = sys_err
        self.fcontam = fcontam

        self.MW_EBV = MW_EBV

        self.h0 = self.hdulist[0].header.copy()
        #self.Ngrism = self.h0['NGRISM']
        self.grisms = []
        self.ext = []
        for i in range(self.h0['NGRISM']):
            g = self.h0['GRISM{0:03d}'.format(i+1)]
            self.grisms.append(g)
            if fit_stacks:
                if extensions is not None:
                    if g not in extensions:
                        continue

                self.ext.append(g)
            else:
                ng = self.h0['N{0}'.format(g)]
                for j in range(ng):
                    pa = self.h0['{0}{1:02d}'.format(g, j+1)]

                    if PAs is not None:
                        if pa not in PAs:
                            continue

                    ext = '{0},{1}'.format(g,pa)
                    if extensions is not None:
                        if ext not in extensions:
                            continue

                    self.ext.append(ext)

        self.N = len(self.ext)
        self.beams = []
        pop = []
        for i in range(self.N):
            E_i = StackedSpectrum(file=self.file, sys_err=sys_err,
                                  mask_min=mask_min, extver=self.ext[i],
                                  mask_threshold=-1, fcontam=fcontam,
                                  min_ivar=min_ivar, MW_EBV=MW_EBV)
            E_i.compute_model()

            if np.isfinite(E_i.kernel.sum()) & (E_i.DoF >= min_DoF):
                self.beams.append(E_i)
            else:
                pop.append(i)

        for i in pop[::-1]:
            self.N -= 1
            p = self.ext.pop(i)

        # Get some parameters from the beams
        self.id = self.h0['ID']
        self.ra = self.h0['RA']
        self.dec = self.h0['DEC']

        self.Asave = {}

        ## Photometry
        self.is_spec = 1
        self.Nphot = 0

        ## Parse the beam data
        self._parse_beams_list()

        if not fit_stacks:
            # self.mask_drizzle_overlaps(threshold=overlap_threshold,
            #                           verbose=verbose)
            if chi2_threshold > 0:
                orig_ext = [e for e in self.ext]
                fit_log, keep_dict, has_bad = self.check_for_bad_PAs(poly_order=3, chi2_threshold=chi2_threshold, fit_background=True, reinit=True, verbose=False)
                if has_bad & verbose:
                    print('Found bad PA.  New list: {0}'.format(keep_dict))

        if verbose:
            print('  {0}'.format(' '.join(self.ext)))

        # Read multiple
        if isinstance(files, list):
            if len(files) > 1:
                for file in files[1:]:
                    extra = StackFitter(files=file, sys_err=sys_err, mask_min=mask_min, fit_stacks=fit_stacks, fcontam=fcontam, pas=pas, extensions=extensions, min_ivar=min_ivar, overlap_threshold=overlap_threshold, eazyp=eazyp, eazy_ix=eazy_ix, chi2_threshold=chi2_threshold, verbose=verbose)
                    self.extend(extra)

        self.idf = np.hstack([b.scif*0+ib for ib, b in enumerate(self.beams)])
        self.idf = np.cast[int](self.idf)

        # if eazyp is not None:
        #     self.eazyp = eazyp
        #
        #     # TBD: do matching to eazyp.cat directly?
        #     self.eazy_ix = eazy_ix
        #
        #     ok_phot = (eazyp.efnu[eazy_ix,:] > 0) & (eazyp.fnu[eazy_ix,:] > eazyp.param['NOT_OBS_THRESHOLD']) & np.isfinite(eazyp.fnu[eazy_ix,:]) & np.isfinite(eazyp.efnu[eazy_ix,:])
        #     ok_phot = np.squeeze(ok_phot)
        #     self.ok_phot = ok_phot
        #
        #     self.Nphot = ok_phot.sum()
        #     if self.Nphot > 0:
        #
        #         # F-lambda photometry, 1e-19 erg/s/cm2/A
        #         self.photom_eflam = (eazyp.efnu[eazy_ix,:]*eazyp.to_flam*eazyp.zp*eazyp.ext_corr/100.)[ok_phot]
        #         self.photom_flam = (eazyp.fnu[eazy_ix,:]*eazyp.to_flam*eazyp.zp*eazyp.ext_corr/100.)[ok_phot]
        #         self.photom_lc = eazyp.lc[ok_phot]
        #
        #         self.scif = np.hstack((self.scif, self.photom_flam))
        #         self.ivarf = np.hstack((self.ivarf, 1/self.photom_eflam**2))
        #         self.sivarf = np.hstack((self.sivarf, 1/self.photom_eflam))
        #         self.wavef = np.hstack((self.wavef, self.photom_lc))
        #
        #         self.weightf = np.hstack((self.weightf, np.ones(self.Nphot)))
        #         self.fit_mask = np.hstack((self.fit_mask, np.ones(self.Nphot, dtype=bool)))
        #         self.DoF += self.Nphot
        #         self.phot_scale = np.array([10.])

    def _parse_beams_list(self):
        """
        """
        # Parse from self.beams list
        self.N = len(self.beams)
        self.ext = [E.extver for E in self.beams]

        self.Ngrism = OrderedDict()
        for beam in self.beams:
            if beam.grism in self.Ngrism:
                self.Ngrism[beam.grism] += 1
            else:
                self.Ngrism[beam.grism] = 1

        # Make "PA" attribute
        self.PA = OrderedDict()
        for g in self.Ngrism:
            self.PA[g] = OrderedDict()

        for i in range(self.N):
            grism = self.ext[i].split(',')[0]
            if ',' in self.ext[i]:
                PA = float(self.ext[i].split(',')[1])
            else:
                PA = 0

            if PA in self.PA[grism]:
                self.PA[grism][PA].append(i)
            else:
                self.PA[grism][PA] = [i]

        self.grisms = list(self.PA.keys())

        self.Ntot = np.sum([E.size for E in self.beams])
        self.scif = np.hstack([E.scif for E in self.beams])
        self.ivarf = np.hstack([E.ivarf for E in self.beams])
        self.wavef = np.hstack([E.wavef for E in self.beams])

        self.weightf = np.hstack([E.weightf for E in self.beams])
        #self.ivarf *= self.weightf

        self.sivarf = np.sqrt(self.ivarf)

        self.fit_mask = np.hstack([E.fit_mask for E in self.beams])
        self.fit_mask &= self.ivarf > self.min_ivar*self.ivarf.max()

        # Dummy parameter.  Implemented for MultiBeam
        self.sens_mask = 1.

        self.DoF = int((self.fit_mask*self.weightf).sum())
        #self.Nmask = self.fit_mask.sum()
        self.Nmask = np.sum([E.fit_mask.sum() for E in self.beams])

        self.slices = self._get_slices(masked=False)
        self.A_bg = self._init_background(masked=False)
        self.Asave = {}

        self._update_beam_mask()
        self.A_bgm = self._init_background(masked=True)

        self.flat_flam = np.hstack([E.flat_flam for E in self.beams])

        self.initialize_masked_arrays()

    def extend(self, st):
        """
        Append second StackFitter objects to `self`.
        """
        self.beams.extend(st.beams)
        self._parse_beams_list()

        # self.grisms.extend(st.grisms)
        # self.grisms = list(np.unique(self.grisms))
        #
        # self.Ngrism = {}
        # for grism in self.grisms:
        #     self.Ngrism[grism] = 0
        #
        # for beam in self.beams:
        #     self.Ngrism[beam.grism] += 1
        #
        # #self.Ngrism = len(self.grisms)
        #
        # self.N += st.N
        # self.ext.extend(st.ext)
        # self.files.extend(st.files)
        #
        # # Re-init
        # self.Ntot = np.sum([E.size for E in self.beams])
        # self.scif = np.hstack([E.scif for E in self.beams])
        # self.ivarf = np.hstack([E.ivarf for E in self.beams])
        # self.wavef = np.hstack([E.wavef for E in self.beams])
        #
        # self.weightf = np.hstack([E.weightf for E in self.beams])
        # #self.ivarf *= self.weightf
        #
        # self.sivarf = np.sqrt(self.ivarf)
        #
        # self.fit_mask = np.hstack([E.fit_mask for E in self.beams])
        # self.fit_mask &= self.ivarf > self.min_ivar*self.ivarf.max()
        #
        # self.slices = self._get_slices(masked=False)
        # self.A_bg = self._init_background(masked=False)
        #
        # self._update_beam_mask()
        # self.A_bgm = self._init_background(masked=True)
        #
        # self.Nmask = self.fit_mask.sum()
        # self.DoF = int((self.fit_mask*self.weightf).sum())
        #
        # self.flat_flam = np.hstack([E.flat_flam for E in self.beams])

    def check_for_bad_PAs(self, poly_order=1, chi2_threshold=1.5, fit_background=True, reinit=True, verbose=False):
        """
        """

        wave = np.linspace(2000,2.5e4,100)
        poly_templates = utils.polynomial_templates(wave, order=poly_order)

        fit_log = OrderedDict()
        keep_dict = OrderedDict()
        has_bad = False

        keep_beams = []

        for g in self.PA:
            fit_log[g] = OrderedDict()
            keep_dict[g] = []

            for pa in self.PA[g]:
                extensions = [self.ext[i] for i in self.PA[g][pa]]
                mb_i = StackFitter(self.file, fcontam=self.fcontam,
                                   sys_err=self.sys_err,
                                   extensions=extensions, fit_stacks=False,
                                   verbose=verbose, chi2_threshold=-1)

                try:
                    chi2, _, _, _ = mb_i.xfit_at_z(z=0,
                                                   templates=poly_templates,
                                                fit_background=fit_background)
                except:
                    chi2 = 1e30

                if False:
                    p_i = mb_i.template_at_z(z=0, templates=poly_templates, fit_background=fit_background, fitter='lstsq', fwhm=1400, get_uncertainties=2)

                fit_log[g][pa] = {'chi2': chi2, 'DoF': mb_i.DoF,
                                  'chi_nu': chi2/np.maximum(mb_i.DoF, 1)}

            min_chinu = 1e30
            for pa in self.PA[g]:
                min_chinu = np.minimum(min_chinu, fit_log[g][pa]['chi_nu'])

            fit_log[g]['min_chinu'] = min_chinu

            for pa in self.PA[g]:
                fit_log[g][pa]['chinu_ratio'] = fit_log[g][pa]['chi_nu']/min_chinu

                if fit_log[g][pa]['chinu_ratio'] < chi2_threshold:
                    keep_dict[g].append(pa)
                    keep_beams.extend([self.beams[i] for i in self.PA[g][pa]])
                else:
                    has_bad = True

        if reinit:
            self.beams = keep_beams
            self._parse_beams_list()
            #self._parse_beams(psf=self.psf_param_dict is not None)

        return fit_log, keep_dict, has_bad

    def compute_model(self, spectrum_1d=None):
        """
        TBD
        """
        return False

    def get_sky_coords(self):
        """Get WCS coordinates of the center of the direct image

        Returns
        -------
        ra, dec : float
            Center coordinates of the beam thumbnail in decimal degrees
        """
        ra = self.beams[0].h0['RA']
        dec = self.beams[0].h0['DEC']
        return ra, dec

    def fit_combined_at_z(self, z=0, fitter='nnls', get_uncertainties=False, eazyp=None, ix=0, order=1, scale_fit=None, method='BFGS'):
        """Fit the 2D spectra with a set of templates at a specified redshift.
        TBD
        Parameters
        ----------
        z : float
            Redshift.

        templates : list
            List of templates to fit.

        fitter : str
            Minimization algorithm to compute template coefficients.
            The default 'nnls' uses non-negative least squares.
            The other option is standard 'leastsq'.

        get_uncertainties : bool
            Compute coefficient uncertainties from the covariance matrix

        Returns
        -------
        chi2 : float
            Chi-squared of the fit

        background : `~np.ndarray`
            Background model

        full_model : `~np.ndarray`
            Best-fit 2D model.

        coeffs, err : `~np.ndarray`
            Template coefficients and uncertainties.

        """
        import scipy.optimize

        # Photometry
        ok_phot = (eazyp.efnu[ix,:] > 0) & (eazyp.fnu[ix,:] > eazyp.param['NOT_OBS_THRESHOLD']) & np.isfinite(eazyp.fnu[ix,:]) & np.isfinite(eazyp.efnu[ix,:])
        ok_phot = np.squeeze(ok_phot)
        self.ok_phot = ok_phot

        Nphot = ok_phot.sum()
        photom_eflam = (eazyp.efnu[ix,:]*eazyp.to_flam*eazyp.ext_corr/100.)[ok_phot]
        photom_flam = (eazyp.fnu[ix,:]*eazyp.to_flam*eazyp.ext_corr/100.)[ok_phot]

        templates = eazyp.templates

        NTEMP = len(templates)
        A = np.zeros((self.N+NTEMP, self.Ntot+Nphot))
        A[:self.N,:-Nphot] += self.A_bg

        pedestal = 0.04

        sivarf = np.hstack((self.sivarf, 1/photom_eflam))
        dataf = np.hstack((self.scif+pedestal, photom_flam))
        fit_mask = np.hstack((self.fit_mask, np.ones(Nphot, dtype=bool)))
        #sivarf *= fit_mask

        # Photometry
        Aphot = (eazyp.tempfilt(z)*3.e18/eazyp.lc**2*(1+z))[:,ok_phot]
        A[self.N:,-Nphot:] += Aphot

        for i, ti in enumerate(templates):
            #ti = templates[t]
            s = [ti.wave*(1+z), ti.flux/(1+z)]

            for j, E in enumerate(self.beams):
                clip = E.ivar.sum(axis=0) > 0
                if (s[0][0] > E.wave[clip].max()) | (s[0][-1] < E.wave[clip].min()):
                    continue

                sl = self.slices[j]
                A[self.N+i, sl] = E.compute_model(spectrum_1d=s)

        oktemp = (A*fit_mask).sum(axis=1) != 0

        # ATA = np.dot(A[oktemp,:]*(sivarf*fit_mask)**2, A[oktemp,:].T)
        # ATy = np.dot(A[oktemp,:]*(sivarf*fit_mask)**2, dataf)

        Ax = A[oktemp,:]*sivarf
        #AxT = Ax[:,fit_mask].T
        #data = (dataf*sivarf)[fit_mask]

        # Run the optimizer
        #method = 'Powell'
        #method = 'BFGS'
        tol = 1.e-4
        init = np.zeros(order+1)
        init[0] = 10.

        if scale_fit is None:
            scale_fit = scipy.optimize.minimize(self.objective_scale, init, args=(Ax, dataf*sivarf, self.wavef, fit_mask, sivarf, Nphot, self.N, 0), method=method, jac=None, hess=None, hessp=None, bounds=(), constraints=(), tol=tol, callback=None, options=None)

            if order == 0:
                scale_fit.x = np.array([np.float(scale_fit.x)])

        coeffs, full, resid, chi2, AxT = self.objective_scale(scale_fit.x, Ax, dataf*sivarf, self.wavef, fit_mask, sivarf, Nphot, self.N, True)

        background = np.dot(coeffs[:self.N], A[:self.N,:])
        full -= background

        background -= pedestal
        background[-Nphot:] = 0

        # Uncertainties from covariance matrix
        if get_uncertainties:
            try:
                covar = np.matrix(np.dot(AxT.T, AxT)).I
                covard = np.sqrt(covar.diagonal()).A.flatten()

                # covarf = np.matrix(np.dot(ATA.T, ATA)).I
                # covardf = np.sqrt(covarf.diagonal()).A.flatten()

            except:
                print('Except!')
                covard = np.zeros(oktemp.sum())#-1.
        else:
            covard = np.zeros(oktemp.sum())#-1.

        full_coeffs = np.zeros(NTEMP)
        full_coeffs[oktemp[self.N:]] = coeffs[self.N:]

        full_coeffs_err = np.zeros(NTEMP)
        full_coeffs_err[oktemp[self.N:]] = covard[self.N:]

        return chi2, background, full, full_coeffs, full_coeffs_err, scale_fit

    @staticmethod
    def scale_AxT(p, Ax, spec_wave, Nphot, Next):
        """
        Scale spectrum templates by polynomial function
        """
        from scipy import polyval

        scale = np.ones(Ax.shape[1])
        scale[:-Nphot] = polyval(p[::-1]/10., (spec_wave-1.e4)/1000.)
        AxT = Ax*scale
        for i in range(Next):
            AxT[i,:] /= scale

        return AxT

    @staticmethod
    def objective_scale(p, Ax, data, spec_wave, fit_mask, sivarf, Nphot, Next, return_coeffs):
        """
        Objective function for fitting for a scale term between photometry and
        spectra
        """
        import scipy.optimize
        from scipy import polyval

        scale = np.ones(Ax.shape[1])
        scale[:-Nphot] = polyval(p[::-1]/10., (spec_wave-1.e4)/1000.)
        AxT = Ax*scale

        # Remove scaling from background component
        for i in range(Next):
            AxT[i,:] /= scale

        #AxT = AxT[:,fit_mask].T
        #(Ax*scale)[:,fit_mask].T
        #AxT[:,:Next] = 1.

        coeffs, rnorm = scipy.optimize.nnls(AxT[:,fit_mask].T, data[fit_mask])

        full = np.dot(coeffs, AxT/sivarf)
        resid = data/sivarf - full# - background
        chi2 = np.sum(resid[fit_mask]**2*sivarf[fit_mask]**2)

        #print('{0} {1}'.format(p, chi2))

        if return_coeffs:
            return coeffs, full, resid, chi2, AxT
        else:
            return chi2

        # Testing
        # method = 'COBYLA'
        # out = scipy.optimize.minimize(objective_scale, [10.], args=(Ax, dataf*sivarf, fit_mask, sivarf, Nphot, 0), method=method, jac=None, hess=None, hessp=None, bounds=None, constraints=(), tol=None, callback=None, options=None)
        # print(method, out.nfev, out.x)
        # out = scipy.optimize.minimize(objective_scale, [10.], args=(Ax, dataf*sivarf, fit_mask, sivarf, Nphot, 0), method='COBYLA', jac=None, hess=None, hessp=None, bounds=None, constraints=(), tol=None, callback=None, options=None)

    def fit_at_z(self, z=0, templates=[], fitter='nnls', get_uncertainties=False):
        """Fit the 2D spectra with a set of templates at a specified redshift.

        Parameters
        ----------
        z : float
            Redshift.

        templates : list
            List of templates to fit.

        fitter : str
            Minimization algorithm to compute template coefficients.
            The default 'nnls' uses non-negative least squares.
            The other option is standard 'leastsq'.

        get_uncertainties : bool
            Compute coefficient uncertainties from the covariance matrix

        Returns
        -------
        chi2 : float
            Chi-squared of the fit

        background : `~np.ndarray`
            Background model

        full_model : `~np.ndarray`
            Best-fit 2D model.

        coeffs, err : `~np.ndarray`
            Template coefficients and uncertainties.

        """
        import scipy.optimize

        NTEMP = len(templates)
        A = np.zeros((self.N+NTEMP, self.Ntot))
        A[:self.N,:] += self.A_bg

        for i, t in enumerate(templates):
            ti = templates[t]
            try:
                import eazy.igm
                if z > 7:
                    igm = eazy.igm.Inoue14()
                    igmz = igm.full_IGM(z, ti.wave*(1+z))
                else:
                    igmz = 1.

            except:
                igmz = 1.

            s = [ti.wave*(1+z), ti.flux/(1+z)*igmz]

            for j, E in enumerate(self.beams):
                clip = E.ivar.sum(axis=0) > 0
                if (s[0][0] > E.wave[clip].max()) | (s[0][-1] < E.wave[clip].min()):
                    continue

                sl = self.slices[j]
                A[self.N+i, sl] = E.compute_model(spectrum_1d=s)

        oktemp = (A*self.fit_mask).sum(axis=1) != 0

        Ax = A[oktemp,:]*self.sivarf

        pedestal = 0.04

        AxT = Ax[:,self.fit_mask].T
        data = ((self.scif+pedestal)*self.sivarf)[self.fit_mask]

        if fitter == 'nnls':
            coeffs, rnorm = scipy.optimize.nnls(AxT, data)
        else:
            coeffs, residuals, rank, s = np.linalg.lstsq(AxT, data)

        background = np.dot(coeffs[:self.N], A[:self.N,:]) - pedestal
        full = np.dot(coeffs[self.N:], Ax[self.N:,:]/self.sivarf)

        resid = self.scif - full - background
        chi2 = np.sum(resid[self.fit_mask]**2*self.sivarf[self.fit_mask]**2)

        # Uncertainties from covariance matrix
        if get_uncertainties:
            try:
                covar = np.matrix(np.dot(AxT.T, AxT)).I
                covard = np.sqrt(covar.diagonal()).A.flatten()
            except:
                print('Except: covar!')
                covard = np.zeros(oktemp.sum())#-1.
        else:
            covard = np.zeros(oktemp.sum())#-1.

        full_coeffs = np.zeros(NTEMP)
        full_coeffs[oktemp[self.N:]] = coeffs[self.N:]

        full_coeffs_err = np.zeros(NTEMP)
        full_coeffs_err[oktemp[self.N:]] = covard[self.N:]

        return chi2, background, full, full_coeffs, full_coeffs_err

    def fit_zgrid(self, dz0=0.005, zr=[0.4, 3.4], fitter='nnls', make_plot=True, save_data=True, prior=None, templates_file='templates.npy', verbose=True, outlier_threshold=1e30, eazyp=None, ix=0, order=0, scale_fit=None):
        """Fit templates on a redshift grid.

        Parameters
        ----------
        dz0 : float
            Initial step size of the redshift grid (dz/1+z).

        zr : list
            Redshift range to consider.

        fitter : str
            Minimization algorithm.  Default is non-negative least-squares.

        make_plot : bool
            Make the diagnostic plot.

        prior : list
            Naive prior to add to the nominal chi-squared(z) of the template
            fits.  The example below is a simple Gaussian prior centered
            at z=1.5.

                >>> z_prior = np.arange(0,3,0.001)
                >>> chi_prior = (z_prior-1.5)**2/2/0.1**2
                >>> prior = [z_prior, chi_prior]

        templates_file : str
            Filename of the `~numpy` save file containing the templates.  Use
            the `make_templates` script to generate this.

        verbose : bool
            Print the redshift grid steps.

        Returns
        -------
        hdu : `~astropy.io.fits.HDUList`
            Multi-extension FITS file with the result of the redshift fits.

        """
        import os
        #import grizli

        import matplotlib.gridspec
        import matplotlib.pyplot as plt
        import numpy as np

        from . import utils

        t_complex, t_i = np.load(templates_file, allow_pickle=True)

        z = utils.log_zgrid(zr=zr, dz=dz0)
        chi2 = z*0.
        for i in range(len(z)):
            if eazyp:
                out = self.fit_combined_at_z(z=z[i], eazyp=eazyp, ix=ix, order=order, scale_fit=scale_fit)
                chi2[i], bg, full, coeffs, err, scale_fit = out
            else:
                out = self.fit_at_z(z=z[i], templates=t_complex)
                chi2[i], bg, full, coeffs, err = out

            if verbose:
                print('{0:.4f} - {1:10.1f}'.format(z[i], chi2[i]))

        # Zoom in on the chi-sq minimum.
        ci = chi2
        zi = z
        for iter in range(1,7):
            if prior is not None:
                pz = np.interp(zi, prior[0], prior[1])
                cp = ci+pz
            else:
                cp = ci

            iz = np.argmin(cp)
            z0 = zi[iz]
            dz = dz0/2.02**iter
            zi = utils.log_zgrid(zr=[z0-dz*4, z0+dz*4], dz=dz)
            ci = zi*0.
            for i in range(len(zi)):

                if eazyp:
                    out = self.fit_combined_at_z(z=zi[i], eazyp=eazyp, ix=ix, order=order, scale_fit=scale_fit)
                    ci[i], bg, full, coeffs, err, scale_fit = out
                else:
                    out = self.fit_at_z(z=zi[i], templates=t_complex, fitter=fitter)
                    ci[i], bg, full, coeffs, err = out

                # out = self.fit_at_z(z=zi[i], templates=t_complex,
                #                     fitter=fitter)
                #
                # ci[i], bg, full, coeffs, err = out

                if verbose:
                    print('{0:.4f} - {1:10.1f}'.format(zi[i], ci[i]))

            z = np.append(z, zi)
            chi2 = np.append(chi2, ci)

        so = np.argsort(z)
        z = z[so]
        chi2 = chi2[so]

        # Apply the prior
        if prior is not None:
            pz = np.interp(z, prior[0], prior[1])
            chi2 += pz

        # Get the fit with the individual line templates at the best redshift
        chi2x, bgz, fullz, coeffs, err = self.fit_at_z(z=z[np.argmin(chi2)], templates=t_i, fitter=fitter, get_uncertainties=True)

        # Mask outliers
        if outlier_threshold > 0:
            resid = self.scif - fullz - bgz
            outlier_mask = (resid*self.sivarf < outlier_threshold)
            #outlier_mask &= self.fit_mask
            #self.sivarf[outlier_mask] = 1/resid[outlier_mask]
            #print('Mask {0} pixels with resid > {1} sigma'.format((outlier_mask).sum(), outlier_threshold))

            print('Mask {0} pixels with resid > {1} sigma'.format((~outlier_mask & self.fit_mask).sum(), outlier_threshold))
            self.fit_mask &= outlier_mask
            #self.DoF = self.fit_mask.sum() #(self.ivar > 0).sum()
            self.DoF = int((self.fit_mask*self.weightf).sum())

        # Table with z, chi-squared
        t = utils.GTable()
        t['z'] = z
        t['chi2'] = chi2

        if prior is not None:
            t['prior'] = pz

        # "area" parameter for redshift quality.
        num = np.trapz(np.clip(chi2-chi2.min(), 0, 25), z)
        denom = np.trapz(z*0+25, z)
        area25 = 1-num/denom

        # "best" redshift
        zbest = z[np.argmin(chi2)]

        # Metadata will be stored as header keywords in the FITS table
        t.meta = OrderedDict()
        t.meta['ID'] = (self.h0['ID'], 'Object ID')
        t.meta['RA'] = (self.h0['RA'], 'Right Ascension')
        t.meta['DEC'] = (self.h0['DEC'], 'Declination')
        t.meta['Z'] = (zbest, 'Best-fit redshift')
        t.meta['CHIMIN'] = (chi2.min(), 'Min Chi2')
        t.meta['CHIMAX'] = (chi2.max(), 'Max Chi2')
        t.meta['DOF'] = (self.DoF, 'Degrees of freedom')
        t.meta['AREA25'] = (area25, 'Area under CHIMIN+25')
        t.meta['FITTER'] = (fitter, 'Minimization algorithm')
        t.meta['HASPRIOR'] = (prior is not None, 'Was prior specified?')

        # Best-fit templates
        tc, tl = self.generate_1D_templates(coeffs,
                                            templates_file=templates_file)
        # for i, te in enumerate(t_i):
        #     if i == 0:
        #         tc = t_i[te].zscale(0, scalar=coeffs[i])
        #         tl = t_i[te].zscale(0, scalar=coeffs[i])
        #     else:
        #         if te.startswith('line'):
        #             tc += t_i[te].zscale(0, scalar=0.)
        #         else:
        #             tc += t_i[te].zscale(0, scalar=coeffs[i])
        #
        #         tl += t_i[te].zscale(0, scalar=coeffs[i])

        # Get line fluxes, uncertainties and EWs
        il = 0
        for i, te in enumerate(t_i):
            if te.startswith('line'):
                il+=1

                if coeffs[i] == 0:
                    EW = 0.
                else:
                    tn = (t_i[te].zscale(0, scalar=coeffs[i]) +
                              tc.zscale(0, scalar=1))

                    td = (t_i[te].zscale(0, scalar=0) +
                             tc.zscale(0, scalar=1))

                    clip = (td.wave <= t_i[te].wave.max())
                    clip &= (td.wave >= t_i[te].wave.min())

                    EW = np.trapz((tn.flux/td.flux)[clip]-1, td.wave[clip])
                    if not np.isfinite(EW):
                        EW = -1000.

                t.meta['LINE{0:03d}F'.format(il)] = (coeffs[i],
                                       '{0} line flux'.format(te[5:]))

                t.meta['LINE{0:03d}E'.format(il)] = (err[i],
                            '{0} line flux uncertainty'.format(te[5:]))

                #print('xxx EW', EW)
                t.meta['LINE{0:03d}W'.format(il)] = (EW,
                            '{0} line rest EQW'.format(te[5:]))

        tfile = self.file.replace('.fits', '.zfit.fits')
        if os.path.exists(tfile):
            os.remove(tfile)

        t.write(tfile)

        ### Add image HDU and templates
        hdu = pyfits.open(tfile, mode='update')
        hdu[1].header['EXTNAME'] = 'ZFIT'

        # oned_templates = np.array([tc.wave*(1+zbest), tc.flux/(1+zbest),
        #                            tl.flux/(1+zbest)])
        header = pyfits.Header()
        header['TEMPFILE'] = (templates_file, 'File with stored templates')
        hdu.append(pyfits.ImageHDU(data=coeffs, name='COEFFS'))

        for i in range(self.N):
            E = self.beams[i]
            model_i = fullz[self.slices[i]].reshape(E.sh)
            bg_i = bgz[self.slices[i]].reshape(E.sh)

            model_i[~np.isfinite(model_i)] = 0
            bg_i[~np.isfinite(bg_i)] = 0

            hdu.append(pyfits.ImageHDU(data=model_i, header=E.header,
                                       name='MODEL'))

            hdu.append(pyfits.ImageHDU(data=bg_i, header=E.header,
                                       name='BACKGROUND'))

        hdu.flush()

        if make_plot:
            self.make_fit_plot(hdu)

        return hdu

    @classmethod
    def generate_1D_templates(self, coeffs, templates_file='templates.npy'):
        """
        TBD
        """
        t_complex, t_i = np.load(templates_file, allow_pickle=True)

        # Best-fit templates
        for i, te in enumerate(t_i):
            if i == 0:
                tc = t_i[te].zscale(0, scalar=coeffs[i])
                tl = t_i[te].zscale(0, scalar=coeffs[i])
            else:
                if te.startswith('line'):
                    tc += t_i[te].zscale(0, scalar=0.)
                else:
                    tc += t_i[te].zscale(0, scalar=coeffs[i])

                tl += t_i[te].zscale(0, scalar=coeffs[i])

        return tc, tl

    def mask_drizzle_overlaps(self, threshold=3, verbose=True):
        """
        TBD

        mask pixels in beams of a given grism that are greater than
        `threshold` sigma times the minimum of all beams of that
        grism
        """
        min_grism = {}
        for grism in self.grisms:
            for E in self.beams:
                if not E.extver.startswith(grism):
                    continue

                if grism not in min_grism:
                    min_grism[grism] = E.scif*E.fit_mask
                else:
                    empty = (min_grism[grism] == 0) & E.fit_mask
                    min_grism[grism][E.fit_mask] = np.minimum(min_grism[grism][E.fit_mask], E.scif[E.fit_mask])
                    min_grism[grism][empty] = E.scif[empty]

        for grism in self.grisms:
            for E in self.beams:
                if not E.extver.startswith(grism):
                    continue

                new_mask = (E.scif - min_grism[grism]) < threshold/E.sivarf
                if verbose:
                    print('Mask {0} additional pixels for ext {1}'.format((~new_mask & E.fit_mask).sum(), E.extver))

                E.fit_mask &= new_mask

    def make_fit_plot(self, hdu, scale_pz=True, show_2d=False):
        """Make a plot showing the fit

        Parameters
        ----------
        hdu : `~astropy.io.fits.HDUList`
            Fit results from `fit_zgrid`.

        Returns
        -------
        fig : `~matplotlib.figure.Figure`
            The figure object.

        """
        import matplotlib.pyplot as plt
        import matplotlib.gridspec
        from matplotlib.ticker import MultipleLocator

        #import grizli
        from . import utils

        zfit = utils.GTable.read(hdu[1])
        z = zfit['z']
        chi2 = zfit['chi2']

        # Initialize plot window
        Ng = len(self.grisms)
        if show_2d:
            height_ratios = [0.25]*self.N
            height_ratios.append(1)
            gs = matplotlib.gridspec.GridSpec(self.N+1,2,
                                width_ratios=[1,1.5+0.5*(Ng>1)],
                                height_ratios=height_ratios, hspace=0.)

            fig = plt.figure(figsize=[8+4*(Ng>1), 3.5+0.5*self.N])
        else:
            gs = matplotlib.gridspec.GridSpec(1,2,
                            width_ratios=[1,1.5+0.5*(Ng>1)],
                            hspace=0.)

            fig = plt.figure(figsize=[8+4*(Ng>1), 3.5])

        # Chi-squared
        axz = fig.add_subplot(gs[-1,0]) #121)

        axz.text(0.95, 0.04, self.file + '\n'+'z={0:.4f}'.format(z[np.argmin(chi2)]), ha='right', va='bottom', transform=axz.transAxes, fontsize=9)

        # Scale p(z) to DoF=1.
        if scale_pz:
            scale_nu = chi2.min()/self.DoF
            scl_label = '_s'
        else:
            scale_nu = 1.
            scl_label = ''

        axz.plot(z, (chi2-chi2.min())/scale_nu, color='k')
        axz.fill_between(z, (chi2-chi2.min())/scale_nu, 27, color='k', alpha=0.5)
        axz.set_ylim(-4,27)
        axz.set_xlabel(r'$z$')
        axz.set_ylabel(r'$\Delta\chi^2{2}$ ({0:.0f}/$\nu$={1:d})'.format(chi2.min(), self.DoF, scl_label))
        axz.set_yticks([1,4,9,16,25])

        axz.set_xlim(z.min(), z.max())
        axz.grid()

        # 2D spectra
        if show_2d:
            twod_axes = []
            for i in range(self.N):
                ax_i = fig.add_subplot(gs[i,1])

                model = hdu['MODEL', self.ext[i]].data
                ymax = model[np.isfinite(model)].max()
                #print('MAX', ymax)

                cmap = 'viridis_r'
                cmap = 'cubehelix_r'

                clean = self.beams[i].sci - hdu['BACKGROUND', self.ext[i]].data
                clean *= self.beams[i].fit_mask.reshape(self.beams[i].sh)

                w = self.beams[i].wave/1.e4

                ax_i.imshow(clean, vmin=-0.02*ymax, vmax=1.1*ymax, origin='lower',
                            extent=[w[0], w[-1], 0., 1.], aspect='auto',
                            cmap=cmap)

                ax_i.text(0.04, 0.92, self.ext[i], ha='left', va='top', transform=ax_i.transAxes, fontsize=8)

                ax_i.set_xticklabels([]); ax_i.set_yticklabels([])
                twod_axes.append(ax_i)

        axc = fig.add_subplot(gs[-1,1]) #224)

        # 1D Spectra + fit
        ymin = 1.e30
        ymax = -1.e30
        wmin = 1.e30
        wmax = -1.e30

        for i in range(self.N):

            E = self.beams[i]

            clean = E.sci - hdu['BACKGROUND', self.ext[i]].data
            w, fl, er = E.optimal_extract(clean)
            w, flm, erm = E.optimal_extract(hdu['MODEL', self.ext[i]].data)
            w = w/1.e4

            # Do we need to convert to F-lambda units?
            if E.is_flambda:
                unit_corr = 1.
                clip = (er > 0) & np.isfinite(er) & np.isfinite(flm)
                clip[:10] = False
                clip[-10:] = False

                if clip.sum() == 0:
                    clip = (er > -1)

            else:
                unit_corr = 1./E.sens
                clip = (E.sens > 0.1*E.sens.max())
                clip &= (np.isfinite(flm)) & (er > 0)

            if clip.sum() == 0:
                continue

            # fl *= 100*unit_corr
            # er *= 100*unit_corr
            # flm *= 100*unit_corr

            fl *= unit_corr/1.e-19
            er *= unit_corr/1.e-19
            flm *= unit_corr/1.e-19

            f_alpha = 1./self.Ngrism[E.grism]**0.5 #self.h0['N{0}'.format(E.header['GRISM'])]**0.5

            axc.errorbar(w[clip], fl[clip], er[clip], color=GRISM_COLORS[E.grism], alpha=0.3*f_alpha, marker='.', linestyle='None')
            #axc.fill_between(w[clip], (fl+er)[clip], (fl-er)[clip], color='k', alpha=0.2)
            axc.plot(w[clip], flm[clip], color='r', alpha=0.6*f_alpha, linewidth=2)

            # Plot limits
            # ymax = np.maximum(ymax, (flm+np.median(er))[clip].max())
            # ymin = np.minimum(ymin, (flm-er*0.)[clip].min())

            ymax = np.maximum(ymax,
                              np.percentile((flm+np.median(er))[clip], 98))

            ymin = np.minimum(ymin, np.percentile((flm-er*0.)[clip], 2))

            wmax = np.maximum(wmax, w.max())
            wmin = np.minimum(wmin, w.min())

        axc.set_xlim(wmin, wmax)
        axc.semilogx(subsx=[wmax])
        #axc.set_xticklabels([])
        axc.set_xlabel(r'$\lambda$')
        axc.set_ylabel(r'$f_\lambda \times 10^{-19}$')
        #axc.xaxis.set_major_locator(MultipleLocator(0.1))

        axc.set_ylim(ymin-0.2*ymax, 1.2*ymax)
        axc.grid()

        for ax in [axc]: #[axa, axb, axc]:

            labels = np.arange(np.ceil(wmin*10), np.ceil(wmax*10))/10.
            ax.set_xticks(labels)
            ax.set_xticklabels(labels)
            #ax.set_xticklabels([])
            #print(labels, wmin, wmax)

        if show_2d:
            for ax in twod_axes:
                ax.set_xlim(wmin, wmax)

        gs.tight_layout(fig, pad=0.1, w_pad=0.1)
        fig.savefig(self.file.replace('.fits', '.zfit.png'))
        return fig

class StackedSpectrum(object):
    def __init__(self, file='gnt_18197.stack.G141.285.fits', sys_err=0.02, mask_min=0.1, extver='G141', mask_threshold=7, fcontam=1., min_ivar=0.001, MW_EBV=0.):
        import os
        #import grizli
        from . import GRIZLI_PATH
        from . import grismconf

        self.sys_err = sys_err
        self.mask_min = mask_min
        self.extver = extver
        self.grism = self.extver.split(',')[0]
        self.mask_threshold=mask_threshold

        self.MW_EBV = MW_EBV
        self.init_galactic_extinction(MW_EBV)

        self.file = file
        self.hdulist = pyfits.open(file)

        self.h0 = self.hdulist[0].header.copy()
        self.header = self.hdulist['SCI',extver].header.copy()
        self.sh = (self.header['NAXIS2'], self.header['NAXIS1'])
        self.wave = self.get_wavelength_from_header(self.header)
        self.wavef = np.dot(np.ones((self.sh[0],1)), self.wave[None,:]).flatten()

        if 'BEAM' in self.header:
            self.beam_name = self.header['BEAM']
        else:
            self.beam_name = 'A'

        # Configuration file
        self.is_flambda = self.header['ISFLAM']
        self.conf_file = self.header['CONF']
        try:
            self.conf = grismconf.aXeConf(self.conf_file)
        except:
            # Try global path
            base = os.path.basename(self.conf_file)
            localfile = os.path.join(GRIZLI_PATH, 'CONF', base)
            self.conf = grismconf.aXeConf(localfile)

        self.conf.get_beams()

        self.sci = self.hdulist['SCI',extver].data*1.
        self.ivar0 = self.hdulist['WHT',extver].data*1
        self.size = self.sci.size
        self.thumbs = {}

        self.scif = self.sci.flatten()
        self.ivarf0 = self.ivar0.flatten()

        self.ivarf = 1/(1/self.ivarf0 + (sys_err*self.scif)**2)
        self.ivar = self.ivarf.reshape(self.sh)

        self.sivarf = np.sqrt(self.ivar).flatten()

        self.fit_mask = (self.ivarf > min_ivar*self.ivarf.max())
        self.fit_mask &= np.isfinite(self.scif) & np.isfinite(self.ivarf)

        # Dummy parameter.  Implemented for MultiBeam
        self.sens_mask = 1.

        # Contamination weighting
        self.fcontam = fcontam
        if ('CONTAM',extver) in self.hdulist:
            self.contam = np.abs(self.hdulist['CONTAM',extver].data*1.)
            self.weight = np.exp(-fcontam*self.contam*np.sqrt(self.ivar0))
            self.contamf = self.contam.flatten()
            self.weightf = self.weight.flatten()
        else:
            self.contam = self.sci*0.
            self.contamf = self.scif*0.

            self.weight = self.sci*0.+1
            self.weightf = self.scif*0.+1

        # Spatial kernel
        self.kernel = self.hdulist['KERNEL',extver].data*1
        self.kernel /= self.kernel.sum()

        self._build_model()

        self.flat = self.compute_model()

        if self.is_flambda:
            elec = self.flat*self.ivarf0/self.ivarf0.max()
        else:
            elec = self.flat

        self.fit_mask &= elec > mask_min*elec.max()

        self.DoF = self.fit_mask.sum() #(self.ivar > 0).sum()

        if mask_threshold > 0:
            self.drizzle_mask(mask_threshold=mask_threshold)

        self.flat_flam = self.compute_model(in_place=False, is_cgs=True)

    def init_galactic_extinction(self, MW_EBV=0., R_V=utils.MW_RV):
        """
        Initialize Fitzpatrick 99 Galactic extinction

        Parameters
        ----------
        MW_EBV : float
            Local E(B-V)

        R_V : float
            Relation between specific and total extinction,
            ``a_v = r_v * ebv``.

        Returns
        -------
        Sets `self.MW_F99` attribute, which is a callable function that
        returns the extinction for a supplied array of wavelengths.

        If MW_EBV <= 0, then sets `self.MW_F99 = None`.

        """
        self.MW_F99 = None
        if MW_EBV > 0:
            self.MW_F99 = utils.MW_F99(MW_EBV*R_V, r_v=R_V)

            # try:
            #     from specutils.extinction import ExtinctionF99
            #     self.MW_F99 = ExtinctionF99(MW_EBV*R_V, r_v=R_V)
            # except(ImportError):
            #     print('Couldn\'t import `specutils.extinction`, MW extinction not implemented')

    @classmethod
    def get_wavelength_from_header(self, h):
        """
        Generate wavelength array from WCS keywords
        """
        w = (np.arange(h['NAXIS1'])+1-h['CRPIX1'])*h['CD1_1'] + h['CRVAL1']

        # Now header keywords scaled to microns
        if w.max() < 3:
            w *= 1.e4

        return w

    def init_optimal_profile(self):
        """Initilize optimal extraction profile
        """
        m = self.compute_model(in_place=False)
        m = m.reshape(self.sh)
        m[m < 0] = 0
        self.optimal_profile = m/m.sum(axis=0)

    def optimal_extract(self, data, bin=0, ivar=None, weight=None, loglam=False):
        """
        Optimally-weighted 1D extraction

        xx Dummy parameters ivar, weight, loglam
        """
        import scipy.ndimage as nd

        # flatf = self.flat.reshape(self.sh).sum(axis=0)
        # prof = self.flat.reshape(self.sh)/flatf
        #
        if not hasattr(self, 'optimal_profile'):
            self.init_optimal_profile()

        prof = self.optimal_profile

        if ivar is None:
            ivar = self.ivar

        num = prof*data*ivar*self.weight
        den = prof**2*ivar*self.weight
        opt_flux = num.sum(axis=0)/den.sum(axis=0)
        opt_var = 1./den.sum(axis=0)

        opt_rms = np.sqrt(opt_var)
        clip = (opt_var == 0) | ~np.isfinite(opt_var)
        opt_rms[clip] = 0
        opt_flux[clip] = 0

        if bin > 1:
            kern = np.ones(bin, dtype=float)/bin
            opt_flux = nd.convolve(opt_flux, kern)[bin // 2::bin]
            opt_var = nd.convolve(opt_var, kern**2)[bin // 2::bin]
            opt_rms = np.sqrt(opt_var)
            wave = self.wave[bin // 2::bin]
        else:
            wave = self.wave

        return wave, opt_flux, opt_rms

    def _build_model(self):
        """
        Initiazize components for generating 2D model
        """
        from .utils_c.interp import interp_conserve_c

        NY = self.sh[0]
        data = np.zeros((self.header['NAXIS1'], self.header['NAXIS2'], self.header['NAXIS1']))

        for j in range(NY//2):
            data[j,:,:j+NY//2] += self.kernel[:, -NY//2-j:]

        for j in range(self.sh[1]-NY//2, self.sh[1]):
            data[j,:,-NY//2+j:] += self.kernel[:, :self.sh[1]-j+NY//2]

        for j in range(NY//2, self.sh[1]-NY//2):
            #print(j)
            data[j,:,j-NY//2:j+NY//2] += self.kernel

        self.fit_data = data.reshape(self.sh[1],-1)

        if not self.is_flambda:
            conf_sens = self.conf.sens[self.beam_name]
            if self.MW_F99 is not None:
                MWext = 10**(-0.4*(self.MW_F99(conf_sens['WAVELENGTH']*u.AA)))
            else:
                MWext = 1.

            sens = interp_conserve_c(self.wave, conf_sens['WAVELENGTH'],
                                     conf_sens['SENSITIVITY']*MWext)

            if 'DLAM0' in self.header:
                #print('xx Header')
                dlam = self.header['DLAM0']
            else:
                dlam = np.median(np.diff(self.wave))

            #dlam = np.median(np.diff(self.wave))
            dlam = np.gradient(self.wave)

            # G800 variable dispersion, which isn't coming out right
            # in the 2D drizzled spectrum
            if self.grism == 'G800L':
                ## Linear fit to differential G800L dispersion
                # y = np.diff(beam.wave)/np.diff(beam.wave)[0]
                # x = beam.wave[1:]
                # c_i = np.polyfit(x/1.e4, y, 1)
                c_i = np.array([ 0.07498747,  0.98928126])
                scale = np.polyval(c_i, self.wave/1.e4)
                sens *= scale

            self.sens = sens*dlam #*1.e-17
            self.fit_data = (self.fit_data.T*self.sens).T

    def compute_model(self, spectrum_1d=None, is_cgs=None, in_place=False):
        """
        Generate the model spectrum

        xxx is_cgs and in_place are dummy parameters to match `MultiBeam.compute_model`.
        """
        from .utils_c.interp import interp_conserve_c

        if spectrum_1d is None:
            fl = self.wave*0+1
        else:
            fl = interp_conserve_c(self.wave, spectrum_1d[0], spectrum_1d[1])

        model = np.dot(fl, self.fit_data)#.reshape(self.sh)
        #self.model = model
        return model

import numpy,import scipy,from scipy,import astropy


"""

Framework for GO analysis in glbase

"""

import sys, os, math
import numpy
import matplotlib.pyplot as plot
from matplotlib.lines import Line2D
import matplotlib.lines as mpllines

from . import config
from .genelist import genelist
from .location import location
from .draw import draw

# Available formats
# go_DAVID

# Bind the GO data into glbase?
# If so make the data available here:

class glgo(genelist):
    def __init__(self, filename=None,  **kargs):
        assert "loadable_list" not in kargs, "glgo() does not support loadable_lists"
        assert "format" in kargs and kargs["format"], "glgo requries a 'format' argument"
        assert "qvalue" in kargs["format"] or 'pvalue' in kargs["format"], "the format must contain a 'qvalue' or 'pvalue' key"

        genelist.__init__(self, filename=filename, **kargs)

        # Tidy up '~' symbols in GO terms from typical DAVID/R output
        if 'name' in self.linearData[0]:
            for item in self.linearData:
                if not isinstance(item['name'], location): # weird but true. Occasionally appears as an enriched term
                    item['name'] = item['name'].replace("~", ":")
            self._optimiseData()

    def bind_GO_set(self, filename):
        """
        **Purpose**
            bind a GO set (or list) of data into glgo and make it available for testing

        **Arguments**
            filename
                the filename for the GO set.
        """
        pass

    def barh(self, filename, ontology=None, maximum=None, pad=1e-105, key="qvalue",
        label_fontsize=7, mark_significance=True, tight_layout=True, qvalue_limit=0.1, **kargs):
        """
        **Purpose**
            Draw a barchart for some ontology

        **Arguments**
            filename (Required)
                THe filename to save the image to

            ontology (Optional)
                The name of the ontology class to use. If not specified then use all of the data
                in the list.

            maximum (Optional)
                the maximum number of GO items to plot. If not specified then use all of the items
                in the list

            qvalue_limit (Optional, default=0.1)
                The significance limit for this data

            pad (Optional, default=1e-105)
                By default glgo will -log10 transform the pvalues/qvalues.
                Occasionally some GO discovery programs will output a p/q-value
                of 0 or NA. Pad allows these values to be used. But also sets a maximum
                size of the bar chart.

                Warning:
                If your GO program has higher precision than pad (1E-105) then items with a p/q-value
                of 0 will be set to pad and so may appear in the middle of your list and
                not at the top, as you would expect.

            label_fontsize (Optional, default=&)
                The font size of the labels

            mark_significance (Optional, default=True)
                mark lines of significance at 0.05 and 0.01 p-value/q-value.

            tight_layout (Optional, default=True)
                use tight_layout to maximise space usage, but the location of the y axis
                will shift around to make space for the labels and may not be in the same
                position across multiple graphs.

                If tight_layout is False though then the lables may be truncated on the left hand edge.
                Use the 'position' argument to modify the position of the graph

        **Returns**
            None
        """
        self.sort(key)
        if ontology:
            data = self.get(key="ontology", value=ontology, mode="greedy") # get only for this ontology
        else:
            data = self.linearData # take all data

        bar_width = 0.5
        half_bar_width = bar_width / 2.0

        if not data: # Check that we got some entries
            config.log.warning("No items found for ontology '%s'" % ontology)
            # Don't throw an error just warn and ignore.
            return(None)

        if maximum:
            data = data[0:maximum]
        data.reverse()

        #print data

        vals = [i[key] for i in data] # data ends up unfolded.
        labs = [i["name"] for i in data]

        if qvalue_limit:
            newl = []
            newv = []
            for l, v in zip(labs, vals):
                if v < qvalue_limit:
                    newl.append(l)
                    newv.append(v)
            labs = newl
            vals = [-math.log10(i+pad) for i in newv]
            print(labs, vals)
        else:
            vals = [-math.log10(i[key]+pad) for i in data] # now padded, etc

        if "figsize" not in kargs:
            kargs["figsize"] = (8, 5)

        fig = self.draw.getfigure(**kargs)
        ax = fig.add_subplot(111)

        x = numpy.arange(0, len(vals))
        ax.barh(x, vals, bar_width, ec="none", color="grey")
        ax.set_yticklabels(labs)
        ax.set_yticks(x+half_bar_width)
        ax.set_frame_on(False)
        if ontology:
            ax.set_title(ontology)
        ax.set_ylim([-half_bar_width, len(vals)-1+bar_width+half_bar_width]) # num_items + bar_width
        ax.set_xlim([0, max(vals + [-math.log10(0.05)])])
        if mark_significance:
            ax.axvline(-math.log10(0.05), ls="-", color="black")
            ax.axvline(-math.log10(0.01), ls=":", color="grey")
        [t.set_fontsize(label_fontsize) for t in ax.get_yticklabels()]
        [t.set_fontsize(8) for t in ax.get_xticklabels()]
        ax.tick_params(top=False, bottom=True, left=False, right=False)
        #ax.axes.get_yaxis().set_visible(False)
        for line in ax.get_xticklines():
            line.set_marker(mpllines.TICKDOWN)
        #[item.set_markeredgewidth(0.0) for item in ax.yaxis.get_ticklines()]
        #[item.set_markeredgewidth(0.0) for item in ax.xaxis.get_ticklines()]

        # put back x axis:
        xs = ax.get_xaxis().get_view_interval()
        ymin = ax.get_yaxis().get_view_interval()[0]
        ax.add_artist(Line2D(xs, (ymin, ymin), color='black', linewidth=2))

        if tight_layout:
            fig.tight_layout()
        elif "position" in kargs and kargs["position"] and len(kargs["position"]) == 4:
            ax.set_position(kargs["position"])
        else:
            ax.set_position([0.7, 0.05, 0.25, 0.90]) # My favourite position. Uhm...
        actual_filename = self.draw.savefigure(fig, filename)
        config.log.info("barh(): saved '%s'" % actual_filename)
        return(None)


import numpy as np
from gym.spaces import Box

import pyflex
from softgym.envs.fluid_rigid_env import FluidTorusEnv
import copy
from softgym.utils.misc import quatFromAxisAngle
from softgym.action_space.robot_env import RobotBase
import pickle
import os.path as osp

class TransportTorus1D(FluidTorusEnv):
    def __init__(self, observation_mode, action_mode, config=None, cached_states_path='transport_torus_init_states.pkl', **kwargs):
        '''
        This class implements a transport torus task.
        The torus is put on a box. You need to move the box to a target location.
        This is a 1D task.

        observation_mode: "cam_rgb" or "full_state"
        action_mode: "direct"

        TODO: add more description of the task.
        '''
        assert observation_mode in ['cam_rgb', 'point_cloud', 'key_point']

        self.observation_mode = observation_mode
        self.action_mode = action_mode
        self.wall_num = 5  # number of box walls. floor/left/right/front/back
        self.distance_coef = 1.
        self.torus_penalty_coef = 10.
        self.terminal_x = 1.2
        self.min_x = -0.25
        self.max_x = 1.4
        self.prev_reward = 0

        self.reward_min = self.distance_coef * min(self.min_x - self.terminal_x, self.terminal_x - self.max_x) - \
            self.torus_penalty_coef * 1
        self.reward_max = 0 # reach target position, with no water spilled
        self.reward_range = self.reward_max - self.reward_min

        super().__init__(**kwargs)
        self.get_cached_configs_and_states(cached_states_path, self.num_variations)

        if observation_mode in ['point_cloud', 'key_point']:
            if observation_mode == 'key_point':
                obs_dim = 0
            else:
                max_particle_num = 13 * 13 * 13 * 4
                obs_dim = max_particle_num * 3
                self.particle_obs_dim = obs_dim
            # z and theta of the second cup (poured_box) does not change and thus are omitted.
            obs_dim += 8      # Pos (x) and shape (w, h, l) reset of the cup, torus x, torus y, if on box.
            self.observation_space = Box(low=np.array([-np.inf] * obs_dim), high=np.array([np.inf] * obs_dim), dtype=np.float32)
        elif observation_mode == 'cam_rgb':
            self.observation_space = Box(low=-np.inf, high=np.inf, shape=(self.camera_height, self.camera_width, 3),
                                         dtype=np.float32)

        default_config = self.get_default_config()
        if action_mode == 'direct':
            self.action_direct_dim = 1
            action_low = np.array([-0.011])
            action_high = np.array([0.011])
            self.action_space = Box(action_low, action_high, dtype=np.float32)
        elif action_mode in ['sawyer', 'franka']:
            self.action_tool = RobotBase(action_mode)
        else:
            raise NotImplementedError

    def get_default_config(self):
        config = {
            'torus': {
                'radius': 0.03, # if self.action_mode in ['sawyer', 'franka'] else 0.1,
                'rest_dis_coef': 0.5,
                'num': 5,
                'size': 0.2,
            },
            'box': { # all these will be overwritten by generate_env_variations
                'box_dis_x': 0.6,
                'box_dis_z': 0.6,
                'height': 0.6,
            },
            'static_friction': 0.5,
            'dynamic_friction': 1.0,
            'camera_name': 'default_camera',
        }
        return config

    def generate_env_variation(self, num_variations=5,  **kwargs):
        """
        TODO: add more randomly generated configs instead of using manually specified configs.
        """
        num_low = 1
        num_high = 2
        size_low = 0.12
        size_high = 0.28
        box_height_low = 0.1
        box_height_high = 0.3
        generated_configs = []
        generated_init_states = []
        config = self.get_default_config()
        config_variations = [copy.deepcopy(config) for _ in range(num_variations)]

        idx = 0
        while idx < num_variations:
            num = np.random.randint(num_low, num_high)
            size = np.random.uniform(size_low, size_high)

            particle_radius = config['torus']['radius'] * config['torus']['rest_dis_coef']
            estimated_particle_num_height = int(size / 0.05 -1)
            print("estimated_particle_num_height: ", estimated_particle_num_height)
            height = num * estimated_particle_num_height * particle_radius + num * size * 4 * particle_radius

            estimated_particle_num_width = int(size / 0.05 * 3 + 2)
            box_dis_x = estimated_particle_num_width * particle_radius + particle_radius  # box floor length
            box_dis_z = estimated_particle_num_width * particle_radius + particle_radius  # box width

            config_variations[idx]['torus']['lower_x'] = - box_dis_x / 2.
            config_variations[idx]['torus']['lower_z'] = - box_dis_z / 2.

            print("num {} size {}".format(num, config['torus']['size']))
            config_variations[idx]['torus']['num'] = num
            config_variations[idx]['torus']['size'] = size

            config_variations[idx]['box']['height'] = np.random.uniform(box_height_low, box_height_high)
            config_variations[idx]['torus']['height'] = estimated_particle_num_height * particle_radius + \
                config_variations[idx]['box']['height']
            config_variations[idx]['box']['box_dis_x'] = box_dis_x
            config_variations[idx]['box']['box_dis_z'] = box_dis_z
            print("box x, y, z: ", box_dis_x, height, box_dis_z)

            if self.set_scene(config_variations[idx]):
                init_state = copy.deepcopy(self.get_state())

                generated_configs.append(config_variations[idx])
                generated_init_states.append(init_state)
                idx += 1

        combined = [generated_configs, generated_init_states]

        return generated_configs, generated_init_states

    def get_config(self):
        if self.deterministic:
            config_idx = 0
        else:
            config_idx = np.random.randint(len(self.config_variations))

        self.config = self.config_variations[config_idx]
        return self.config

    def _reset(self):
        '''
        reset to environment to the initial state.
        return the initial observation.
        '''
        self.inner_step = 0
        self.performance_init = None
        info = self._get_info()
        self.performance_init = info['performance']
        return self._get_obs()

    def get_state(self):
        '''
        get the postion, velocity of flex particles, and postions of flex shapes.
        '''
        particle_pos = pyflex.get_positions()
        particle_vel = pyflex.get_velocities()
        shape_position = pyflex.get_shape_states()
        return {'particle_pos': particle_pos, 'particle_vel': particle_vel, 'shape_pos': shape_position,
                'box_x': self.box_x, 'box_states': self.box_states, 'box_params': self.box_params, 'config_id': self.current_config_id}

    def set_state(self, state_dic):
        '''
        set the postion, velocity of flex particles, and postions of flex shapes.
        '''
        self.box_params = state_dic['box_params']
        pyflex.set_positions(state_dic["particle_pos"])
        pyflex.set_velocities(state_dic["particle_vel"])
        pyflex.set_shape_states(state_dic["shape_pos"])
        self.box_x = state_dic['box_x']
        self.box_states = state_dic['box_states']
        for _ in range(5):
            pyflex.step()

    def initialize_camera(self):
        '''
        set the camera width, height, position and angle.
        **Note: width and height is actually the screen width and screen height of FLex.
        I suggest to keep them the same as the ones used in pyflex.cpp.
        '''
        # if self.action_mode in ['sawyer', 'franka']:
        self.camera_params = {
            'default_camera': {'pos': np.array([0.64, 1.8, 0.85]),
                            'angle': np.array([0 * np.pi, -65 / 180. * np.pi, 0]),
                            'width': self.camera_width,
                            'height': self.camera_height},
            'cam_2d': {'pos': np.array([0.5, .7, 4.]),
                    'angle': np.array([0, 0, 0.]),
                    'width': self.camera_width,
                    'height': self.camera_height}
        }
        # else:
        #     self.camera_params = {
        #     'default': {'pos': np.array([2.2, 2.7, 0.3]),
        #                        'angle': np.array([0.45 * np.pi, -60 / 180. * np.pi, 0]),
        #                        'width': self.camera_width,
        #                        'height': self.camera_height},
        #     'cam_2d': {'pos': np.array([0.5, .7, 4.]),
        #                'angle': np.array([0, 0, 0.]),
        #                'width': self.camera_width,
        #                'height': self.camera_height}
        # }

    def set_box_params(self, config=None):
        params = config['box']

        self.height = params['height']

        # TODO: correctly determine the box size
        particle_radius = config['torus']['radius'] * config['torus']['rest_dis_coef']
        self.box_dis_x = params['box_dis_x']
        self.box_dis_z = params['box_dis_z']

        self.x_center = 0
        params['box_x_center'] = 0

        self.box_params = params

    def set_scene(self, config, states=None):
        '''
        Construct the passing water scence.
        '''
        # create fluid
        super().set_scene(config)  # do not sample fluid parameters, as it's very likely to generate very strange fluid

        # compute box params
        if states is None:
            self.set_box_params(config)
        else:
            box_params = states['box_params']
            self.height = box_params['height']
            self.box_dis_x = box_params['box_dis_x']
            self.box_dis_z = box_params['box_dis_z']
            self.box_params = box_params
            self.x_center = 0

        # create box
        self.create_box(self.box_dis_x, self.box_dis_z, self.height)

        # move box to be at ground or on the table
        self.box_states = self.init_box_state(self.x_center, 0, self.box_dis_x, self.box_dis_z, self.height)

        pyflex.set_shape_states(self.box_states)

        # record box floor center x
        self.box_x = self.x_center

        # no cached init states passed in
        if states is None:
            for _ in range(50):
                pyflex.step()
                pyflex.render()

            return True

        else:  # set to passed-in cached init states
            self.set_state(states)

    def _get_obs(self):
        '''
        return the observation based on the current flex state.
        '''
        if self.observation_mode == 'cam_rgb':
            return self.get_image(self.camera_width, self.camera_height)
        elif self.observation_mode in ['point_cloud', 'key_point']:
            if self.observation_mode == 'point_cloud':
                particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3].flatten()
                pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float)
                pos[:len(particle_pos)] = particle_pos
            else:
                pos = np.empty(0, dtype=np.float)

            particle_state = pyflex.get_positions().reshape((-1, self.dim_position))
            torus_center_y = np.mean(particle_state[:, 1])
            torus_center_x = np.mean(particle_state[:, 0])

            on_box = float(torus_center_y >= self.height)
            within_box_bound = float(torus_center_x > self.box_x - 0.5 * self.box_dis_x and
                torus_center_x < self.box_x + self.box_dis_x * 0.5)
            cup_state = np.array([self.box_x, self.box_dis_x, self.box_dis_z, self.height,
                torus_center_x, torus_center_y, on_box, within_box_bound])
            return np.hstack([pos, cup_state]).flatten()
        else:
            raise NotImplementedError

    def compute_reward(self, obs=None, action=None, set_prev_reward=False):
        """
        The reward is the negative distance to the target.
        If the torus falls off the box, give a large negative penalty.
        """

        state_dic = self.get_state()
        particle_state = state_dic['particle_pos'].reshape((-1, self.dim_position))
        torus_center_y = np.mean(particle_state[:, 1])
        torus_center_x = np.mean(particle_state[:, 0])

        reward = -self.distance_coef * np.abs((self.terminal_x - self.box_x))
        box_x_low = self.box_x - self.box_dis_x / 2.
        box_x_high = self.box_x + self.box_dis_x / 2.
        if torus_center_y < self.height or torus_center_x < box_x_low or torus_center_x > box_x_high:
            reward -= self.torus_penalty_coef
        return reward

    def _get_info(self):
        state_dic = self.get_state()
        particle_state = state_dic['particle_pos'].reshape((-1, self.dim_position))
        torus_center_y = np.mean(particle_state[:, 1])
        torus_center_x = np.mean(particle_state[:, 0])

        reward = -self.distance_coef * np.abs((self.terminal_x - self.box_x))
        box_x_low = self.box_x - self.box_dis_x / 2.
        box_x_high = self.box_x + self.box_dis_x / 2.
        if torus_center_y < self.height or torus_center_x < box_x_low or torus_center_x > box_x_high:
            reward -= self.torus_penalty_coef

        performance = reward
        performance_init =  performance if self.performance_init is None else self.performance_init  # Use the original performance
        normalized_performance = (performance - performance_init) / (self.reward_max - performance_init)

        return {'performance': performance,
                'normalized_performance': normalized_performance,
                'distance_to_target': np.abs((self.terminal_x - self.box_x)),
                'torus_on': float(torus_center_y >= self.height)}

    def _step(self, action):
        '''
        action: np.ndarray of dim 1x1, dx, which specifies how much to move on the x-axis.
        '''
        # make action as increasement, clip its range
        dx = action[0]
        dx = np.clip(dx, a_min=self.action_space.low[0], a_max=self.action_space.high[0])
        x = self.box_x + dx

        # move the box
        new_states = self.move_box(self.box_states, x)
        self.box_states = new_states
        self.box_x = x
        self.box_x = np.clip(self.box_x, a_min=self.min_x, a_max=self.max_x)

        # pyflex takes a step to update the box and the water fluid
        pyflex.set_shape_states(self.box_states)
        pyflex.step()

        self.inner_step += 1

    def create_box(self, box_dis_x, box_dis_z, height):
        """
        create a box.
        dis_x: the length of the box
        dis_z: the width of the box
        height: the height of the box.

        the halfEdge determines the center point of each wall.
        """
        center = np.array([0., 0., 0.])
        quat = quatFromAxisAngle([0, 0, -1.], 0.)
        boxes = []

        # a single box
        halfEdge = np.array([box_dis_x / 2., height / 2., box_dis_z / 2.])
        boxes.append([halfEdge, center, quat])

        for i in range(len(boxes)):
            halfEdge = boxes[i][0]
            center = boxes[i][1]
            quat = boxes[i][2]
            pyflex.add_box(halfEdge, center, quat)

        return boxes

    def move_box(self, prev_states, x):
        '''
        given the previous states of the box, move it in 1D along x-axis.

        state:
        0-3: current (x, y, z) coordinate of the center point
        3-6: previous (x, y, z) coordinate of the center point
        6-10: current quat
        10-14: previous quat
        '''
        quat_curr = quatFromAxisAngle([0, 0, -1.], 0.)
        states = np.zeros((1, self.dim_shape_state))

        for i in range(1):
            states[i][3:6] = prev_states[i][:3]
            states[i][10:] = prev_states[i][6:10]

        x_center = x
        y = self.height / 2.

        # floor: center position does not change
        states[0, :3] = np.array([x_center, y, 0.])
        states[:, 6:10] = quat_curr

        return states

    def init_box_state(self, x, y, box_dis_x, box_dis_z, height):
        '''
        set the initial state of the box.
        '''
        dis_x, dis_z = box_dis_x, box_dis_z
        x_center, y_curr, y_last = x, y, 0.
        if self.action_mode in ['sawyer', 'franka']:
            y_curr = y_last = 0.56 # NOTE: robotics table
        quat = quatFromAxisAngle([0, 0, -1.], 0.)

        # states of a single box
        states = np.zeros((1, self.dim_shape_state))

        # a single box
        states[0, :3] = np.array([x_center, height / 2., 0.])
        states[0, 3:6] = np.array([x_center, height / 2., 0.])

        states[:, 6:10] = quat
        states[:, 10:] = quat

        return states

    def in_box(self, water, box_states, border, height, return_sum=True):
        '''
        judge whether a water particle is in the poured box
        water: [x, y, z, 1/m] water particle state.
        '''
        # floor, left, right, back, front
        # state:
        # 0-3: current (x, y, z) coordinate of the center point
        # 3-6: previous (x, y, z) coordinate of the center point
        # 6-10: current quat
        # 10-14: previous quat
        x_lower = box_states[1][0] - border / 2.
        x_upper = box_states[2][0] + border / 2.
        z_lower = box_states[3][2] - border / 2.
        z_upper = box_states[4][2] + border / 2
        # y_lower = box_states[0][1] - border / 2.
        # y_upper = box_states[0][1] + height + border / 2.
        # x, y, z = water[:, 0], water[:, 1], water[:, 2]
        x, z = water[:, 0], water[:, 2]

        # * (y >= y_lower) * (y <= y_upper)
        res = (x >= x_lower) * (x <= x_upper) * (z >= z_lower) * (z <= z_upper)
        if return_sum:
            res = np.sum(res)
            return res
        return res

    def _get_current_water_height(self):
        pos = pyflex.get_positions().reshape(-1, 4)
        return np.max(pos[:, 1])


##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
## Created by: Yaoyao Liu
## Modified from: https://github.com/cbfinn/maml
## Tianjin University
## liuyaoyao@tju.edu.cn
## Copyright (c) 2019
##
## This source code is licensed under the MIT-style license found in the
## LICENSE file in the root directory of this source tree
##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

""" Data generator for pre-train phase. """
import numpy as np
import os
import random
import tensorflow as tf
from tqdm import trange

from tensorflow.python.platform import flags
from utils.misc import get_pretrain_images

FLAGS = flags.FLAGS

class PreDataGenerator(object):
    """The class to generate episodes for pre-train phase."""
    def __init__(self):
        self.num_classes = FLAGS.way_num
        self.img_size = (FLAGS.img_size, FLAGS.img_size)
        self.dim_input = np.prod(self.img_size)*3
        self.pretrain_class_num = FLAGS.pretrain_class_num
        self.pretrain_batch_size = FLAGS.pretrain_batch_size
        pretrain_folder = FLAGS.pretrain_folders

        pretrain_folders = [os.path.join(pretrain_folder, label) for label in os.listdir(pretrain_folder) if os.path.isdir(os.path.join(pretrain_folder, label))]
        self.pretrain_character_folders = pretrain_folders

    def make_data_tensor(self):
        """The function to make tensor for the tensorflow model."""
        print('Generating pre-training data')
        all_filenames_and_labels = []
        folders = self.pretrain_character_folders

        for idx, path in enumerate(folders):
            all_filenames_and_labels += get_pretrain_images(path, idx)
        random.shuffle(all_filenames_and_labels)
        all_labels = [li[0] for li in all_filenames_and_labels]
        all_filenames = [li[1] for li in all_filenames_and_labels]
        filename_queue = tf.train.string_input_producer(tf.convert_to_tensor(all_filenames), shuffle=False)
        label_queue = tf.train.slice_input_producer([tf.convert_to_tensor(all_labels)], shuffle=False)
        image_reader = tf.WholeFileReader()
        _, image_file = image_reader.read(filename_queue)

        image = tf.image.decode_jpeg(image_file, channels=3)
        image.set_shape((self.img_size[0],self.img_size[1],3))
        image = tf.reshape(image, [self.dim_input])
        image = tf.cast(image, tf.float32) / 255.0

        num_preprocess_threads = 1
        min_queue_examples = 256
        batch_image_size = self.pretrain_batch_size
        image_batch, label_batch = tf.train.batch([image, label_queue], batch_size = batch_image_size, num_threads=num_preprocess_threads,capacity=min_queue_examples + 3 * batch_image_size)
        label_batch = tf.one_hot(tf.reshape(label_batch, [-1]), self.pretrain_class_num)
        return image_batch, label_batch


# -*- coding: utf-8 -*-

from __future__ import print_function
import base64
from PIL import Image
import os
from collections import defaultdict
from itertools import product
from sklearn.model_selection import train_test_split
import shutil
import re
import glob
from scipy import ndimage
from six.moves import cPickle as pickle
import tensorflow as tf
import numpy as np
from six.moves import range
import sys
from ctypes import *
import math
import random
from box import BOX
from image import IMAGE
from metadata import METADATA
from detection import DETECTION

def sample(probs):
    s = sum(probs)
    probs = [a/s for a in probs]
    r = random.uniform(0, 1)
    for i in range(len(probs)):
        r = r - probs[i]
        if r <= 0:
            return i
    return len(probs)-1

def c_array(ctype, values):
    arr = (ctype*len(values))()
    arr[:] = values
    return arr

def get_setting():
    subpath = os.path.join(os.path.dirname(os.path.abspath(__file__)))
    path = os.path.abspath(os.path.join(subpath, os.pardir))
    path = path + "/setting.conf"
    with open(path) as file:
        lines = [line.rstrip('\n') for line in file]
    d = {}
    for line in lines:
        t = line.split(':')
        d[t[0]] = t[1]
    return d

setting_dict = get_setting()
libdarknet_path = setting_dict['libdarknet_path']
yolo_config_path = setting_dict['yolo_config_path']
yolo_weight_path = setting_dict['yolo_weight_path']
meta_path = setting_dict['meta_path']

lib = CDLL(libdarknet_path, RTLD_GLOBAL)
lib.network_width.argtypes = [c_void_p]
lib.network_width.restype = c_int
lib.network_height.argtypes = [c_void_p]
lib.network_height.restype = c_int

predict = lib.network_predict
predict.argtypes = [c_void_p, POINTER(c_float)]
predict.restype = POINTER(c_float)

set_gpu = lib.cuda_set_device
set_gpu.argtypes = [c_int]

make_image = lib.make_image
make_image.argtypes = [c_int, c_int, c_int]
make_image.restype = IMAGE

get_network_boxes = lib.get_network_boxes
get_network_boxes.argtypes = [c_void_p, c_int, c_int, c_float, c_float, POINTER(c_int), c_int, POINTER(c_int)]
get_network_boxes.restype = POINTER(DETECTION)

make_network_boxes = lib.make_network_boxes
make_network_boxes.argtypes = [c_void_p]
make_network_boxes.restype = POINTER(DETECTION)

free_detections = lib.free_detections
free_detections.argtypes = [POINTER(DETECTION), c_int]

free_ptrs = lib.free_ptrs
free_ptrs.argtypes = [POINTER(c_void_p), c_int]

network_predict = lib.network_predict
network_predict.argtypes = [c_void_p, POINTER(c_float)]

reset_rnn = lib.reset_rnn
reset_rnn.argtypes = [c_void_p]

load_net = lib.load_network
load_net.argtypes = [c_char_p, c_char_p, c_int]
load_net.restype = c_void_p

do_nms_obj = lib.do_nms_obj
do_nms_obj.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]

do_nms_sort = lib.do_nms_sort
do_nms_sort.argtypes = [POINTER(DETECTION), c_int, c_int, c_float]

free_image = lib.free_image
free_image.argtypes = [IMAGE]

letterbox_image = lib.letterbox_image
letterbox_image.argtypes = [IMAGE, c_int, c_int]
letterbox_image.restype = IMAGE

load_meta = lib.get_metadata
lib.get_metadata.argtypes = [c_char_p]
lib.get_metadata.restype = METADATA

load_image = lib.load_image_color
load_image.argtypes = [c_char_p, c_int, c_int]
load_image.restype = IMAGE

rgbgr_image = lib.rgbgr_image
rgbgr_image.argtypes = [IMAGE]

predict_image = lib.network_predict_image
predict_image.argtypes = [c_void_p, IMAGE]
predict_image.restype = POINTER(c_float)

net = load_net(yolo_config_path, yolo_weight_path, 0)
meta = load_meta(meta_path)

def classify(net, meta, im):
    out = predict_image(net, im)
    res = []
    for i in range(meta.classes):
        res.append((meta.names[i], out[i]))
    res = sorted(res, key=lambda x: -x[1])
    return res

def detect(net, meta, image, thresh=.5, hier_thresh=.5, nms=.45):
    im = load_image(image, 0, 0)
    num = c_int(0)
    pnum = pointer(num)
    predict_image(net, im)
    dets = get_network_boxes(net, im.w, im.h, thresh, hier_thresh, None, 0, pnum)
    num = pnum[0]
    if (nms): do_nms_obj(dets, num, meta.classes, nms);

    res = []
    for j in range(num):
        for i in range(meta.classes):
            if dets[j].prob[i] > 0:
                b = dets[j].bbox
                res.append((meta.names[i], dets[j].prob[i], (b.x, b.y, b.w, b.h)))
    res = sorted(res, key=lambda x: -x[1])
    free_image(im)
    free_detections(dets, num)
    return res


"""Class that schedules motion on can bus."""
import asyncio
from collections import defaultdict
import logging
from typing import List, Set, Tuple, Iterator, Union
import numpy as np

from opentrons_hardware.firmware_bindings import ArbitrationId
from opentrons_hardware.firmware_bindings.constants import NodeId
from opentrons_hardware.drivers.can_bus.can_messenger import CanMessenger
from opentrons_hardware.firmware_bindings.messages import MessageDefinition
from opentrons_hardware.firmware_bindings.messages.message_definitions import (
    ClearAllMoveGroupsRequest,
    AddLinearMoveRequest,
    MoveCompleted,
    ExecuteMoveGroupRequest,
    HomeRequest,
    GripperGripRequest,
    GripperHomeRequest,
    TipActionRequest,
    TipActionResponse,
)
from opentrons_hardware.firmware_bindings.messages.payloads import (
    AddLinearMoveRequestPayload,
    ExecuteMoveGroupRequestPayload,
    HomeRequestPayload,
    GripperMoveRequestPayload,
    TipActionRequestPayload,
)
from .constants import interrupts_per_sec
from opentrons_hardware.hardware_control.motion import (
    MoveGroups,
    MoveGroupSingleAxisStep,
    MoveGroupSingleGripperStep,
    MoveGroupTipActionStep,
    MoveType,
    SingleMoveStep,
)
from opentrons_hardware.firmware_bindings.utils import (
    UInt8Field,
    UInt32Field,
    Int32Field,
)
from opentrons_hardware.firmware_bindings.messages.fields import (
    PipetteTipActionTypeField,
)
from opentrons_hardware.hardware_control.motion import MoveStopCondition
from opentrons_hardware.hardware_control.motion_planning.move_utils import (
    MoveConditionNotMet,
)
from .types import NodeDict

log = logging.getLogger(__name__)

_AcceptableMoves = Union[MoveCompleted, TipActionResponse]
_CompletionPacket = Tuple[ArbitrationId, _AcceptableMoves]
_Completions = List[_CompletionPacket]

class MoveGroupRunner:
    """A move command scheduler."""

    def __init__(self, move_groups: MoveGroups, start_at_index: int = 0) -> None:
        """Constructor.

        Args:
            move_groups: The move groups to run.
            start_at_index: The index the MoveGroupManager will start at
        """
        self._move_groups = move_groups
        self._start_at_index = start_at_index
        self._is_prepped: bool = False

    @staticmethod
    def _has_moves(move_groups: MoveGroups) -> bool:
        for move_group in move_groups:
            for move in move_group:
                for node, step in move.items():
                    return True
        return False

    async def prep(self, can_messenger: CanMessenger) -> None:
        """Prepare the move group. The first thing that happens during run().

        prep() and execute() can be used to replace a single call to run() to
        ensure tighter timing, if you want something else to start as soon as
        possible to the actual execution of the move.
        """
        if not self._has_moves(self._move_groups):
            log.debug("No moves. Nothing to do.")
            return
        await self._clear_groups(can_messenger)
        await self._send_groups(can_messenger)
        self._is_prepped = True

    async def execute(self, can_messenger: CanMessenger) -> NodeDict[float]:
        """Execute a pre-prepared move group. The second thing that run() does.

        prep() and execute() can be used to replace a single call to run() to
        ensure tighter timing, if you want something else to start as soon as
        possible to the actual execution of the move.
        """
        if not self._has_moves(self._move_groups):
            log.debug("No moves. Nothing to do.")
            return {}
        if not self._is_prepped:
            raise RuntimeError("A group must be prepped before it can be executed.")
        move_completion_data = await self._move(can_messenger)
        return self._accumulate_move_completions(move_completion_data)

    async def run(self, can_messenger: CanMessenger) -> NodeDict[float]:
        """Run the move group.

        Args:
            can_messenger: a can messenger

        Returns:
            The current position after the move for all the axes that
            acknowledged completing moves.

        This function first prepares all connected devices to move (by sending
        all the data for the moves over) and then executes the move with a
        single call.

        prep() and execute() can be used to replace a single call to run() to
        ensure tighter timing, if you want something else to start as soon as
        possible to the actual execution of the move.
        """
        await self.prep(can_messenger)
        return await self.execute(can_messenger)

    @staticmethod
    def _accumulate_move_completions(completions: _Completions) -> NodeDict[float]:
        position: NodeDict[List[Tuple[Tuple[int, int], float]]] = defaultdict(list)
        for arbid, completion in completions:
            position[NodeId(arbid.parts.originating_node_id)].append(
                (
                    (
                        completion.payload.group_id.value,
                        completion.payload.seq_id.value,
                    ),
                    float(completion.payload.current_position_um.value) / 1000.0,
                )
            )
        # for each node, pull the position from the completion with the largest
        # combination of group id and sequence id
        return {
            node: next(
                reversed(
                    sorted(poslist, key=lambda position_element: position_element[0])
                )
            )[1]
            for node, poslist in position.items()
        }

    async def _clear_groups(self, can_messenger: CanMessenger) -> None:
        """Send commands to clear the message groups.

        Args:
            can_messenger: a can messenger
        """
        await can_messenger.send(
            node_id=NodeId.broadcast,
            message=ClearAllMoveGroupsRequest(),
        )

    async def _send_groups(self, can_messenger: CanMessenger) -> None:
        """Send commands to set up the message groups."""
        for group_i, group in enumerate(self._move_groups):
            for seq_i, sequence in enumerate(group):
                for node, step in sequence.items():
                    await can_messenger.send(
                        node_id=node,
                        message=self._get_message_type(
                            step, group_i + self._start_at_index, seq_i
                        ),
                    )

    def _convert_velocity(
        self, velocity: Union[float, np.float64], interrupts: int
    ) -> Int32Field:
        return Int32Field(int((velocity / interrupts) * (2**31)))

    def _get_message_type(
        self, step: SingleMoveStep, group: int, seq: int
    ) -> MessageDefinition:
        """Return the correct payload type."""
        if isinstance(step, MoveGroupSingleAxisStep):
            return self._get_stepper_motor_message(step, group, seq)
        elif isinstance(step, MoveGroupTipActionStep):
            return self._get_tip_action_motor_message(step, group, seq)
        else:
            return self._get_brushed_motor_message(step, group, seq)

    def _get_brushed_motor_message(
        self, step: MoveGroupSingleGripperStep, group: int, seq: int
    ) -> MessageDefinition:
        if step.move_type == MoveType.home:
            home_payload = GripperMoveRequestPayload(
                group_id=UInt8Field(group),
                seq_id=UInt8Field(seq),
                duration=UInt32Field(int(step.duration_sec * step.pwm_frequency)),
                freq=UInt32Field(int(step.pwm_frequency)),
                duty_cycle=UInt32Field(int(step.pwm_duty_cycle)),
            )
            return GripperHomeRequest(payload=home_payload)

        else:
            linear_payload = GripperMoveRequestPayload(
                group_id=UInt8Field(group),
                seq_id=UInt8Field(seq),
                duration=UInt32Field(int(step.duration_sec * step.pwm_frequency)),
                freq=UInt32Field(int(step.pwm_frequency)),
                duty_cycle=UInt32Field(int(step.pwm_duty_cycle)),
            )
            return GripperGripRequest(payload=linear_payload)

    def _get_stepper_motor_message(
        self, step: MoveGroupSingleAxisStep, group: int, seq: int
    ) -> MessageDefinition:
        if step.move_type == MoveType.home:
            home_payload = HomeRequestPayload(
                group_id=UInt8Field(group),
                seq_id=UInt8Field(seq),
                duration=UInt32Field(int(step.duration_sec * interrupts_per_sec)),
                velocity=self._convert_velocity(
                    step.velocity_mm_sec, interrupts_per_sec
                ),
            )
            return HomeRequest(payload=home_payload)
        else:
            linear_payload = AddLinearMoveRequestPayload(
                request_stop_condition=UInt8Field(step.stop_condition),
                group_id=UInt8Field(group),
                seq_id=UInt8Field(seq),
                duration=UInt32Field(int(step.duration_sec * interrupts_per_sec)),
                acceleration=Int32Field(
                    int(
                        (
                            step.acceleration_mm_sec_sq
                            / interrupts_per_sec
                            / interrupts_per_sec
                        )
                        * (2**31)
                    )
                ),
                velocity=Int32Field(
                    int((step.velocity_mm_sec / interrupts_per_sec) * (2**31))
                ),
            )
            return AddLinearMoveRequest(payload=linear_payload)

    def _get_tip_action_motor_message(
        self, step: MoveGroupTipActionStep, group: int, seq: int
    ) -> TipActionRequest:
        tip_action_payload = TipActionRequestPayload(
            group_id=UInt8Field(group),
            seq_id=UInt8Field(seq),
            duration=UInt32Field(int(step.duration_sec * interrupts_per_sec)),
            velocity=self._convert_velocity(step.velocity_mm_sec, interrupts_per_sec),
            action=PipetteTipActionTypeField(step.action),
            request_stop_condition=UInt8Field(step.stop_condition),
        )
        return TipActionRequest(payload=tip_action_payload)

    async def _move(self, can_messenger: CanMessenger) -> _Completions:
        """Run all the move groups."""
        scheduler = MoveScheduler(self._move_groups)
        try:
            can_messenger.add_listener(scheduler)
            completions = await scheduler.run(can_messenger)
        finally:
            can_messenger.remove_listener(scheduler)
        return completions

class MoveScheduler:
    """A message listener that manages the sending of execute move group messages."""

    def __init__(self, move_groups: MoveGroups) -> None:
        """Constructor."""
        # For each move group create a set identifying the node and seq id.
        self._moves: List[Set[Tuple[int, int]]] = []
        self._durations: List[float] = []
        self._stop_condition: List[MoveStopCondition] = []
        for move_group in move_groups:
            move_set = set()
            duration = 0.0
            for seq_id, move in enumerate(move_group):
                move_set.update(set((k.value, seq_id) for k in move.keys()))
                duration += float(list(move.values())[0].duration_sec)
                for step in move_group[seq_id]:
                    self._stop_condition.append(move_group[seq_id][step].stop_condition)

            self._moves.append(move_set)
            self._durations.append(duration)
        log.debug(f"Move scheduler running for groups {move_groups}")
        self._completion_queue: asyncio.Queue[_CompletionPacket] = asyncio.Queue()
        self._event = asyncio.Event()

    def _remove_move_group(
        self, message: _AcceptableMoves, arbitration_id: ArbitrationId
    ) -> None:
        seq_id = message.payload.seq_id.value
        group_id = message.payload.group_id.value
        node_id = arbitration_id.parts.originating_node_id
        log.info(
            f"Received completion for {node_id} group {group_id} seq {seq_id}"
            ", which "
            f"{'is' if (node_id, seq_id) in self._moves[group_id] else 'isn''t'}"
            " in group"
        )
        try:
            self._moves[group_id].remove((node_id, seq_id))
            self._completion_queue.put_nowait((arbitration_id, message))
        except KeyError:
            log.warning(
                f"Got a move ack for ({node_id}, {seq_id}) which is not in this "
                "group; may have leaked from an earlier timed-out group"
            )

        if not self._moves[group_id]:
            log.info(f"Move group {group_id} has completed.")
            self._event.set()

    def _handle_move_completed(self, message: MoveCompleted) -> None:
        group_id = message.payload.group_id.value
        ack_id = message.payload.ack_id.value
        if self._stop_condition[
            group_id
        ] == MoveStopCondition.limit_switch and ack_id != UInt8Field(2):
            if ack_id == UInt8Field(1):
                condition = "Homing timed out."
                log.warning(f"Homing failed. Condition: {condition}")
                raise MoveConditionNotMet()

    def _handle_tip_action(self, message: TipActionResponse) -> None:
        group_id = message.payload.group_id.value
        ack_id = message.payload.ack_id.value
        limit_switch = bool(
            self._stop_condition[group_id] == MoveStopCondition.limit_switch
        )
        success = message.payload.success.value
        # TODO need to add tip action type to the response message.
        if limit_switch and limit_switch != ack_id and not success:
            condition = "Tip still detected."
            log.warning(f"Drop tip failed. Condition {condition}")
            raise MoveConditionNotMet()
        elif not limit_switch and not success:
            condition = "Tip not detected."
            log.warning(f"Pick up tip failed. Condition {condition}")
            raise MoveConditionNotMet()

    def __call__(
        self, message: MessageDefinition, arbitration_id: ArbitrationId
    ) -> None:
        """Incoming message handler."""
        if isinstance(message, MoveCompleted):
            self._remove_move_group(message, arbitration_id)
            self._handle_move_completed(message)
        elif isinstance(message, TipActionResponse):
            self._remove_move_group(message, arbitration_id)
            self._handle_tip_action(message)

    async def run(self, can_messenger: CanMessenger) -> _Completions:
        """Start each move group after the prior has completed."""
        for group_id in range(len(self._moves)):
            self._event.clear()

            log.info(f"Executing move group {group_id}.")
            await can_messenger.send(
                node_id=NodeId.broadcast,
                message=ExecuteMoveGroupRequest(
                    payload=ExecuteMoveGroupRequestPayload(
                        group_id=UInt8Field(group_id),
                        # TODO (al, 2021-11-8): The triggers should be populated
                        #  with actual values.
                        start_trigger=UInt8Field(0),
                        cancel_trigger=UInt8Field(0),
                    )
                ),
            )

            try:
                # TODO: The max here can be removed once can_driver.send() no longer
                # returns before the message actually hits the bus. Right now it
                # returns when the message is enqueued in the kernel, meaning that
                # for short move durations we can see the timeout expiring before
                # the execute even gets sent.
                await asyncio.wait_for(
                    self._event.wait(), max(1.0, self._durations[group_id] * 1.1)
                )
            except asyncio.TimeoutError:
                log.warning("Move set timed out")

        def _reify_queue_iter() -> Iterator[_CompletionPacket]:
            while not self._completion_queue.empty():
                yield self._completion_queue.get_nowait()

        return list(_reify_queue_iter())