#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SPECEX - SPECtra EXtractor.

Extract spectra from spectral data cubes.

Copyright (C) 2022-2023  Maurizio D'Addona <mauritiusdadd@gmail.com>
"""
from __future__ import absolute_import, division, print_function

import os
import sys
import json
import argparse
import logging

from typing import Optional, Any, List, Union

import numpy as np
import multiprocessing

from scipy.signal import savgol_filter
from astropy import wcs
from astropy.table import Table, Column
from astropy.io import fits
from astropy.io.registry import IORegistryError
from astropy.coordinates import SkyCoord
from astropy.time import Time
from astropy import units

import matplotlib.pyplot as plt

from specex.utils import (
    plot_zfit_check, get_log_img, stack, plot_scandata,
    get_spectrum_snr, get_spectrum_snr_emission, parse_regionfile,
    HAS_REGION, HAS_REPROJECT, simple_pbar_callback, do_reproject
)
from specex.sources import get_spectrum
from specex.cube import get_hdu, SpectralCube
from specex.exceptions import exception_handler

from specex.cube import (
    KNOWN_SPEC_EXT_NAMES,
    KNOWN_VARIANCE_EXT_NAMES,
    KNOWN_INVAR_EXT_NAMES,
    KNOWN_MASK_EXT_NAMES,
    KNOWN_RCURVE_EXT_NAMES
)

try:
    from specex.rrspecex import rrspecex, get_templates
except Exception:
    HAS_RR = False
else:
    HAS_RR = True

try:
    from IPython.core import ultratb
    from IPython import embed
except Exception:
    HAS_IPYTHON = False
else:
    HAS_IPYTHON = True

SPECTRA_WEIGHTING_MODES = {
    'none': 'no weighting',
    'whitelight': 'weighting using cube whitelight image',
    'log-whitelight': 'weighting using cube whitelight image in log10 scale',
}

if HAS_REPROJECT:
    SPECTRA_WEIGHTING_MODES['<file url>'] = 'weighting using an external image'
def __argshandler(options: Optional[List[str]] = None) -> Any:
    """
    Parse the arguments given by the user.

    Returns
    -------
    args: Namespace
        A Namespace containing the parsed arguments. For more information see
        the python documentation of the argparse module.
    """
    parser = argparse.ArgumentParser(
        description='Extract spectra from a datacube and optionally compute '
        'the redshift using rrspecex backend.'
    )
    parser.add_argument(
        'input_cube', metavar='SPEC_CUBE', type=str, nargs=1,
        help='The spectral cube in fits format from which spectra will be '
        'extracted.'
    )

    parser.add_argument(
        '--mode', metavar='WORKING_MODE', type=str, default='auto',
        help='Set the working mode of the program. Can be one of the following'
        ' modes: "auto", "kron_ellipse", "kron_circle", "circular_aperture". '
        'When using "kron_ellipse" mode, the input catalog must contain the '
        'column for objetcs kron radius, semi-major and semi-minor axes and '
        'for the rotation angle of the objects. In this mode all the pixels '
        'whitin this ellipse will be added together to compute the spectra. '
        '"kron_circular" mode requires the same parameters of the previous '
        'mode, but a circular aperture is used instead of an ellipse and the '
        'radius of this aperture is the circularized kron radius (i.e. the '
        'kron darius multiplied by the axis ratio q = minor-axis/major-axis. '
        'In circular_aperture mode, no morphological parameters are required '
        'to be in the input catalog and the aperture size is set using the '
        'input parameter --aperture-size. In "auto" mode, the best working '
        'mode is selected based on the parameters available in the input '
        'catalog.'
    )

    parser.add_argument(
        '--catalog', metavar='CATALOG', type=str, default=None,
        help='A catalog containing the position and the (optional) morphology '
        'of the objects of which you want to extract the spectra. The catalog '
        'must contain at least the RA and DEC coordinates and optionally the '
        'dimensions of the isophotal axes and the rotation angle of the '
        ' isophotal ellipse, as well as the kron radius. The way how these '
        'parameter are used depends on the program working mode (see the '
        '--mode section for more information). The catalog format is'
        ' by default compatible by the one generated by Sextractor, so the '
        'column name used to read the aformentioned qauntities are '
        'respectively ALPHA_J2000, DELTA_J2000, A_IMAGE, B_IMAGE, THETA_IMAGE '
        'and KRON_RADIUS. If the catalog has different column names for these '
        'quantities, they can be overridden with a corresponding --key-* '
        'argument (see below). The sky frame of the RA and DEC coordinates can'
        ' be set using --cat-skyframe.'
    )

    parser.add_argument(
        '--regionfile', metavar='REGFILE', type=str, default=None,
        help='A region file containing regions from which extract spectra. '
        'The name of the refion will be used as ID for the spectra.'
    )

    parser.add_argument(
        '--nspectra', metavar='N', type=int, default=0,
        help='Number of spectra to extract (usefull for debugging purposes).'
    )

    parser.add_argument(
        '--sn-threshold', metavar='SN_THRESH', type=float, default=1.5,
        help='Set the signal to Noise ratio threshold: objects with their '
        'spectrum having a SN ratio lower than threshold will be ignored. '
        'If SN_THRESH is set to a value equal or lower than 0, SN filtering '
        'is disabled'
    )

    parser.add_argument(
        '--no-nans', action='store_true', default=False,
        help="This option tells the program to replace eventual NaNs in the"
        "extracted spectrum with linarly interpolated data. In this case the "
        "output fits files contain an additional extension named 'NAN_MASK' "
        "that identifies the previously bad pixels in the spectrum."
    )

    parser.add_argument(
        '--info-hdu', metavar='INFO_HDU', type=int, default=0,
        help='The HDU containing cube metadata. If this argument '
        'Set this to -1 to automatically detect the HDU containing the info. '
        'NOTE that this value is zero indexed (i.e. firts HDU has index 0).'
    )

    parser.add_argument(
        '--spec-hdu', metavar='SPEC_HDU', type=int, default=None,
        help='The HDU containing the spectral data to use. If this argument '
        'Set this to -1 to automatically detect the HDU containing spectra. '
        'NOTE that this value is zero indexed (i.e. second HDU has index 1).'
    )

    parser.add_argument(
        '--var-hdu', metavar='VAR_HDU', type=int, default=None,
        help='The HDU containing the variance of the spectral data. '
        'Set this to -1 if no variance data is present in the cube. '
        'The default value is %(metavar)s=%(default)s.'
        'NOTE that this value is zero indexed (i.e. third HDU has index 2).'
    )

    parser.add_argument(
        '--mask-hdu', metavar='MASK_HDU', type=int, default=None,
        help='The HDU containing the valid pixel mask of the spectral data. '
        'Set this to -1 if no mask is present in the cube. '
        'The default value is %(metavar)s=%(default)s.'
        'NOTE that this value is zero indexed (i.e. fourth HDU has index 3).'
    )

    parser.add_argument(
        '--invert-mask', action='store_true', default=False,
        help='Set whether to use the inverse of the cube mask.'
    )

    parser.add_argument(
        '--aperture-size', metavar='APERTURE', type=str, default='0.8mas',
        help='Set the diameter of a fixed circula aperture used in the '
        '"fixed-aperture" mode. The pixels within this aperture will be '
        'used for spectra extraction. Units can be specified using the names '
        'compatible with astropy.units (eg. 5mas, 1arcsec, etc.).'
        'The default value is %(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--anulus-size', metavar='R1,R2', type=str, default='None',
        help='Set the inner and outer radii in arcsecs of a fixed anular '
        'aperture used in  the "fixed-aperture" mode. The pixels within this '
        'anulus will be used to compute the background nearby the objects. If '
        'The default value is %(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--weighting', metavar='MODE', type=str, default='None',
        help='Set the spatial weighting for the spectra. Available modes are:'
        + ';'.join([
            f"{mode_name} for {mode_desc}"
            for mode_name, mode_desc in SPECTRA_WEIGHTING_MODES.items()
        ]) +
        ". Note that in order to use extrenal images as weightmap, the python "
        "package 'reproject' is needed. (Is reproject installed: " +
        "YES" if HAS_REPROJECT else "NO" +
        ")."
    )

    parser.add_argument(
        '--cat-pixelscale', metavar='PIXEL_SCALE', type=str, default='1mas',
        help='Set the scale size in mas of the major and minor axes. Useful if'
        'A and B are given in pixels instead of physical units (like in the '
        'case of catalogs extracted by sextractor). Units can be specified '
        'using the astropy.units names (ie. 5mas, 1arcsec, etc.). If it is '
        'adimensional, then the value is interpreted as mas per pixel. '
        'NOTE that this is not the pixel scale of the cube but it is the size '
        ' of a pixel of the image used to generate the input catalog. '
        'This parameter will be ignored if using a region file as input.'
        'The default value is %(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--cat-format', metavar='FORMAT', type=str, default=None,
        help='Manually set the format of the input catalog to %(metavar)s. '
        'If not specified the program will try to automatically detect the '
        'format.'
    )

    parser.add_argument(
        '--cutouts-image', metavar='IMAGE', type=str, default=None,
        help='The path of a FITS image (grayscale or single/multiext RGB) from'
        'which to extract cutout of the objects. If this option is not '
        'specified then the cutouts will be extracted directly from the '
        'input cube.'
    )

    parser.add_argument(
        '--cutouts-size', metavar='SIZE', type=float, default=10,
        help='Size of the cutouts of the object in pixel or arcseconds. '
        'The default value is %(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--key-ra', metavar='RA_KEY', type=str, default='ALPHA_J2000',
        help='Set the name of the column from which to read the RA of the '
        'objects. The default value is %(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--key-dec', metavar='DEC_KEY', type=str, default='DELTA_J2000',
        help='Set the name of the column from which to read the RA of the '
        'objects. The default value is %(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--key-a', metavar='A_KEY', type=str, default='A_IMAGE',
        help='Set the name of the column from which to read the size of the'
        'semi-major isophotal axis. The default value is '
        '%(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--key-b', metavar='B_KEY', type=str, default='B_IMAGE',
        help='Set the name of the column from which to read the size of the'
        'semi-minor isophotal axis. The default value is '
        '%(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--key-theta', metavar='THETA_KEY', type=str, default='THETA_IMAGE',
        help='Set the name of the column from which to read the rotation angle'
        ' of the isophotal ellipse. The default value is '
        '%(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--key-kron', metavar='KRON_KEY', type=str, default='KRON_RADIUS',
        help='Set the name of the column from which to read the size of the '
        'Kron radius (in units of A or B). The default value is '
        '%(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--key-id', metavar='OBJID_KEY', type=str, default=None,
        help='Set the name of the column from which to read the id of the '
        'object to which the spectral data belong. If this argument is not '
        'specified then a suitable column will be selected according to its '
        'name or, if none is found, a new progressive id will be generated. '
        'This id will be included in the fits header.'
    )

    parser.add_argument(
        '--cat-skyframe', metavar='SKY_FRAME', type=str, default='icrs',
        help='Set the sky frame RA and DEC values of the input catalog. This '
        'option is ignored if a region file is uesed for the extraction. '
        'The default value is %(metavar)s=%(default)s.'
    )

    parser.add_argument(
        '--outdir', metavar='DIR', type=str, default=None,
        help='Set the directory where extracted spectra will be outputed. '
        'If this parameter is not specified, then a new directory will be '
        'created based on the name of the input cube.'
    )

    parser.add_argument(
        '--check-images', action='store_true', default=False,
        help='Whether or not to generate check images (cutouts, etc.).'
    )

    parser.add_argument(
        "--checkimg-outdir", type=str, default=None, required=False,
        help='Set the directory where check images are saved (when they are '
        'enabled thorugh the appropriate parameter).'
    )

    if HAS_RR:
        rr_help_note = ''
    else:
        rr_help_note = ' (WARNING: this option is unavailable since redrock' \
                       ' seems not to be installed).'

    parser.add_argument(
        "--zbest", type=str, default=None, required=False,
        help='Whether to find redshift from the spectra. ' + rr_help_note
    )

    parser.add_argument(
        "--priors", type=str, default=None, required=False,
        help='optional redshift prior file. ' + rr_help_note
    )

    parser.add_argument(
        "--nminima", type=int, default=3, required=False,
        help="the number of redshift minima to search. " + rr_help_note
    )

    parser.add_argument(
        "--mp", type=int, default=0, required=False,
        help="if not using MPI, the number of multiprocessing processes to use"
        " (defaults to half of the hardware threads). " + rr_help_note
    )

    parser.add_argument(
        "-t", "--templates", type=str, default=None, required=False,
        help="template file or directory. " + rr_help_note
    )

    parser.add_argument(
        "--debug", action='store_true', default=False,
        help="Enable some debug otupout on plots and on the console."
    )

    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)

    if args.cutouts_image and not os.path.exists(args.cutouts_image):
        print(f"The file {args.cutouts_image} does not exist!")
        sys.exit(1)

    if args.mp == 0:
        args.mp = int(multiprocessing.cpu_count() / 4)

    return args


# TODO: resample to constant log(lambda)
def get_spplate_fits(cube_header, spec_header, obj_ids, spec_data,
                     var_data=None, and_mask_data=None, or_mask_data=None,
                     wdisp_data=None, sky_data=None):
    """
    Generate a plate fits file similar to the ones used in SDSS.

    Parameters
    ----------
    cube_header : TYPE
        DESCRIPTION.
    spec_header : TYPE
        DESCRIPTION.
    obj_ids : TYPE
        DESCRIPTION.
    spec_data : TYPE
        DESCRIPTION.
    var_data : TYPE, optional
        DESCRIPTION. The default is None.
    and_mask_data : TYPE, optional
        DESCRIPTION. The default is None.
    or_mask_data : TYPE, optional
        DESCRIPTION. The default is None.
    wdisp_data : TYPE, optional
        DESCRIPTION. The default is None.
    sky_data : TYPE, optional
        DESCRIPTION. The default is None.

    Returns
    -------
    hdul : TYPE
        DESCRIPTION.

    """
    spec_hdu = fits.PrimaryHDU()
    ivar_hdu = fits.ImageHDU()
    andmsk_hdu = fits.ImageHDU()
    ormsk_hdu = fits.ImageHDU()
    wdisp_hdu = fits.ImageHDU()
    tbl_hdu = fits.BinTableHDU()
    sky_hdu = fits.ImageHDU()

    if isinstance(spec_data, np.ma.MaskedArray):
        spec_data = spec_data.filled(np.nan)

    if isinstance(var_data, np.ma.MaskedArray):
        var_data = var_data.filled(np.nan)

    spec_hdu.data = spec_data.T
    spec_hdu.header['PLATEID'] = 0

    try:
        spec_hdu.header['MJD'] = cube_header['MJD']
    except KeyError:
        try:
            spec_hdu.header['MJD'] = cube_header['MJD-OBS']
        except KeyError:
            print(
                "WARNING: no MJD found in the cube metadata, using extraction "
                "time instead!",
                file=sys.stderr
            )
            spec_hdu.header['MJD'] = Time.now().mjd

    # See the following link for COEFF0 and COEFF1 definition
    # https://classic.sdss.org/dr7/products/spectra/read_spSpec.html
    try:
        spec_hdu.header["COEFF0"] = spec_header["COEFF0"]
        spec_hdu.header["COEFF1"] = spec_header["COEFF1"]
    except KeyError:
        print(
            "WARING: wavelength binning of the input cube is not compatible "
            "with boss log10 binning data model.\n        Computing the best "
            "values of COEFF0 and COEFF1 for this wavelengt range...",
            file=sys.stderr
        )
        # Compute coeff0 and coeff1 that best approximate wavelength range
        # of the datacube
        crpix = spec_header["CRPIX3"] - 1
        crval = spec_header["CRVAL3"]
        cd1 = spec_header["CD3_3"]

        w_end = spec_data.shape[1] - crpix

        coeff0 = np.log10(crval + cd1*(crpix))
        coeff1 = np.log10(1 + (w_end*cd1)/crval) / w_end

        spec_hdu.header['COEFF0'] = coeff0
        spec_hdu.header['COEFF1'] = coeff1

    spec_hdu.header["CRVAL1"] = spec_header["CRVAL3"]
    spec_hdu.header["CD1_1"] = spec_header["CD3_3"]
    spec_hdu.header["CDELT1"] = spec_header["CD3_3"]
    spec_hdu.header["CRPIX1"] = spec_header["CRPIX3"]
    spec_hdu.header['BUNIT'] = spec_header['BUNIT']
    spec_hdu.header["CUNIT1"] = spec_header["CUNIT3"]
    spec_hdu.header["CTYPE1"] = spec_header["CTYPE3"]

    if var_data is not None:
        ivar_hdu.data = 1 / var_data.T

    if wdisp_data is None:
        wdisp_hdu.data = np.zeros_like(spec_data.T)
    else:
        wdisp_hdu.data = wdisp_data.T

    info_table = Table()
    info_table.add_column(
        Column(data=obj_ids, name='FIBERID', dtype='>i4')
    )

    tbl_hdu.data = info_table.as_array()

    hdul = fits.HDUList([
        spec_hdu, ivar_hdu, andmsk_hdu, ormsk_hdu, wdisp_hdu, tbl_hdu, sky_hdu
    ])

    return hdul


def get_spsingle_fits(main_header, spec_wcs_header, obj_spectrum,
                      spec_hdu_header, obj_spectrum_var=None, no_nans=False):
    """
    Generate a fits containing the spectral data.

    Parameters
    ----------
    main_header : dict
        Dictionary containing the base information of the spectrum extraction.
        The keys of the dictionaries are the names of the fits CARDS that will
        be written, while the values can be either strings or tuples in the
        form of (cadr value, card comment).
    spec_wcs_header : astropy.io.fits.Header object
        Header containing the WCS information of the spectrum.
    obj_spectrum : numpy.ndarray
        The actual spectral data.
    spec_hdu_header : astropy.io.fits.Header object
        Header of the HDU of the original cube containing the spectral data.
    obj_spectrum_var : numpy.ndarray (default=None)
        The variance of the spectral data.
    no_nans : bool, optional
        If true, nan values in the spectrum are replaced by a linear
        interpolation. In this case an additional extension is added to the
        fits file, containing a mask for identifying the previously bax pixels.

    Returns
    -------
    hdul : astropy.io.fits.HDUList object
        The fits HDUL containing the extracted spectra.

    """
    my_time = Time.now()
    my_time.format = 'isot'

    if isinstance(obj_spectrum, np.ma.MaskedArray):
        obj_spectrum = obj_spectrum.filled(np.nan)

    if isinstance(obj_spectrum_var, np.ma.MaskedArray):
        obj_spectrum_var = obj_spectrum_var.filled(np.nan)

    # The primary hdu contains no data but a header with information on
    # how the spectrum was extracted
    main_hdu = fits.PrimaryHDU()
    for key, val in main_header.items():
        try:
            main_hdu.header[key] = val
        except ValueError:
            print(
                f"Invalid value for spec header: {key} = {val}",
                file=sys.stderr
            )

    spec_header = fits.Header()
    spec_header['BUNIT'] = spec_hdu_header['BUNIT']
    spec_header["OBJECT"] = spec_hdu_header["OBJECT"]
    spec_header["CUNIT1"] = spec_hdu_header["CUNIT3"]
    spec_header["CTYPE1"] = spec_hdu_header["CTYPE3"]
    spec_header["OBJECT"] = spec_hdu_header["OBJECT"]
    spec_header["CRPIX1"] = spec_hdu_header["CRPIX3"]
    spec_header["CRVAL1"] = spec_hdu_header["CRVAL3"]
    spec_header["CDELT1"] = spec_hdu_header["CD3_3"]

    nanmask = np.isnan(obj_spectrum)

    if no_nans:
        obj_spectrum[nanmask] = 0
        if obj_spectrum_var is not None:
            obj_spectrum_var[nanmask] = np.nanmax(obj_spectrum_var) * 100.0

    hdu_list = [
        main_hdu,
        fits.ImageHDU(
            data=obj_spectrum,
            header=spec_header,
            name='SPECTRUM'
        )
    ]

    if obj_spectrum_var is not None:
        var_header = fits.Header()
        var_header['BUNIT'] = spec_hdu_header['BUNIT']
        var_header["OBJECT"] = spec_hdu_header["OBJECT"]
        var_header["CUNIT1"] = spec_hdu_header["CUNIT3"]
        var_header["CTYPE1"] = spec_hdu_header["CTYPE3"]
        var_header["OBJECT"] = spec_hdu_header["OBJECT"]
        var_header["CRPIX1"] = spec_hdu_header["CRPIX3"]
        var_header["CRVAL1"] = spec_hdu_header["CRVAL3"]
        var_header["CDELT1"] = spec_hdu_header["CD3_3"]
        var_header["OBJECT"] = spec_hdu_header["OBJECT"]
        hdu_list.append(
            fits.ImageHDU(
                data=obj_spectrum_var,
                header=var_header,
                name='VARIANCE'
            )
        )

    hdu_list.append(
        fits.ImageHDU(
            data=nanmask.astype('uint8'),
            name='NAN_MASK'
        )
    )

    hdul = fits.HDUList(hdu_list)
    return hdul


def load_specex_file(filename, spec_hdu_index=None, var_hdu_index=None,
                     wd_hdu_index=None, mask_hdu_index=None, smoothing=0):
    """
    Load a spectrum file generated by specex (or any compatible file).

    Parameters
    ----------
    filename : str
        The path of the fits file containing the spectrum.
    spec_hdu_index : int or str, optional
        The index or name of the HDU containing the flux.
        The default is None.
    var_hdu_index : int or str, optional
        The index or name of the HDU containing the variance.
        The default is None.
    wd_hdu_index : int or str, optional
        The index or name of the HDU containing the resolution curve.
        The default is None.
    mask_hdu_index : int or str, optional
        The index or name of the HDU containing the flux nan mask.
        The default is None.
    smoothing : int, optional
        If greather than zero, the the spectrum is smoothed with a window of
        this size. The default is 0.

    Raises
    ------
    ValueError
        When flux and variance have different shape.

    Returns
    -------
    dict
        A dict containing the following keys:

        * 'id': the ID of the spectrum,
        * 'main_header': the FITS header of the primary HDU,
        * 'wavelength': the wavelengths of the dispersion grid,
        * 'flux': the fluxes,
        * 'variance': the variance of the fluxes,
        * 'r_curve': the resolution curve,
        * 'nan_mask': a mask indicating invalid fluxes,
        * 'flux_units': the units of the fluxes,
        * 'wavelength_units': the units of the wavelengths

    """
    with fits.open(filename) as hdul:
        valid_id_keys = [
            f"{i}{j}"
            for i in ['', 'OBJ', 'OBJ_', 'TARGET', 'TARGET_']
            for j in ['ID', 'NUMBER', 'UID', 'UUID']
        ]

        spec_id = None
        for hdu in hdul:
            for key in valid_id_keys:
                try:
                    spec_id = hdu.header[key]
                except KeyError:
                    continue
                else:
                    break

        main_header = get_hdu(
            hdul,
            hdu_index=0,
            valid_names=['PRIMARY', 'primary'],
            msg_index_error="WARNING: No Primary HDU",
            exit_on_errors=False
        ).header.copy()
        spec_hdu = get_hdu(
            hdul,
            hdu_index=spec_hdu_index,
            valid_names=KNOWN_SPEC_EXT_NAMES,
            msg_err_notfound="WARNING: No spectrum HDU",
            exit_on_errors=False
        )
        var_hdu = get_hdu(
            hdul,
            hdu_index=var_hdu_index,
            valid_names=KNOWN_VARIANCE_EXT_NAMES,
            msg_err_notfound="WARNING: No variance HDU",
            exit_on_errors=False
        )

        if var_hdu is None:
            ivar_hdu = get_hdu(
                hdul,
                hdu_index=var_hdu_index,
                valid_names=KNOWN_INVAR_EXT_NAMES,
                msg_err_notfound="WARNING: No inverse variance HDU",
                exit_on_errors=False
            )

        mask_hdu = get_hdu(
            hdul,
            hdu_index=mask_hdu_index,
            valid_names=KNOWN_MASK_EXT_NAMES,
            msg_err_notfound="WARNING: No nan mask HDU",
            exit_on_errors=False
        )

        wd_hdu = get_hdu(
            hdul,
            hdu_index=wd_hdu_index,
            valid_names=KNOWN_RCURVE_EXT_NAMES,
            msg_err_notfound="WARNING: No r curve HDU",
            exit_on_errors=False
        )

        smoothing_window_size = 4*smoothing + 1

        if spec_hdu is None:
            flux = None
            flux_units = None
            lam = None
            wavelength_units = None
        else:
            flux = spec_hdu.data.copy()
            spec_wcs = wcs.WCS(spec_hdu.header)
            if smoothing > 0:
                flux = savgol_filter(flux, smoothing_window_size, 3)
            flux = flux.astype('float32')

            pixel = np.arange(len(flux))
            if spec_wcs.has_spectral:
                lam = (
                    spec_wcs.pixel_to_world(pixel)
                    .Angstrom
                    .astype('float32')
                )
            else:
                try:
                    coeff0 = spec_hdu.header["COEFF0"]
                    coeff1 = spec_hdu.header["COEFF1"]
                except KeyError:
                    lam = None
                else:
                    lam = (10**(coeff0 + coeff1*pixel)).astype('float32')

            try:
                flux_units = spec_hdu.header['BUNIT']
            except KeyError:
                flux_units = None

            try:
                wavelength_units = spec_hdu.header['CUNIT1']
            except KeyError:
                wavelength_units = None

        if mask_hdu is None:
            flux_nan_mask = None
        else:
            flux_nan_mask = mask_hdu.data.astype(bool)
            if flux is not None:
                flux_nan_mask |= np.isnan(flux)

        if var_hdu is None:
            if ivar_hdu is None:
                var = None
            else:
                var = 1 / ivar_hdu.data.copy()
                if smoothing > 0:
                    var = savgol_filter(var, smoothing_window_size, 3)
                var = var.astype('float32')
        else:
            var = var_hdu.data.copy()
            if smoothing > 0:
                var = savgol_filter(var, smoothing_window_size, 3)
            var = var.astype('float32')

        if wd_hdu is None:
            wd = None
        else:
            wd = wd_hdu.data.copy()
            if smoothing > 0:
                wd = savgol_filter(wd, smoothing_window_size, 3)
            wd = wd.astype('float32')

        if (
            (flux is not None) and
            (var is not None) and
            flux.shape != var.shape
        ):
            raise ValueError(
                f"'{filename}' - "
                "Spectral data invalid or corruptede: Flux data shape "
                "do not match variance data one!"
            )

    return {
        'id': spec_id,
        'main_header': main_header,
        'wavelength': lam,
        'flux': flux,
        'variance': var,
        'r_curve': wd,
        'nan_mask': flux_nan_mask,
        'flux_units': flux_units,
        'wavelength_units': wavelength_units
    }


def speclist2ez(spec_files, spec_hdu_name='SPECTRUM', var_hdu_name='VARIANCE'):
    """
    Generate a list of spectra that can be imported into pandora EZ.

    Parameters
    ----------
    spec_files : list of str
        List of file paths.
    spec_hdu_name : str, optional
        The name of the HDU containing the spectral data.
        The default is 'SPECTRUM'.
    var_hdu_name : str, optional
        The name of the HDU containing the variance data.
        The default is 'VARIANCE'.

    Returns
    -------
    txt_data : TYPE
        DESCRIPTION.

    """
    txt_data = ""

    for file in spec_files:
        txt_data += f'{file}[{spec_hdu_name}](1) - '
        txt_data += f'{file}[{var_hdu_name}](1) -\n'
    return txt_data


def specex(options=None):
    """
    Run the main program.

    Returns
    -------
    None.

    """
    global HAS_IPYTHON

    # Get user input from argument parsers helper
    args = __argshandler(options)

    if args.debug and HAS_IPYTHON:
        sys.excepthook = exception_handler

    # Working mode
    mode = None

    if args.catalog is not None:
        if args.cat_format is not None:
            sources = Table.read(args.catalog, format=args.cat_format)
        else:
            try:
                sources = Table.read(args.catalog)
            except IORegistryError:
                try:
                    sources = Table.read(args.catalog, format='ascii')
                except IORegistryError:
                    print(
                        "ERROR: cannot determine the format of the input "
                        "catalog. Try to specify it manually using the "
                        "option --cat-format.",
                        file=sys.stderr
                    )
                    sys.exit(1)

        basename_with_ext = os.path.basename(args.catalog)
        pixelscale = units.Quantity(args.cat_pixelscale)
        coord_frame = args.cat_skyframe  # frame of source coordinates
        if str(pixelscale.unit.physical_type) == 'dimensionless':
            pixelscale = pixelscale * units.mas
    elif args.regionfile is not None:
        if not HAS_REGION:
            logging.error(
                "astropy regions package is needed to handle regionfiles!"
            )
            sys.exit(1)

        sources, coord_frame = parse_regionfile(
            args.regionfile,
            key_ra=args.key_ra,
            key_dec=args.key_dec,
            key_id=args.key_id,
            file_format=args.cat_format,
        )
        basename_with_ext = os.path.basename(args.regionfile)
        pixelscale = 1.0 * units.arcsec
        mode = 'regions'
    else:
        logging.error(
            "You must provide at least an input catalog or a region file!",
        )
        sys.exit(1)

    basename = os.path.splitext(basename_with_ext)[0]

    with SpectralCube.open(args.input_cube[0]) as my_cube:
        if args.debug:
            print(f"Flux HDU: {my_cube.spec_hdu.name}")
            if my_cube.var_hdu is not None:
                print(f"Variance HDU: {my_cube.var_hdu.name}")
            if my_cube.mask_hdu is not None:
                print(f"Mask HDU: {my_cube.mask_hdu.name}")

        # Get spectral and celstial WCS from the cube
        celestial_wcs = my_cube.getSpecWCS().celestial
        spectral_wcs = my_cube.getSpecWCS().spectral

        # If weighting with an external image, read that
        valid_wl_modes = ['none', 'whitelight', 'log-whitelight']
        weight_img: Union[None, np.ndarray] = None
        if args.weighting.lower() not in valid_wl_modes:
            if not os.path.isfile(args.weighting):
                logging.error(
                    f"weighting image does not exist '{args.weighting}' "
                )
                sys.exit(1)
            with fits.open(args.weighting) as hdul:
                wimg_hdu = hdul[0]
                wimg_shape = wimg_hdu.data.shape

                if len(wimg_shape) != 2:
                    logging.error(
                        "ERROR: only grayscale images are supported as "
                        "weightmap!"
                    )
                    sys.exit(1)
                elif HAS_REPROJECT:
                    # Reproject the external weight image to the
                    # datacube footprint
                    weight_img, wimg_mask = do_reproject(
                        wimg_hdu, my_cube
                    )

                    weight_img[wimg_mask == 0] = np.nan
                elif wimg_shape != my_cube.spec_hdu.data.shape[1:]:
                    logging.error(
                        "The shape of the external weightmap image does "
                        "not match the datacube and reproject is not "
                        "installed! Use --help for more information."
                    )
                    sys.exit(1)
                else:
                    weight_img = wimg_hdu.data.copy()

        # TODO: we assume here that the distortion across the sky plane is
        #       negligible, but it is very bold assumption.
        #       A more general approach that takes into account also distortion
        #       in function of the position on the sky is needed.
        # Get the pixelscales in arcsec/pixel.
        cube_pixelscale_x, cube_pixelscale_y = [
            units.Quantity(sc_val, u)
            for sc_val, u in zip(
                wcs.utils.proj_plane_pixel_scales(celestial_wcs),
                celestial_wcs.wcs.cunit
            )
        ]

        if args.debug:
            print(
                "Cube X and Y pixelscales are: "
                f"{cube_pixelscale_x.to(units.arcsec)}, "
                f"{cube_pixelscale_y.to(units.arcsec)}",
                file=sys.stderr
            )

        ra_unit = sources[args.key_ra].unit
        if ra_unit is None:
            if args.debug:
                logging.warning(
                    "RA data has no units, assuming values are in degrees!",
                )
            ra_unit = 'deg'

        dec_unit = sources[args.key_dec].unit
        if dec_unit is None:
            if args.debug:
                logging.warning(
                    "DEC data has no units, assuming values are in degrees!",
                )
            dec_unit = 'deg'

        # Get the objects positions in the cube (in pixels)
        obj_sky_coords = SkyCoord(
            sources[args.key_ra], sources[args.key_dec],
            unit=(ra_unit, dec_unit),
            frame=coord_frame
        )

        obj_x, obj_y = wcs.utils.skycoord_to_pixel(
            coords=obj_sky_coords,
            wcs=celestial_wcs
        )

        sources.add_column(obj_x, name='CUBE_X_IMAGE')
        sources.add_column(obj_y, name='CUBE_Y_IMAGE')

        img_height, img_width = my_cube.getSpatialSizePixels()
        yy, xx = np.meshgrid(
            np.arange(img_height),
            np.arange(img_width),
            indexing='ij'
        )

        extracted_data = np.zeros((img_height, img_width))

        # Chech if there is a footprint mask in the cube
        if my_cube.mask_hdu is not None:
            cube_footprint = my_cube.mask_hdu.data.prod(axis=0) == 0
            if args.invert_mask:
                cube_footprint = ~cube_footprint
        else:
            cube_footprint = None

        if args.outdir is None:
            outdir = f'{basename}_spectra'
        else:
            outdir = args.outdir

        if not os.path.isdir(outdir):
            os.mkdir(outdir)

        if mode is None:
            if args.mode == 'auto':
                if (
                        args.key_a in sources.colnames and
                        args.key_b in sources.colnames and
                        args.key_kron in sources.colnames
                ):
                    mode = 'kron_ellipse'
                else:
                    mode = 'circular_aperture'
            else:
                mode = args.mode

        print(f"Extracting spectra with strategy '{mode}'", file=sys.stderr)

        n_objects = len(sources)
        if args.nspectra:
            n_objects = int(args.nspectra)

        # Try to identify the column with object IDs
        if args.key_id is None:
            valid_id_keys = [
                f"{i}{j}"
                for i in ['', 'OBJ', 'OBJ_', 'TARGET', 'TARGET_']
                for j in ['ID', 'NUMBER', 'UID', 'UUID']
            ]

            for key in valid_id_keys:
                if key.lower() in sources.colnames:
                    key_id = key.lower()
                    break
                elif key.upper() in sources.colnames:
                    key_id = key.upper()
                    break
            else:
                key_id = None
        else:
            key_id = args.key_id

        out_specfiles = []
        specex_apertures = {}  # these values are in physical units (arcsecs)
        specex_anuli = {}  # these values are in physical units (arcsecs)
        source_ids = []

        valid_sources_mask = np.zeros(len(sources), dtype=bool)

        # Check wheter to use anuli to compute background
        anulus_str = args.anulus_size.lower().strip().split(',')

        anulus_r_in = None
        anulus_r_out = None
        if anulus_str and len(anulus_str) == 2:
            try:
                anulus_r1 = float(anulus_str[0])
                anulus_r2 = float(anulus_str[1])
            except ValueError:
                pass
            else:
                anulus_r_in = min(anulus_r1, anulus_r2)
                anulus_r_out = max(anulus_r1, anulus_r2)

        if args.debug:
            print(
                f"Read {len(sources[:n_objects])} sources from "
                "input catalog/regionfile...",
                file=sys.stderr
            )

        for i, source in enumerate(sources[:n_objects]):
            simple_pbar_callback(i + 1, n_objects)

            if key_id is not None:
                obj_id = source[key_id]
            else:
                obj_id = f"{i:06}"

            obj_ra = source[args.key_ra]
            obj_dec = source[args.key_dec]
            xx_tr = xx - source['CUBE_X_IMAGE']
            yy_tr = yy - source['CUBE_Y_IMAGE']

            anulus_mask = None
            if mode == 'regions':
                anulus_mask = None
                pix_reg = source['region'].to_pixel(celestial_wcs)
                reg_mask = pix_reg.to_mask()
                mask = reg_mask.to_image((img_height, img_width)) > 0
                specex_apertures[obj_id] = (
                    1 * pixelscale,
                    1 * pixelscale,
                    0 * units.rad
                )

            elif mode == 'kron_ellipse':
                ang = np.deg2rad(source[args.key_theta])

                x_over_a = xx_tr*np.cos(ang) + yy_tr*np.sin(ang)
                x_over_a /= source[args.key_a]

                y_over_b = xx_tr*np.sin(ang) - yy_tr*np.cos(ang)
                y_over_b /= source[args.key_b]

                specex_apertures[obj_id] = (
                    source[args.key_a] * pixelscale,
                    source[args.key_b] * pixelscale,
                    source[args.key_theta] * units.rad
                )

                mask = (x_over_a**2 + y_over_b**2) < (1.0/source[args.key_kron])

            elif mode == 'kron_circular':
                kron_circular = source[args.key_kron] * source[args.key_b]
                kron_circular /= source[args.key_a]

                specex_apertures[obj_id] = (
                    kron_circular * pixelscale,
                    kron_circular * pixelscale,
                    0 * units.rad
                )

                mask = (xx_tr**2 + yy_tr**2) < (kron_circular ** 2)

            elif mode == 'circular_aperture':
                aperture_size = units.Quantity(args.aperture_size)
                aperture_radius = aperture_size / 2
                if str(aperture_size.unit.physical_type) == 'dimensionless':
                    aperture_radius = aperture_radius * units.mas

                try:
                    specex_apertures[obj_id] = (
                        aperture_radius,
                        aperture_radius,
                        0 * units.rad
                    )
                except TypeError:
                    logging.warning(
                        f"object {obj_id} has an invalid ID, skipping..."
                    )
                    continue

                xx_scaled = xx_tr * cube_pixelscale_x
                yy_scaled = yy_tr * cube_pixelscale_y

                mask = (xx_scaled**2 + yy_scaled**2) < (aperture_radius)**2

                # TODO: fix here
                """
                if anulus_r_in and anulus_r_out:
                    specex_anuli[obj_id] = (
                        anulus_r_in,
                        anulus_r_out
                    )

                    anulus_mask_out = rad_coords <= anulus_r_out / cube_pixelscale
                    anulus_mask_in = rad_coords > anulus_r_in / cube_pixelscale

                    anulus_mask = anulus_mask_in & anulus_mask_out
                else:
                    specex_anuli[obj_id] = None
                    anulus_mask = None
                """

            if cube_footprint is not None:
                mask &= cube_footprint

            if np.sum(mask) == 0:
                logging.warning(
                    f"object {obj_id} has no footprint, skipping..."
                )
                continue

            obj_extraction_mask = np.zeros_like(mask, dtype='float32')
            if args.weighting.lower() == 'none':
                weights = None
                obj_extraction_mask[mask] = 1.0
            else:
                if args.weighting.lower() == 'whitelight':
                    # summing selected spaxels along spectral axis to get a
                    # whitelight weightmap
                    w_arr = np.ma.array(
                        my_cube.spec_hdu.data[:, mask],
                        mask=~np.isfinite(my_cube.spec_hdu.data[:, mask]),
                        dtype='float32'
                    )
                    weights = np.ma.mean(w_arr, axis=0)
                    weights -= np.ma.min(weights)
                elif args.weighting.lower() == 'log-whitelight':
                    w_arr = np.ma.array(
                        my_cube.spec_hdu.data[:, mask],
                        mask=~np.isfinite(my_cube.spec_hdu.data[:, mask]),
                        dtype='float32'
                    )
                    weights = np.ma.mean(w_arr, axis=0)
                    weights -= np.ma.min(weights)
                    weights = np.ma.log10(1 + weights)
                elif weight_img is not None:
                    weights = weight_img[mask]
                    weights -= np.ma.min(weights)

                obj_extraction_mask[mask] = weights

            var_spaxels = None
            if my_cube.var_hdu is not None:
                var_spaxels = my_cube.var_hdu.data[:, mask]

            obj_spectrum, obj_spectrum_var = get_spectrum(
                flux_spaxels=my_cube.spec_hdu.data[:, mask],
                var_spaxels=var_spaxels,
                weights=weights
            )

            obj_spectrum = obj_spectrum.filled(np.nan)
            obj_spectrum_var = obj_spectrum_var.filled(np.nan)

            if np.sum(
                np.isnan(obj_spectrum_var[~np.isnan(obj_spectrum)])
            ) != 0:
                logging.warning(
                    f"\nWARNING: object {obj_id}  is weird! Skipping...\n"
                )
                continue

            if np.sum(~np.isnan(obj_spectrum)) < 100:
                logging.warning(
                    f"WARNING: object {obj_id} has invalid data..."
                )

            if anulus_mask is not None:
                # TODO: fix here!
                # Here we suppose that background vaies very slowly within
                # object aperture. In this case we can copute the mean
                # background and subtract it total contribution within the
                # aperture.
                bkg_spectrum = np.median(
                    my_cube.spec_hdu.data[:, mask], axis=0
                )
                smoothed_bkg_spectrum = savgol_filter(bkg_spectrum, 51, 11)
                obj_spectrum -= smoothed_bkg_spectrum*np.sum(mask)

            if np.sum(~np.isnan(obj_spectrum)) == 0:
                if args.debug:
                    logging.warning(
                        f"WARNING: object {obj_id} has no spectral data, "
                        "skipping..."
                    )
                continue

            sn_spec = get_spectrum_snr(obj_spectrum, obj_spectrum_var)
            sn_em = get_spectrum_snr_emission(obj_spectrum, obj_spectrum_var)

            if (
                    (args.sn_threshold > 0) and
                    (max(sn_spec, sn_em) < args.sn_threshold)
               ):

                if args.debug:
                    logging.warning(
                        f"WARNING: object {obj_id} has SNR={sn_spec:.2f} < "
                        f"{args.sn_threshold}, skipping..."
                    )
                continue

            my_time = Time.now()
            my_time.format = 'isot'

            outname = f"spec_{obj_id}.fits"

            apertures_info = [
                x.to_string() for x in specex_apertures[obj_id]
            ]

            main_header = {
                'CUBE': (
                    basename_with_ext, "Spectral cube used for extraction"
                ),
                'RA': (
                    obj_ra, "Ra of the center of the object"
                ),
                'DEC': (
                    obj_dec, "Dec of the center of the object"
                ),
                'FRAME': (
                    coord_frame, "Ra and Dec sky frame"
                ),
                'NPIX': (
                    np.sum(mask), 'Number of pixels used for this spectra'
                ),
                'HISTORY': f'Extracted on {str(my_time)}',
                'ID': obj_id,
                'SN': (
                    sn_spec, "SNR of the total spectrum"
                ),
                'SN_EMISS': (
                    sn_em, "SNR due to emission lines only"
                ),
                'EXT_MODE': (
                    mode, "Spex extraction mode"
                ),
                'EXT_APER': (
                    json.dumps(apertures_info),
                    "Apertures used in specex"
                )
            }

            # Copy the spectral part of the WCS into the new FITS
            spec_wcs_header = spectral_wcs.to_header()

            hdul = get_spsingle_fits(
                main_header, spec_wcs_header,
                obj_spectrum, my_cube.spec_hdu.header,
                obj_spectrum_var, no_nans=args.no_nans
            )

            out_file_name = os.path.join(outdir, outname)
            hdul.writeto(
                out_file_name,
                overwrite=True
            )

            # If we arrived here, then the object has a vail spectrum and can
            # be saved. Let's also mark the object as valid and store its ID
            out_specfiles.append(os.path.realpath(out_file_name))
            source_ids.append(obj_id)
            valid_sources_mask[i] = True

            # Add also the extracted pixels to the extraction map
            extracted_data += obj_extraction_mask
            if anulus_mask is not None:
                extracted_data -= anulus_mask

        # Discard all invalid sources
        sources = sources[valid_sources_mask]
        source_ids = np.array(source_ids)

        with open("spec_list_pandora_ez.txt", 'w') as f:
            f.write(speclist2ez(out_specfiles))

        if args.debug:
            print(
                f"\nExtracted {len(sources)} valid sources\n",
                file=sys.stderr
            )

        if args.checkimg_outdir is not None:
            check_images_outdir = args.checkimg_outdir
        else:
            check_images_outdir = os.path.join(outdir, 'checkimages')

        if not os.path.isdir(check_images_outdir):
            os.mkdir(check_images_outdir)

        # NOTE: for some weird reason redrock hangs if any figure has been
        #       created using matplotlib. Do all the plottings only after
        #       redrock has finished!
        if args.zbest:
            rrspecex_options = [
                '--zbest', args.zbest,
                '--checkimg-outdir', check_images_outdir
            ]

            if args.priors is not None:
                rrspecex_options += ['--priors', args.priors]

            if args.nminima is not None:
                rrspecex_options += ['--nminima', f'{args.nminima:d}']

            if args.mp is not None:
                rrspecex_options += ['--mp', f'{args.mp:d}']

            if args.templates is not None:
                rrspecex_options += ['--templates', f'{args.templates}']

            rrspecex_options += out_specfiles
            targets, zfit, scandata = rrspecex(options=rrspecex_options)

            if args.debug or not args.cutouts_image:
                stacked_cube = stack(my_cube.spec_hdu.data)

            if args.debug:
                fig = plt.figure(figsize=(10, 10))
                ax = plt.subplot(projection=celestial_wcs)
                logimg, cvmin, cvmax = get_log_img(stacked_cube)
                ax.imshow(
                    logimg,
                    origin='lower',
                    vmin=cvmin,
                    vmax=cvmax,
                    cmap='jet'
                )
                fig.savefig("stacked_cube.png")
                plt.tight_layout()
                plt.close(fig)

            if args.debug:
                plot_templates = get_templates(templates=args.templates)
            else:
                plot_templates = None

            print("", file=sys.stderr)

            # Creating directory to save Chi-2 plots
            if args.debug:
                x2_plots_out_dir = os.path.join(
                    check_images_outdir,
                    "X2_plots"
                )
                if not os.path.isdir(x2_plots_out_dir):
                    os.mkdir(x2_plots_out_dir)

            for i, target in enumerate(targets):
                logging.info("Generating previews...")
                simple_pbar_callback(i, n_objects-1)

                fig, axs = plot_zfit_check(
                    target,
                    zfit,
                    plot_template=plot_templates,
                    wavelengt_units=my_cube.spec_hdu.header['CUNIT3'],
                    flux_units=my_cube.spec_hdu.header['BUNIT'],
                )

                figname = f'spectrum_{target.spec_id}.png'
                figname = os.path.join(check_images_outdir, figname)
                fig.savefig(figname, dpi=150)
                plt.close(fig)

                if args.debug:
                    figname = f'scandata_{target.spec_id}.png'
                    figname = os.path.join(x2_plots_out_dir, figname)
                    fig, axs = plot_scandata(target, scandata)
                    fig.savefig(figname, dpi=150)
                    plt.close(fig)

        specex_map = fits.PrimaryHDU(
            data=extracted_data,
            header=celestial_wcs.to_header()
        )
        specex_map.writeto(
            f"{my_cube.getBaseName()}_extmap.fits",
            overwrite=True
        )

        if args.check_images:
            fig, ax = plt.subplots(1, 1, figsize=(10, 10))
            ax.imshow(extracted_data, origin='lower')
            fig.savefig("sext_extraction_map.png", dpi=150)
            plt.close(fig)


if __name__ == '__main__':
    specex()
