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

redrock wrapper tools for python-specex spectra.
This program is based on the structure of redrock.external.boss function.

Copyright (C) 2022-2023  Maurizio D'Addona <mauritiusdadd@gmail.com>
"""

from __future__ import absolute_import, division, print_function

import os
import sys
import traceback
import argparse
from glob import glob
from packaging import version

import numpy as np
from scipy import sparse

from astropy.io import fits
from astropy.table import Table, join
import astropy.wcs as wcs
import matplotlib.pyplot as plt

from scipy.signal import savgol_filter

try:
    import redrock
    from redrock.utils import elapsed, get_mp
    from redrock.targets import Spectrum, Target, DistTargetsCopy
    from redrock.templates import load_dist_templates, find_templates, Template
    from redrock.results import write_zscan
    from redrock.zfind import zfind
    from redrock._version import __version__
    from redrock.archetypes import All_archetypes
except (ModuleNotFoundError, Exception):
    HAS_REDROCK = False
else:
    HAS_REDROCK = True

from specex.exceptions import exception_handler
from specex.utils import (
    plot_zfit_check, plot_scandata, get_pbar
)

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

MUSE_RESOLUTION_ANG = 2.51
RR_GPU_MIN_VER = "0.16.0"

def get_templates(template_types=[], filepath=False, templates=None):
    """
    Get avilable templates.

    Parameters
    ----------
    template_types : list of str, optional
        List of template types to retrieve. If it's empty all available
        templates will be returned.
        The default is [].
    filepath : boot, optional
        If it's true then return the file paths instead of actual templates.
    templates : str, optional
        The path of a template file or of a directory containing templates
        files. If None, templates are searched in the default redrock path.
        The default value is None.

    Returns
    -------
    available_templates : list of redrock.templates.Template or file paths
        The available templates or the corresponding file paths.

    """
    if templates is not None and os.path.isfile(templates):
        return [Template(templates), ]

    available_templates = []
    for t in find_templates(templates):
        templ = Template(t)
        if not template_types or templ.template_type in template_types:
            if filepath:
                available_templates.append(t)
            else:
                available_templates.append(templ)

    return available_templates


def get_template_types():
    """
    Get the available types of templates.

    Returns
    -------
    types : list of str
        List of types of available templates.

    """
    templates = [
        t.template_type
        for t in get_templates()
    ]
    types = set(templates)
    return types


def write_zbest(outfile, zbest, template_version, archetype_version):
    """
    Write zbest Table to outfile.

    Parameters
    ----------
        outfile : str
            The output file path.
        zbest : astropy.table.Table
            The output redshift fitting results.

    """
    header = fits.Header()
    header['RRVER'] = (__version__, 'Redrock version')

    for i, fulltype in enumerate(template_version.keys()):
        header['TEMNAM'+str(i).zfill(2)] = fulltype
        header['TEMVER'+str(i).zfill(2)] = template_version[fulltype]

    if archetype_version is not None:
        for i, fulltype in enumerate(archetype_version.keys()):
            header['ARCNAM'+str(i).zfill(2)] = fulltype
            header['ARCVER'+str(i).zfill(2)] = archetype_version[fulltype]

    zbest.meta['EXTNAME'] = 'ZBEST'

    hx = fits.HDUList()
    hx.append(fits.PrimaryHDU(header=header))
    hx.append(fits.convenience.table_to_hdu(zbest))
    hx.writeto(os.path.expandvars(outfile), overwrite=True)
    return


def read_spectra(spectra_fits_list, spec_hdu=None, var_hdu=None, wd_hdu=None,
                 memmap=True, resolution=2.51, smoothing=0, quite=False):
    """
    Read input spectra fits files.

    Parameters
    ----------
    spectra_fits_list : list
        List of fits files containing the input spectra.
    spec_hdu : int or None, optional
        The index of the HDU that contains the spectral data itself.
        If it is None, then the index is determined automatically by the name
        of the HDU. If this operation fails ValueError exception is raised.
        The default value is None.
    var_hdu : int or None, optional
        The index of the HDU that contains the  variance of the spectral data.
        If it is None, then the index is determined automatically by the name
        of the HDU. If this operation fails ValueError exception is raised.
        The default value is None.
    wd_hdu : int or None, optional
        The index of the HDU that contains the wavelength dispersion in pixels.
        If it is None, then the index is determined automatically by the name
        of the HDU. If this operation fails, no wavelenght dispersion will be
        used and the spectra will be considered having a uniform resolution.
        The default value is None.
    memmap : bool, optional
        Whether to use memmap or not. The default value is False.
    resolution : float, optional
        A fixed spectral resolution in Angrstrom to be used if no dispersion
        information is available from the spectrum file itself. The default
        value is 2.51 (the MUSE resolution).
    smoothing : int, optional
        A smooting to be applyed before redshift estimation.
        A value of 0 means no smoothing. The default value is 0.
    quite : bool, optional
        Reduce the verbosity of the output.

    Raises
    ------
    ValueError
        If cannot automatically determine the HDU containing the specral data
        or its variance.

    Returns
    -------
    targets : list of redrock.targets.Target
        The target spectra for which redshift will be computed.
    metatable : astropy.table.Table
        A table containing metadata.

    """
    targets = []
    targetids = []
    target_file_names = []
    specids = []
    sn_vals = []
    sn_var_vals = []
    sn_em_vals = []

    # crop to templates limits
    lmin = 3500.
    lmax = 10000.

    for j, fits_file in enumerate(spectra_fits_list):
        if not quite:
            progress = (j + 1) / len(spectra_fits_list)
            sys.stderr.write(f"\r{get_pbar(progress)} {progress:.2%}\r")
            sys.stderr.flush()
        hdul = fits.open(fits_file, memmap=memmap)

        valid_id_keys = [
            f"{i}{j}"
            for i in ['', 'OBJ', 'OBJ_', 'TARGET', 'TARGET_']
            for j in ['ID', 'NUMBER', 'UID', 'UUID']
        ]

        target_id = f"{j:09}"
        target_file_names.append(fits_file)
        spec_id = target_id
        for hdu in hdul:
            for key in valid_id_keys:
                try:
                    spec_id = hdu.header[key]
                except KeyError:
                    continue
                else:
                    break

        if spec_hdu is None:
            for hdu in hdul:
                if hdu.name.lower() in KNOWN_SPEC_EXT_NAMES:
                    flux = hdu.data
                    if smoothing > 0:
                        window_size = 4*smoothing + 1
                        flux = savgol_filter(flux, window_size, 3)
                    spec_header = hdu.header
                    spec_wcs = wcs.WCS(spec_header)
                    break
            else:
                raise ValueError(
                    "Cannot determine the HDU containing spectral data: "
                    f"'{fits_file}'"
                )
        else:
            flux = hdul[spec_hdu].data
            spec_wcs = wcs.WCS(hdul[spec_hdu].header)

        for hdu in hdul:
            if hdu.name.lower() in KNOWN_MASK_EXT_NAMES:
                nanmask = hdu.data.astype(bool)
                break
        else:
            nanmask = None

        if var_hdu is None:
            for hdu in hdul:
                if hdu.name.lower() in KNOWN_VARIANCE_EXT_NAMES:
                    ivar = 1 / hdu.data
                    break
                elif hdu.name.lower() in KNOWN_INVAR_EXT_NAMES:
                    ivar = hdu.data
                    break
            else:
                print(
                    "WARNING: Cannot determine the HDU containing variance "
                    f"data in '{fits_file}'! Using dumb constan variance...",
                )
                ivar = np.ones_like(flux)
        else:
            ivar = 1 / hdul[var_hdu].data

        if wd_hdu is None:
            for hdu in hdul:
                if hdu.name.lower() in KNOWN_RCURVE_EXT_NAMES:
                    wd = hdu.data
                    break
            else:
                wd = None
        else:
            wd = hdul[wd_hdu].data

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

        main_header = hdul[0].header

        # NOTE: Wavelenghts must be in Angstrom units
        pixel = np.arange(len(flux))
        if spec_wcs.has_spectral:
            lam = spec_wcs.pixel_to_world(pixel).Angstrom
        else:
            try:
                coeff0 = spec_header["COEFF0"]
                coeff1 = spec_header["COEFF1"]
            except KeyError:
                continue
            lam = 10**(coeff0 + coeff1*pixel)
        flux = flux.astype('float32')

        if nanmask is None:
            flux_not_nan_mask = ~np.isnan(flux)
        else:
            flux_not_nan_mask = ~nanmask

        flux = flux[flux_not_nan_mask]
        ivar = ivar[flux_not_nan_mask]
        lam = lam[flux_not_nan_mask]
        if wd is not None:
            wd = wd[flux_not_nan_mask]
        else:
            # If now wavelenght dispersion information is present, then
            # compute it using the wavelenght
            delta_lambda = np.ones_like(lam)
            delta_lambda[1:] = (lam[1:] - lam[:-1])
            wd = resolution / delta_lambda
            wd[0] = wd[1]
        wd[wd < 1e-3] = 2.

        imin = abs(lam-lmin).argmin()
        imax = abs(lam-lmax).argmin()

        lam = lam[imin:imax]
        flux = flux[imin:imax]
        ivar = ivar[imin:imax]
        wd = wd[imin:imax]

        ndiag = int(4*np.ceil(wd.max())+1)
        nbins = wd.shape[0]

        ii = np.arange(lam.shape[0])
        di = ii-ii[:, None]
        di2 = di**2

        # build resolution from wdisp
        reso = np.zeros([ndiag, nbins])

        for idiag in range(ndiag):
            offset = ndiag//2-idiag
            d = np.diagonal(di2, offset=offset)
            if offset < 0:
                reso[idiag, :len(d)] = np.exp(-d/2/wd[:len(d)]**2)
            else:
                reso[idiag, nbins-len(d):nbins] = np.exp(
                    -d/2/wd[nbins-len(d):nbins]**2
                )

        reso /= np.sum(reso, axis=0)
        offsets = ndiag//2 - np.arange(ndiag)
        nwave = reso.shape[1]
        R = sparse.dia_matrix((reso, offsets), (nwave, nwave))

        try:
            s_n = main_header['SN']
        except KeyError:
            s_n = -1

        try:
            s_n_var = main_header['SN_VAR']
        except KeyError:
            s_n_var = -1

        try:
            s_n_em = main_header['SN_EMISS']
        except KeyError:
            s_n_em = -1

        rrspec = Spectrum(lam, flux, ivar, R, None)
        target = Target(target_id, [rrspec])
        target.input_file = fits_file
        target.spec_id = spec_id
        targets.append(target)
        targetids.append(target_id)
        specids.append(spec_id)
        sn_vals.append(s_n)
        sn_var_vals.append(s_n_var)
        sn_em_vals.append(s_n_em)

    metatable = Table()
    metatable['TARGETID'] = targetids
    metatable['SPECID'] = specids
    metatable['FILE'] = target_file_names
    metatable['SN'] = sn_vals
    metatable['SN_VAR'] = sn_var_vals
    metatable['SN_EMISS'] = sn_em_vals

    if not quite:
        print("", file=sys.stderr)

    return targets, metatable


def __argshandler(options=None):
    """
    Handle input arguments.

    Parameters
    ----------
    options : TYPE, optional
        DESCRIPTION. The default is None.

    Returns
    -------
    args : TYPE
        DESCRIPTION.

    """
    parser = argparse.ArgumentParser(
        description="Estimate redshifts for spectra extracted with"
        "python-specex using redrock interface."
    )

    parser.add_argument(
        "spectra", metavar='SPECTRA_FITS', type=str, nargs='+',
        help="input spectra fits files"
    )

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

    parser.add_argument(
        "--archetypes", type=str, default=None, required=False,
        help="archetype file or directory for final redshift comparisons"
    )

    parser.add_argument(
        "-o", "--output", type=str, default=None, required=False,
        help="output file"
    )

    parser.add_argument(
        "--zbest", type=str, default=None, required=False,
        help="output zbest FITS file"
    )

    parser.add_argument(
        "--gpu", action="store_true", default=False,
        help="Use GPU computing acceleration (requires redrock >= "
        f"{RR_GPU_MIN_VER})."
    )

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

    parser.add_argument(
        "--chi2-scan", type=str, default=None, required=False,
        help="Load the chi2-scan from the input file"
    )

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

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

    parser.add_argument(
        "--debug", default=False, action="store_true", required=False,
        help="debug with ipython (only if communicator has a single process)"
    )

    parser.add_argument(
        "--plot-zfit", action="store_true", default=False, required=False,
        help="Generate plots of the spectra with infomrazion about the "
        "result of the template fitting (i.e. the best redshift, the position "
        "of most important lines, the best matching template, etc...)."
    )

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

    parser.add_argument(
        "--quite", action='store_true', default=False,
        help="Reduce program output at bare minimum. Do not print fitting "
        "result,"
    )

    parser.add_argument(
        "--no-memmap", action='store_true', default=False,
        help='Disable memmapping to speed the computation at a cost of an '
        'increased memory usage.'
    )

    parser.add_argument(
        "--resolution", type=float, default=2.51, help="A fixed resolution of "
        "the spectrograph to be used if no resolution curve is present in the "
        "input spectrum itself."
    )

    parser.add_argument(
        '--smoothing', metavar='WINDOW_SIZE', type=int,  default=0,
        help='If %(metavar)s >= 0, then the spectrua are smoothed before the '
        'redshift estimation procedure is run.  %(metavar)s = 0 means that '
        'no smoothing is applyed. If not specified, the default '
        '%(metavar)s = %(default)s is used.'
    )

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

    return args


def rrspecex(options=None, comm=None):
    """
    Estimate redshifts for spectra extracted with python-specex using redrock.

    This loads targets serially and copies them into a DistTargets class.
    It then runs redshift fitting and writes the output to a catalog.

    Parameters
    ----------
    options : list, optional
        lLst of commandline options to parse. The default is None.
    comm : mpi4py.Comm, optional
        MPI communicator to use. The default is None.

    Returns
    -------
    targets : list of redrock.targets.Target objects
        list of target spectra.
    zfit : astropy.table.Table
        Table containing the fit results.
    scandata : dict
        A dictionary containing the redshift scanning information for each
        target
    """
    global HAS_REDROCK
    global HAS_IPYTHON

    global_start = elapsed(None, "", comm=comm)
    comm_size = 1
    comm_rank = 0
    if comm is not None:
        comm_size = comm.size
        comm_rank = comm.rank

    args = __argshandler(options)

    if args.debug:
        sys.excepthook = exception_handler

    # Check arguments - all processes have this, so just check on the first
    # process
    if comm_rank == 0:
        if args.debug and comm_size != 1:
            print(
                "--debug can only be used if the communicator has one process"
            )
            sys.stdout.flush()
            if comm is not None:
                comm.Abort()

        if (args.output is None) and (args.zbest is None):
            print("--output or --zbest required")
            sys.stdout.flush()
            if comm is not None:
                comm.Abort()

    # Multiprocessing processes to use if MPI is disabled.
    mpprocs = 0
    if comm is None:
        mpprocs = get_mp(args.mp)
        print("Running with {} processes".format(mpprocs))
        if "OMP_NUM_THREADS" in os.environ:
            nthread = int(os.environ["OMP_NUM_THREADS"])
            if nthread != 1:
                print("WARNING:  {} multiprocesses running, each with "
                      "{} threads ({} total)".format(
                          mpprocs, nthread, mpprocs*nthread
                      ))
                print("WARNING:  Please ensure this is <= the number of "
                      "physical cores on the system")
        else:
            print("WARNING:  using multiprocessing, but the OMP_NUM_THREADS")
            print("WARNING:  environment variable is not set- your system may")
            print("WARNING:  be oversubscribed.")
        sys.stdout.flush()
    elif comm_rank == 0:
        print("Running with {} processes".format(comm_size))
        sys.stdout.flush()

    try:
        # Load and distribute the targets
        if comm_rank == 0:
            print("Loading targets...")
            sys.stdout.flush()

        start = elapsed(None, "", comm=comm)

        # Read the spectra on the root process.  Currently the "meta" Table
        # returned here is not propagated to the output zbest file.  However,
        # that could be changed to work like the DESI write_zbest() function.
        # Each target contains metadata which is propagated to the output zbest
        # table though.

        # Windows prompt does not expand globs, so let's do it
        spectra_list = []
        for globbed_fname in args.spectra:
            for fname in glob(globbed_fname):
                spectra_list.append(fname)

        targets, meta = read_spectra(
            spectra_list,
            memmap=not args.no_memmap,
            resolution=args.resolution,
            smoothing=args.smoothing,
            quite=args.quite
        )

        if len(targets) == 0:
            raise ValueError("No spectra were loaded!")
        else:
            print(f"Loaded {len(targets)} spectra")

        _ = elapsed(
            start, "Read of {} targets".format(len(targets)), comm=comm
        )

        # Distribute the targets.

        start = elapsed(None, "", comm=comm)

        dtargets = DistTargetsCopy(targets, comm=comm, root=0)

        # Get the dictionary of wavelength grids
        dwave = dtargets.wavegrids()

        _ = elapsed(
            start,
            "Distribution of {} targets".format(len(dtargets.all_target_ids)),
            comm=comm
        )

        opt_zfind_args = {}
        opt_load_dist_templates_args = {}
        if version.parse(redrock.__version__) >= version.parse(RR_GPU_MIN_VER):
            opt_zfind_args['use_gpu'] = args.gpu
            opt_load_dist_templates_args['use_gpu'] = args.gpu

        # Read the template data
        dtemplates = load_dist_templates(
            dwave,
            templates=args.templates,
            comm=comm,
            mp_procs=mpprocs,
            **opt_load_dist_templates_args
        )

        # Compute the redshifts, including both the coarse scan and the
        # refinement.  This function only returns data on the rank 0 process.
        start = elapsed(None, "", comm=comm)
        scandata, zfit = zfind(
            dtargets,
            dtemplates,
            mpprocs,
            nminima=args.nminima,
            archetypes=args.archetypes,
            priors=args.priors,
            chi2_scan=args.chi2_scan,
            **opt_zfind_args
        )

        _ = elapsed(start, "Computing redshifts took", comm=comm)

        # Write the outputs
        if args.output is not None:
            start = elapsed(None, "", comm=comm)
            if comm_rank == 0:
                write_zscan(args.output, scandata, zfit, clobber=True)
            _ = elapsed(start, "Writing zscan data took", comm=comm)

        # Change to upper case like DESI
        for colname in zfit.colnames:
            if colname.islower():
                zfit.rename_column(colname, colname.upper())

        matched_zfit = join(zfit, meta, keys=['TARGETID'], join_type='left')

        zbest = None
        if comm_rank == 0:
            zbest = matched_zfit[matched_zfit['ZNUM'] == 0]
            zbest.add_index('TARGETID')

            if args.plot_zfit:

                if not os.path.isdir(args.checkimg_outdir):
                    os.makedirs(args.checkimg_outdir)

                available_templates = get_templates(
                    templates=args.templates
                )
                for target in targets:
                    orig_file_name = os.path.basename(
                        zbest.loc[target.id]['FILE']
                    )

                    fig, ax = plot_zfit_check(
                        target,
                        zbest,
                        plot_template=available_templates
                    )
                    figname = f'rrspecex_{orig_file_name}.png'
                    if args.checkimg_outdir is not None:
                        figname = os.path.join(args.checkimg_outdir, figname)
                    fig.savefig(figname, dpi=150)
                    plt.close(fig)

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

            zbest.remove_column('TARGETID')

            if args.zbest:
                start = elapsed(None, "", comm=comm)
                if comm_rank == 0:

                    # Remove extra columns not needed for zbest
                    # zbest.remove_columns(['zz', 'zzchi2', 'znum'])
                    # zbest.remove_columns(['ZNUM'])

                    template_version = {
                        t._template.full_type: t._template._version
                        for t in dtemplates
                    }

                    archetype_version = None
                    if args.archetypes is not None:
                        archetypes = All_archetypes(
                            archetypes_dir=args.archetypes
                        ).archetypes
                        archetype_version = {
                            name: arch._version
                            for name, arch in archetypes.items()
                        }

                    write_zbest(
                        args.zbest, zbest, template_version, archetype_version
                    )

                _ = elapsed(start, "Writing zbest data took", comm=comm)

            if (not args.zbest) and (args.output is None) and (not args.quite):
                print("")
                print(zbest)
                print("")

    except Exception:
        exc_type, exc_value, exc_traceback = sys.exc_info()
        lines = traceback.format_exception(exc_type, exc_value, exc_traceback)
        lines = [
            "Proc {}: {}".format(comm_rank, x)
            for x in lines
        ]
        print("".join(lines))
        sys.stdout.flush()
        if comm is not None:
            comm.Abort()

    _ = elapsed(global_start, "Total run time", comm=comm)

    return targets, zbest, scandata


if __name__ == '__main__':
    _ = rrspecex()
