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

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

    fsps_templates : bool
        Use FSPS templates instead of line complex templates

    line_list : list
        List of individual lines to include in the templates

    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.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`.

        PAs : list of int
            List of PAs to fit. If None, fit all PAs.

        extensions : list of str
            List of extensions to fit. If None, fit all extensions.

        min_ivar : float
            Minimum inverse variance threshold for fitting.
            Default = 0.01.

        overlap_threshold : float
            Threshold for masking overlapping regions.
            Default = 3.

        verbose : bool
            If True, print verbose output.

        eazyp : `~eazy.photoz.EazyPhot`
            EazyPhot object for photometric redshift fitting.

        eazy_ix : int
            Index of the EazyPhot object to use for photometric redshift fitting.
            Default = 0.

        MW_EBV : float
            Milky Way E(B-V) reddening value.
            Default = 0.

        chi2_threshold : float
            Threshold for identifying bad PAs based on chi-squared values.
            Default = 1.5 . 

        min_DoF : int
            Minimum degrees of freedom threshold for identifying bad PAs.
            Default = 200.

        """
        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.asarray(self.idf, dtype=int)

        # 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):
        """
        Parses the `self.beams` list and extracts various information such 
        as the number of beams, the extension versions and other attributes.
        
        """

        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.0

        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`.

        Parameters
        ----------
        st : StackFitter
            The StackFitter object to be appended.

        """

        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,
    ):
        """
        Check for bad PAs and return fit log, keep dictionary, and 
        a flag indicating if there are bad PAs.

        Parameters
        ----------
        poly_order : int
            The order of the polynomial templates used for fitting. 
            Default is 1.

        chi2_threshold : float
            The threshold value for the chi-square ratio below which a PA 
            is considered good. Default is 1.5.

        fit_background : bool
            Whether to fit the background when calculating the chi-square ratio. 
            Default is True.

        reinit : bool
            Whether to reinitialize the `self` object with the beams 
            corresponding to the good PAs. Default is True.

        verbose : bool
            Whether to print verbose output. Default is False.

        Returns
        -------
        fit_log : OrderedDict
            A dictionary containing the chi-square values, degrees of freedom, 
            and chi-square ratios for each PA and group.

        keep_dict : OrderedDict
            A dictionary containing the list of PAs to keep for each group.
            
        has_bad : bool
            A flag indicating if there are bad PAs.

        """

        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.

        Parameters
        ----------
        z : float
            Redshift. Default is 0.

        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.
            Default is False.

        eazyp : `~eazy.photoz.EazyPhot`
            EazyPhot object containing the EAZY templates and filter information.
            Default in None.

        ix : int
            Index of the object in the EazyPhot object.
            Default is 0.

        order : int
            The order of the polynomial templates used for fitting. 
            Default is 1.

        scale_fit : `~scipy.optimize.OptimizeResult`
            Object containing the result of the optimization process.
            Default is None.

        method : str
            The optimization method to use. Default is 'BFGS'.

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

        background : `~np.ndarray`
            Background model

        full : `~np.ndarray`
            Full model

        full_coeffs : `~np.ndarray`
            Coefficients of the templates

        full_coeffs_err : `~np.ndarray`
            Uncertainties of the template coefficients

        scale_fit : `~scipy.optimize.OptimizeResult`
            Object containing the result of the optimization process.

        """
        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.0)[
            ok_phot
        ]
        photom_flam = (eazyp.fnu[ix, :] * eazyp.to_flam * eazyp.ext_corr / 100.0)[
            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.0e18 / 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.0e-4
        init = np.zeros(order + 1)
        init[0] = 10.0

        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([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
                covar = utils.safe_invert(np.dot(AxT.T, AxT))
                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.

        Parameters
        ----------
        p : array_like
            The coefficients of the polynomial function.

        Ax : array_like
            The scaled template matrix.

        spec_wave : array_like
            The wavelength array.

        Nphot : int
            The number of photometric points.

        Next : int
            The number of template components.

        Returns
        -------
        AxT : array_like
            The scaled template matrix.

        """
        from np.polynomial import Polynomial

        scale = np.ones(Ax.shape[1])
        scale[:-Nphot] = Polynomial(p / 10.0)((spec_wave - 1.0e4) / 1000.0)
        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

        Parameters
        ----------
        p : float
            Scale term between photometry and spectra.

        Ax : `~np.ndarray`
            Array of shape (M, N) representing the design matrix.

        data : `~np.ndarray`
            Array of shape (N,) representing the observed data.

        spec_wave : `~np.ndarray`
            Array of shape (N,) representing the wavelength of the spectra.

        fit_mask : `~np.ndarray`
            Boolean array of shape (N,) representing the mask for fitting.

        sivarf : `~np.ndarray`
            Array of shape (N,) representing the inverse variance of the spectra.

        Nphot : int
            Number of photometric data points.

        Next : int
            Number of background components.
            
        return_coeffs : bool
            Flag indicating whether to return the coefficients.

        Returns
        -------
        if return_coeffs:
        
            coeffs : `~np.ndarray`
                Array of shape (M,) representing the coefficients of the fit.

            full : `~np.ndarray`
                Array of shape (N,) representing the full model.
            
            resid : `~np.ndarray`
                Array of shape (N,) representing the residuals.

            chi2 : float
                The chi-squared value of the fit.
            
            AxT : `~np.ndarray`
                The scaled template matrix.
        
        else:
        
            chi2 : float
                The chi-squared value of the fit.

        """
        import scipy.optimize
        from np.polynomial import Polynomial

        scale = np.ones(Ax.shape[1])
        scale[:-Nphot] = Polynomial(p / 10.0)((spec_wave - 1.0e4) / 1000.0)
        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)
        if return_coeffs:
            return coeffs, full, resid, chi2, AxT
        else:
            return chi2
        
    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.0

            except:
                igmz = 1.0

            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
                covar = utils.safe_invert(np.dot(AxT.T, AxT))
                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).
            Default is 0.005.

        zr : list
            Redshift range to consider.
            Default is [0.4, 3.4].

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

        make_plot : bool
            Make the diagnostic plot. Default is True.

        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. 
            Default is 'templates.npy'.

        verbose : bool
            Print the redshift grid steps. Default is True.

        outlier_threshold : float
            Threshold for masking outliers in the fit. Default is 1e30.

        eazyp : `~eazy.photoz.EazyTemplate`
            EazyTemplate object for combined fits. Default is None.

        ix : int
            Index of the object in the eazyp object. Default is 0.

        order : int
            Order of the polynomial fit to the combined spectrum. Default is 0.

        scale_fit : float
            Scale factor for the combined spectrum. Default is None.

        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.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.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.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.0

                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"):
        """
        Generate 1D templates based on the given coefficients.
        
        Parameters
        ----------
        coeffs : array_like
            Coefficients used to scale the templates.
            
        templates_file : str
            Path to the file containing the templates. 
            Default is "templates.npy".

        Returns
        -------
        tc : ndarray
            The generated 1D template for the continuum component.

        tl : ndarray
            The generated 1D template for the line component.

        """

        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.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):
        """
        Mask pixels in beams of a given grism that are greater than
        `threshold` sigma times the minimum of all beams of that
        grism

        Parameters
        ----------

        threshold : float
            Threshold for masking. Default is 3.

        verbose : bool
            Print the number of pixels masked. Default is True.

        """
        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`.

        scale_pz : bool
            Whether to scale the p(z) to DoF=1. Default is True.

        show_2d : bool
            Whether to show the 2D spectra. Default is False.

        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.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.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.0
            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.0e4

                ax_i.imshow(
                    clean,
                    vmin=-0.02 * ymax,
                    vmax=1.1 * ymax,
                    origin="lower",
                    extent=[w[0], w[-1], 0.0, 1.0],
                    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.0e30
        ymax = -1.0e30
        wmin = 1.0e30
        wmax = -1.0e30

        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.0e4

            # Do we need to convert to F-lambda units?
            if E.is_flambda:
                unit_corr = 1.0
                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.0 / 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.0e-19
            er *= unit_corr / 1.0e-19
            flm *= unit_corr / 1.0e-19

            f_alpha = (
                1.0 / 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.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.0
            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.0,
        min_ivar=0.001,
        MW_EBV=0.0,
    ):
        """
        The StackedSpectrum class.

        Parameters
        ----------
        file : str
            The path to the FITS file containing the data. 
            Default is "gnt_18197.stack.G141.285.fits".

        sys_err : float
            The systematic error to be used in the calculations. 
            Default is 0.02.

        mask_min : float
            The minimum value for the mask. Default is 0.1.

        extver : str
            The extension version. Default is "G141".

        mask_threshold : int
            The threshold value for the mask. Default is 7.

        fcontam : float
            The contamination weighting factor. Default is 1.0.

        min_ivar : float
            The minimum inverse variance value. Default is 0.001.

        MW_EBV : float
            The Milky Way extinction value. Default is 0.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.0
        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.0

        # Contamination weighting
        self.fcontam = fcontam
        if ("CONTAM", extver) in self.hdulist:
            self.contam = np.abs(self.hdulist["CONTAM", extver].data * 1.0)
            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.0
            self.contamf = self.scif * 0.0

            self.weight = self.sci * 0.0 + 1
            self.weightf = self.scif * 0.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.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

        Parameters
        ----------
        h : `astropy.io.fits.Header`
            The FITS header containing the WCS keywords.

        Returns
        -------
        w : `~np.array`
            The generated wavelength array.

        """
        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.0e4

        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, **kwargs):
        """
        Optimally-weighted 1D extraction

        Parameters
        ----------
        data : `~np.ndarray`
            2D array of the data to extract.

        bin : int
            Binning factor. Default is 0.

        ivar : array-like
            Inverse variance to replace `self.ivar`

        Returns
        -------
        wave : `~np.array`
            Wavelength array
        
        opt_flux : `~np.array`
            Optimal flux array

        opt_rms : `~np.array`
            Optimal RMS array

        """
        import scipy.ndimage as nd

        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.0 / 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_numba.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):
            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.0

            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.polynomial.Polynomial.fit(x/1.e4, y, deg=1)
                c_i = np.array([0.98928126, 0.07498747])
                scale = np.polynomial.Polynomial(c_i)(self.wave / 1.0e4)
                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

        Parameters
        ----------
        spectrum_1d : tuple
            A tuple containing the wavelength and flux arrays. Default is None.

        is_cgsm, in_place : dummy parameters to match `MultiBeam.compute_model`.

        Returns
        -------
        model : `~np.ndarray`
            The generated model spectrum.

        """
        from .utils_numba.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)
        return model
