# Path: py/desisim/spec_qa/utils.py
# def elg_flux_lim(z, oii_flux):
#     '''Assess which objects pass the ELG flux limit
#     Uses DESI document 318 from August 2014
# 
#     Parameters
#     ----------
#     z : ndarray
#       ELG redshifts
#     oii_flux : ndarray
#       [OII] fluxes
#     '''
#     # Init mask
#     mask = np.array([False]*len(z))
#     #
#     # Flux limits from document
#     OII_lim = np.array([10., 9., 9., 9., 9])*1e-17
#     # Get bin
#     #zbin = (z-0.6)/zstep
#     zbins = [0.6, 0.8, 1., 1.2, 1.4, 5.6]
#     z_i = np.digitize(z, zbins) - 1
#     gd_elg = (z_i >= 0) & (oii_flux>0.)
#     #gd = np.where((zbin>0.) & (zbin<len(OII_lim)) & (oii_flux>0.))[0]
#     # Fill gd
#     OII_match = np.ones_like(gd_elg) * 9e99
#     for ss in range(len(zbins)-1):
#         OII_match[z_i==ss] = OII_lim[ss]
#     # Query
#     really_gd = (oii_flux > OII_match) & gd_elg
#     mask[really_gd] = True
#     # Return
#     return mask
# 
# def get_sty_otype():
#     '''Styles for plots'''
#     sty_otype = dict(ELG={'color':'green', 'pcolor':'Greens', 'lbl':'ELG'},
#         LRG={'color':'red', 'lbl':'LRG', 'pcolor':'Reds'},
#                      BGS={'color':'orange', 'lbl':'BGS', 'pcolor':'Oranges'},
#                      MWS={'color':'purple', 'lbl':'MWS', 'pcolor':'Purples'},
#         QSO={'color':'blue', 'lbl':'QSO', 'pcolor':'Blues'},
#         QSO_L={'color':'blue', 'lbl':'QSO z>2.1', 'pcolor':'Blues'},
#         QSO_T={'color':'cyan', 'lbl':'QSO z<2.1', 'pcolor':'GnBu'})
#     return sty_otype
# 
# def catastrophic_dv(objtype):
#     '''Pass back catastrophic velocity limit for given objtype
#     From DESI document 318 (August 2014) in docdb
# 
#     Parameters
#     ----------
#     objtype : str
#       Object type, e.g. 'ELG', 'LRG'
#     '''
#     cat_dict = dict(ELG=1000., LRG=1000., QSO_L=1000., QSO_T=1000., STAR=1000., WD=1000.)
#     #
#     return cat_dict[objtype]
# 
# def match_otype(tbl, objtype):
#     """ Generate a mask for the input objtype
#     :param tbl:
#     :param objtype: str
#     :return: targets: bool mask
#     """
#     if objtype in ['BGS']:
#         targets = (tbl['DESI_TARGET'] & desi_mask['BGS_ANY']) != 0
#     elif objtype in ['MWS']:
#         targets = (tbl['DESI_TARGET'] & desi_mask['MWS_ANY']) != 0
#     elif objtype in ['QSO_L']:
#         targets =  match_otype(tbl, 'QSO') & (tbl['TRUEZ'] >= 2.1)
#     elif objtype in ['QSO_T']:
#         targets =  match_otype(tbl, 'QSO') & (tbl['TRUEZ'] < 2.1)
#     else:
#         targets = (tbl['DESI_TARGET'] & desi_mask[objtype]) != 0
#     # Return
#     return targets


import matplotlib
import numpy as np
import sys, os, pdb, glob
import matplotlib.gridspec as gridspec
import desispec.io
from matplotlib import pyplot as plt
from astropy.io import fits
from astropy.table import Table, vstack, hstack, MaskedColumn, join
    from scipy import constants
from .utils import elg_flux_lim, get_sty_otype, catastrophic_dv, match_otype
from desiutil.log import get_logger, DEBUG
    from desimodel.footprint import radec2pix
    from astropy.stats import sigma_clip
    # Init
    nrow = len(simz_tab)
    stypes = np.char.rstrip(simz_tab['TEMPLATETYPE'].astype(str))

    # Object type
    if objtype is None:
        objtype_mask = np.array([True]*nrow)
    else:
        if objtype in ['STAR', 'WD', 'QSO']:
            objtype_mask = stypes == objtype
        else:
            objtype_mask = match_otype(simz_tab, objtype)  # Use DESI_TARGET when possible
    # Redshift analysis
    z_mask = simz_tab['Z'].mask == False  # Not masked in Table
    # Survey
    survey_mask = (simz_tab['Z'].mask == False)
    elg = np.where(match_otype(simz_tab, 'ELG') & survey_mask)[0]
    if len(elg) > 0:
        elg_mask = elg_flux_lim(simz_tab['TRUEZ'][elg],
                                simz_tab['OIIFLUX'][elg])
        # Update
        survey_mask[elg[~elg_mask]] = False
    # zwarn -- Masked array
    zwarn_mask = np.array([False]*nrow)
    idx = np.where((simz_tab['ZWARN'] == 0) & (simz_tab['ZWARN'].mask == False))[0]
    zwarn_mask[idx] = True
    # Catastrophic/Good (This gets a bit more messy...)
    dv_mask = np.array([True]*nrow)
    for obj in np.unique(stypes):
        if obj in ['ELG','LRG','QSO_L','QSO_T', 'BGS', 'MWS']:  # Use DESI_TARGET when possible

# Path: py/desisim/spec_qa/utils.py
# def elg_flux_lim(z, oii_flux):
#     '''Assess which objects pass the ELG flux limit
#     Uses DESI document 318 from August 2014
# 
#     Parameters
#     ----------
#     z : ndarray
#       ELG redshifts
#     oii_flux : ndarray
#       [OII] fluxes
#     '''
#     # Init mask
#     mask = np.array([False]*len(z))
#     #
#     # Flux limits from document
#     OII_lim = np.array([10., 9., 9., 9., 9])*1e-17
#     # Get bin
#     #zbin = (z-0.6)/zstep
#     zbins = [0.6, 0.8, 1., 1.2, 1.4, 5.6]
#     z_i = np.digitize(z, zbins) - 1
#     gd_elg = (z_i >= 0) & (oii_flux>0.)
#     #gd = np.where((zbin>0.) & (zbin<len(OII_lim)) & (oii_flux>0.))[0]
#     # Fill gd
#     OII_match = np.ones_like(gd_elg) * 9e99
#     for ss in range(len(zbins)-1):
#         OII_match[z_i==ss] = OII_lim[ss]
#     # Query
#     really_gd = (oii_flux > OII_match) & gd_elg
#     mask[really_gd] = True
#     # Return
#     return mask
# 
# def get_sty_otype():
#     '''Styles for plots'''
#     sty_otype = dict(ELG={'color':'green', 'pcolor':'Greens', 'lbl':'ELG'},
#         LRG={'color':'red', 'lbl':'LRG', 'pcolor':'Reds'},
#                      BGS={'color':'orange', 'lbl':'BGS', 'pcolor':'Oranges'},
#                      MWS={'color':'purple', 'lbl':'MWS', 'pcolor':'Purples'},
#         QSO={'color':'blue', 'lbl':'QSO', 'pcolor':'Blues'},
#         QSO_L={'color':'blue', 'lbl':'QSO z>2.1', 'pcolor':'Blues'},
#         QSO_T={'color':'cyan', 'lbl':'QSO z<2.1', 'pcolor':'GnBu'})
#     return sty_otype
# 
# def catastrophic_dv(objtype):
#     '''Pass back catastrophic velocity limit for given objtype
#     From DESI document 318 (August 2014) in docdb
# 
#     Parameters
#     ----------
#     objtype : str
#       Object type, e.g. 'ELG', 'LRG'
#     '''
#     cat_dict = dict(ELG=1000., LRG=1000., QSO_L=1000., QSO_T=1000., STAR=1000., WD=1000.)
#     #
#     return cat_dict[objtype]
# 
# def match_otype(tbl, objtype):
#     """ Generate a mask for the input objtype
#     :param tbl:
#     :param objtype: str
#     :return: targets: bool mask
#     """
#     if objtype in ['BGS']:
#         targets = (tbl['DESI_TARGET'] & desi_mask['BGS_ANY']) != 0
#     elif objtype in ['MWS']:
#         targets = (tbl['DESI_TARGET'] & desi_mask['MWS_ANY']) != 0
#     elif objtype in ['QSO_L']:
#         targets =  match_otype(tbl, 'QSO') & (tbl['TRUEZ'] >= 2.1)
#     elif objtype in ['QSO_T']:
#         targets =  match_otype(tbl, 'QSO') & (tbl['TRUEZ'] < 2.1)
#     else:
#         targets = (tbl['DESI_TARGET'] & desi_mask[objtype]) != 0
#     # Return
#     return targets


# Path: py/desisim/spec_qa/redshifts.py
import matplotlib
import numpy as np
import sys, os, pdb, glob
import matplotlib.gridspec as gridspec
import desispec.io
from matplotlib import pyplot as plt
from astropy.io import fits
from astropy.table import Table, vstack, hstack, MaskedColumn, join
    from scipy import constants
from .utils import elg_flux_lim, get_sty_otype, catastrophic_dv, match_otype
from desiutil.log import get_logger, DEBUG
    from desimodel.footprint import radec2pix
    from astropy.stats import sigma_clip

    # Init
    nrow = len(simz_tab)
    stypes = np.char.rstrip(simz_tab['TEMPLATETYPE'].astype(str))

    # Object type
    if objtype is None:
        objtype_mask = np.array([True]*nrow)
    else:
        if objtype in ['STAR', 'WD', 'QSO']:
            objtype_mask = stypes == objtype
        else:
            objtype_mask = match_otype(simz_tab, objtype)  # Use DESI_TARGET when possible
    # Redshift analysis
    z_mask = simz_tab['Z'].mask == False  # Not masked in Table
    # Survey
    survey_mask = (simz_tab['Z'].mask == False)
    elg = np.where(match_otype(simz_tab, 'ELG') & survey_mask)[0]
    if len(elg) > 0:
        elg_mask = elg_flux_lim(simz_tab['TRUEZ'][elg],
                                simz_tab['OIIFLUX'][elg])
        # Update
        survey_mask[elg[~elg_mask]] = False
    # zwarn -- Masked array
    zwarn_mask = np.array([False]*nrow)
    idx = np.where((simz_tab['ZWARN'] == 0) & (simz_tab['ZWARN'].mask == False))[0]
    zwarn_mask[idx] = True
    # Catastrophic/Good (This gets a bit more messy...)
    dv_mask = np.array([True]*nrow)
    for obj in np.unique(stypes):
        if obj in ['ELG','LRG','QSO_L','QSO_T', 'BGS', 'MWS']:  # Use DESI_TARGET when possible

            omask = np.where(match_otype(simz_tab, obj))[0] # & (simz_tab['ZWARN']==0))[0]