import numpy as np
from scipy.interpolate import RegularGridInterpolator
from astropy.io import fits
from functools import reduce
import warnings

from pynot.functions import get_version_number
from pynot.fitsio import load_fits_spectrum, save_fitstable_spectrum, save_fits_spectrum, error_HDU_names
from pynot.txtio import load_ascii_spectrum


__version__ = get_version_number()


def combine_2d(files, output=None, method='mean', scale=False, extended=False, dispaxis=1, trim=False):
    """Combine a list of 2d-spectra using either median or mean combination.
    For median combination, only the overlapping parts of the spectra will be
    combined. The mean combination uses a weighted average over the entire
    spectral range covered.
    The spectra are all assumed to have the same dispersion axis.

    Input
    =====
    files : list of strings
        Input list of files to combine

    output : string
        Output filename, if the filename will be autogenerated from the OBJECT keyword in the header

    method : string
        Combination method, either 'median' or 'mean'

    dispaxis : int   [default=1]
        Dispersion axis. 1: horizontal spectra, 2: vertical spectra (default for most ALFOSC grisms)

    overwrite : bool   [default=False]
        Overwrite existing files.

    """
    wl = list()
    space = list()
    flux = list()
    err = list()
    mask = list()
    size = list()
    scales = list()
    centers = list()
    exp_times = list()
    meta_string = list()

    msg = list()
    # print("\n Spectral Combination 2D\n")
    for fnum, fname in enumerate(files):
        msg.append("          - Loading file: %s" % fname)
        hdr = fits.getheader(fname)
        data2D = fits.getdata(fname)

        for err_name in error_HDU_names:
            try:
                err2D = fits.getdata(fname, err_name)
                err2D[err2D <= 0.] = np.nanmedian(err2D)*100
                msg.append(f"          - Loaded error image extension: {err_name}")
                break
            except KeyError:
                pass
        else:
            msg.append("[WARNING] - No ERR extension could be found in the FITS file!")
            err2D = np.ones_like(data2D)

        for mask_name in ['MASK', 'QUAL']:
            try:
                mask2D = fits.getdata(fname, mask_name)
                msg.append(f"          - Loaded mask image extension: {mask_name}")
                break
            except KeyError:
                pass
        else:
            msg.append("[WARNING] - No MASK extension could be found in the FITS file!")
            mask2D = np.zeros_like(data2D)

        if 'CD1_1' in hdr:
            cdelt1 = hdr['CD1_1']
        elif 'CDELT1' in hdr:
            cdelt1 = hdr['CDELT1']
        else:
            msg.append(" [ERROR]  - The pixel sampling could not be found!")
            msg.append(" [ERROR]  - Neither CD1_1 nor CDELT1 is in the header keywords...")
            raise KeyError('CD1_1 or CDELT1 missing in header')
        crval1 = hdr['CRVAL1']
        crpix1 = hdr['CRPIX1']
        naxis1 = hdr['NAXIS1']

        if 'CD2_2' in hdr:
            cdelt2 = hdr['CD2_2']
        elif 'CDELT2' in hdr:
            cdelt2 = hdr['CDELT2']
        else:
            msg.append(" [ERROR]  - The pixel sampling could not be found!")
            msg.append(" [ERROR]  - Neither CD2_2 nor CDELT2 is in the header keywords...")
            raise KeyError('CD2_2 or CDELT2 missing in header')
        crval2 = hdr['CRVAL2']
        crpix2 = hdr['CRPIX2']
        naxis2 = hdr['NAXIS2']

        if 'EXPTIME' in hdr:
            exptime = hdr['EXPTIME']
            exp_times.append(exptime)
        else:
            msg.append("[WARNING] - No exposure time could be found in the header!")
            exptime = 0
            exp_times.append(1)

        if 'RESPONSE' in hdr:
            # Data have been flux calibrated, so combination is straight-forwards
            weight_by_t = False
        elif 'BUNIT' in hdr and hdr['BUNIT'].lower() in ['count', 'counts', 'adu', 'adus']:
            msg.append("[WARNING] - Exposures have not been flux calibrated.")
            msg.append("            The combination is done in counts/sec and scaled to the total exposure time.")
            weight_by_t = True
        else:
            msg.append("[WARNING] - Could not determine if the spectra are in counts or flux units")
            msg.append("            I assume they are in counts per second or similar units... Please verify the result!")
            msg.append("            If the data are in counts, please set the header keyword BUNIT to ADU or COUNTS")
            weight_by_t = False

        if dispaxis == 2:
            # Vertical spectra:
            # most useful for raw spectra before processing
            spatial = (np.arange(naxis1) - (crpix1 - 1))*cdelt1 + crval1
            wavelength = (np.arange(naxis2) - (crpix2 - 1))*cdelt2 + crval2
            N_pix = len(spatial)
            lower_bound = None
            upper_bound = None
            SPSF = np.nanmedian(data2D[lower_bound:upper_bound, :], axis=0)

        elif dispaxis == 1:
            # Horizontal spectra:
            # default format after processing by PyNOT
            wavelength = (np.arange(naxis1) - (crpix1 - 1))*cdelt1 + crval1
            spatial = (np.arange(naxis2) - (crpix2 - 1))*cdelt2 + crval2
            N_pix = len(spatial)
            lower_bound = None
            upper_bound = None
            SPSF = np.nanmedian(data2D[:, lower_bound:upper_bound], axis=1)
        else:
            msg.append(" [ERROR]  - Invalid dispaxis!")
            raise ValueError("`dispaxis` must be either 1 or 2 (horizontal or vertical)")

        wl.append(wavelength)
        flux.append(data2D)
        err.append(err2D)
        mask.append(mask2D)
        size.append(data2D.shape)

        if not extended:
            # --- Measure the center and width of the trace
            SPSF[SPSF < 0] = 0.
            # Trim edges of SPSF:
            SPSF[:3] = 0.
            SPSF[-3:] = 0.
            # First guess of trace region based on pixel with max-value +/- 5 pixels:
            y_pix = np.arange(N_pix)
            y0 = np.argmax(SPSF)
            trace_region = (y_pix > y0-10) * (y_pix < y0+10)

            # centroid:
            trace_cen = np.sum((spatial*SPSF)[trace_region]) / np.sum(SPSF[trace_region])

            # width:
            sumY2 = np.sum(((spatial-trace_cen)**2 * SPSF)[trace_region])
            sumP = np.sum(SPSF[trace_region])
            trace_fwhm = np.sqrt(sumY2/sumP)*2.35

            # Recenter trace:
            space.append(spatial - trace_cen)
            centers.append(trace_cen)

        else:
            trace_cen = N_pix // 2
            trace_fwhm = 0

            space.append(spatial)
            centers.append(trace_cen)


        if scale:
            this_scale = np.nanmax(SPSF)
        else:
            this_scale = 1.

        if weight_by_t and exptime > 0:
            this_scale = exptime
        scales.append(this_scale)

        meta_data = (fnum+1, trace_cen, trace_fwhm) + data2D.shape + (this_scale, exptime)
        meta_string.append(7*' ' + '       %3i        %+2.1f"      %.2f"    (%i, %i)   %.2e    %.1f' % meta_data)


    msg.append("\n          - IMAGE STATISTICS:")
    msg.append("          --------------------------------------------------------------")
    msg.append("          Image No    Centroid      FWHM       Shape       Flux Scale    EXPTIME")
    for line in meta_string:
        msg.append(line)

    # Determine new wavelength grid:
    # If differences are smaller than one pixel then do not interpolate
    l_min = np.min(wl, 1)
    l_max = np.max(wl, 1)
    diff_lmin = np.max(np.abs(np.diff(l_min)))
    diff_lmax = np.max(np.abs(np.diff(l_max)))
    pix_size1 = np.mean([np.diff(X)[0] for X in wl])

    cen_min = np.min(centers)
    cen_max = np.max(centers)
    pix_size2 = np.mean([np.diff(Y)[0] for Y in space])

    # Rescale to flux values of the order 10^0:
    rescale = 10**int(np.log10(np.max(SPSF)))
    if weight_by_t:
        f0 = 1 / rescale
    else:
        f0 = np.mean(scales) / rescale
    msg.append("          - Rescaling data by a factor of 10^%i" % (-int(np.log10(np.max(SPSF)))))

    # Check whether to use interpolation or not:
    same_size_x = all(x == size[0][0] for x, y in size)
    same_size_y = all(y == size[0][1] for x, y in size)
    array_shape_x = ((diff_lmax < pix_size1) and (diff_lmin < pix_size1))
    array_centroids = (cen_max - cen_min < pix_size2)
    use_interpolation = not (same_size_x & same_size_y & array_shape_x & array_centroids)
    if trim:
        use_interpolation = True

    if not use_interpolation:
        msg.append("          - Combining with no interpolation!")
        final_wl = wl[0]
        final_space = space[0]
        int_flux = list()
        int_var = list()
        int_mask = list()
        weight = list()
        for f, e, s, m in zip(flux, err, scales, mask):
            e = e/s*f0
            f = f/s*f0
            int_flux.append(f)
            int_var.append(e**2)
            int_mask.append(m)
            weight.append(1./e**2)

    else:
        msg.append("          - Interpolating data onto new grid!")
        int_flux = list()
        int_var = list()
        int_mask = list()
        weight = list()
        if method == 'median' or trim:
            xmin = np.max([np.min(ar) for ar in wl])
            xmax = np.min([np.max(ar) for ar in wl])
            ymin = np.max([np.min(ar) for ar in space])
            ymax = np.min([np.max(ar) for ar in space])

        elif method == 'mean':
            overlap_wl = reduce(np.union1d, wl)
            xmin = overlap_wl.min()
            xmax = overlap_wl.max()
            overlap_space = reduce(np.union1d, space)
            ymin = overlap_space.min()
            ymax = overlap_space.max()

        pix = np.mean([np.diff(X)[0] for X in wl])
        pix_y = np.mean([np.diff(Y)[0] for Y in space])
        final_wl = np.arange(xmin, xmax, pix)
        final_space = np.arange(ymin, ymax, pix_y)

        msg.append("          - Creating new linear wavelength grid:")
        msg.append("          - %.1f -- %.1f  |  sampling: %.1f" % (final_wl.min(), final_wl.max(), pix))

        # interpolate
        for this_wl, y, f, e, s, m in zip(wl, space, flux, err, scales, mask):
            e = e/s*f0
            f = f/s*f0
            new_points = tuple(np.meshgrid(final_space, final_wl, indexing='ij'))
            f_i = RegularGridInterpolator((y, this_wl), f, bounds_error=False)(new_points)
            e_i = RegularGridInterpolator((y, this_wl), e, bounds_error=False)(new_points)
            m_i = RegularGridInterpolator((y, this_wl), m, bounds_error=False)(new_points)
            w_i = RegularGridInterpolator((y, this_wl), 1/e**2, bounds_error=False)(new_points)
            int_flux.append(f_i)
            int_var.append(e_i**2)
            int_mask.append(m_i)
            weight.append(w_i)

    # Combine spectra:
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        if method == 'median':
            msg.append("          - Combination method: median")
            hdr['COMBINE'] = "Median"
            final_flux = np.nanmedian(int_flux, 0)
            final_var = np.nansum(int_var, 0)/len(int_flux)
            final_err = np.sqrt(final_var)
            final_mask = np.nansum(int_mask, 0) > 0

        elif method == 'mean':
            msg.append("          - Combination method: inverse variance weighting")
            hdr['COMBINE'] = "Inverse Variance Weighted"
            int_flux = np.array(int_flux)
            weight = np.array(weight)
            final_flux = np.nansum(int_flux*weight, 0) / np.nansum(weight, 0)
            final_var = 1. / np.nansum(weight, 0)
            final_err = np.sqrt(final_var)
            final_mask = np.nansum(int_mask, 0) > 0


    if dispaxis == 1:
        hdr['CRVAL1'] = np.min(final_wl)
        hdr['CD1_1'] = np.diff(final_wl)[0]
        hdr['CRVAL2'] = np.min(final_space)
        hdr['CD2_2'] = np.diff(final_space)[0]
    elif dispaxis == 2:
        hdr['CRVAL2'] = np.min(final_wl)
        hdr['CD2_2'] = np.diff(final_wl)[0]
        hdr['CRVAL1'] = np.min(final_space)
        hdr['CD1_1'] = np.diff(final_space)[0]

    hdr['NCOMBINE'] = len(int_flux)
    hdr['EXPTIME'] = (np.sum(exp_times), "Sum of integration time for all exposures")
    hdr['DATAMIN'] = np.nanmin(final_flux*rescale)
    hdr['DATAMAX'] = np.nanmax(final_flux*rescale)
    hdr['EXTNAME'] = 'DATA'
    hdr['AUTHOR'] = 'PyNOT version %s' % __version__

    if output is None:
        object_name = hdr['OBJECT']
        output = '%s_combined_2d.fits' % object_name
    else:
        if output[-4:].lower() != 'fits':
            output += '.fits'
    msg.append("          - Bringing data back to original flux scale")

    if weight_by_t:
        msg.append("          - Scaling the counts to the effective exposure time: %.1f sec" % np.sum(exp_times))
        rescale = rescale * np.sum(exp_times)
    FLUX = fits.PrimaryHDU(data=final_flux*rescale, header=hdr)
    FLUX.name = 'DATA'
    ERR = fits.ImageHDU(data=final_err*rescale, header=hdr, name='ERR')
    MASK = fits.ImageHDU(data=final_mask.astype(int), header=hdr, name='MASK')
    HDU = fits.HDUList([FLUX, ERR, MASK])
    HDU.writeto(output, overwrite=True)
    msg.append(" [OUTPUT] - Saving the combined spectrum to file: %s" % output)
    output_msg = "\n".join(msg)

    return (final_wl, final_flux, final_err, final_mask, output_msg)



def combine_1d(files, output=None, method='mean', scale=False, table_output=True):
    """Combine a list of 1d-spectra using either median or mean combination.
    For median combination, only the overlapping parts of the spectra will be
    combined. The mean combination uses a weighted average over the entire
    spectral range covered. If the spectra have different wavelength coverage,
    they will be interpolated onto a common grid.

    Input
    =====
    files : list of strings
        Input list of files to combine

    output : string   [default=None]
        Output filename, if None then the filename will be autogenerated

    method : string   [default=mean]
        Combination method, either 'median' or 'mean'

    scale : bool   [default=False]
        Scale the individual spectra to their mean values?

    table_output : bool   [default=True]
        Use FITS table for the output format? Otherwise use a MultiExtension Fits File
    """

    wl_all = list()
    flux_all = list()
    err_all = list()
    mask_all = list()
    size_all = list()
    scales = list()

    msg = list()
    for fname in files:
        if fname.endswith('.fits') or fname.endswith('.fit'):
            # load FITS file.
            wl, flux, err, mask, hdr, load_msg = load_fits_spectrum(fname)
            if load_msg:
                msg.append(load_msg)
            msg.append("          - Loaded FITS spectrum: %s" % fname)

        else:
            wl, flux, err, mask, _, load_msg = load_ascii_spectrum(fname)
            hdr = fits.Header()
            if load_msg:
                msg.append(load_msg)
            msg.append("          - Loaded ASCII spectrum: %s" % fname)

        wl_all.append(wl)
        flux_all.append(flux)
        err_all.append(err)
        mask_all.append(mask)
        size_all.append(len(wl))
        if scale:
            nonzero = flux.nonzero()[0]
            idx_0 = min(nonzero) + len(nonzero)/2
            x1 = idx_0 - 50
            x2 = idx_0 + 51
            scales.append(np.nanmedian(flux[x1:x2]))
        else:
            scales.append(1.)

    # Determine average reference scale:
    f0 = np.nanmean(scales)

    # -- Determine new wavelength grid:
    # If differences are smaller than one pixel then do not interpolate
    same_size = all(x == size_all[0] for x in size_all)
    wl_min = [min(this_wl) for this_wl in wl_all]
    wl_max = [max(this_wl) for this_wl in wl_all]
    wl_diff_lower = np.max(np.abs(np.diff(wl_min)))
    wl_diff_upper = np.max(np.abs(np.diff(wl_max)))
    avg_pix_size = np.mean([np.mean(np.diff(this_wl)) for this_wl in wl_all])
    if same_size and (wl_diff_lower < avg_pix_size) and (wl_diff_upper < avg_pix_size):
        msg.append("          - Combining with no interpolation!")
        final_wl = wl_all[0]
        int_flux = list()
        int_var = list()
        int_mask = mask_all
        weight = list()
        for f, e, m, s in zip(flux_all, err_all, mask_all, scales):
            v = (e/s*f0)**2
            f = f/s*f0
            int_flux.append(f)
            int_var.append(v)
            weight.append(1./v)

    else:
        msg.append("          - Interpolating onto common wavelength grid!")
        int_flux = list()
        int_var = list()
        int_mask = list()
        weight = list()
        if method == 'median':
            xmin = np.max([np.min(ar) for ar in wl_all])
            xmax = np.min([np.max(ar) for ar in wl_all])

        elif method == 'mean':
            overlap = reduce(np.union1d, wl_all)
            xmin = overlap.min()
            xmax = overlap.max()

        pix = np.min([np.min(np.diff(this_wl)) for this_wl in wl_all])
        final_wl = np.arange(xmin, xmax, pix)
        msg.append("          - Creating new linear wavelength grid:")
        msg.append("          - %.3f -- %.3f  |  sampling: %.3f" % (final_wl.min(), final_wl.max(), pix))
        hdr['CRVAL1'] = final_wl.min()
        hdr['CRPIX1'] = 1
        hdr['CDELT1'] = pix
        hdr['CD1_1'] = pix
        for key in ['CD1_2', 'CD2_1', 'CD2_2', 'CDELT2', 'CRVAL2', 'CRPIX2', 'CTYPE2', 'CUNIT2']:
            if key in hdr:
                hdr.remove(key)

        # interpolate
        for this_wl, f, e, m, s in zip(wl_all, flux_all, err_all, mask_all, scales):
            m[e <= 0] = 1
            e[e <= 0] = 1.e4
            v = (e*s/f0)**2
            f = f*s/f0
            f_i = np.interp(final_wl, this_wl, f, left=np.nan, right=np.nan)
            v_i = np.interp(final_wl, this_wl, v, left=np.nan, right=np.nan)
            m_i = np.interp(final_wl, this_wl, m, left=np.nan, right=np.nan) > 0
            w_i = np.interp(final_wl, this_wl, 1./v, left=np.nan, right=np.nan)
            int_flux.append(f_i)
            int_var.append(v_i)
            int_mask.append(m_i)
            weight.append(w_i)

    # Combine spectra:
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        if method == 'median':
            msg.append("          - Combination method: median")
            final_flux = np.nanmedian(int_flux, axis=0)
            final_var = np.nansum(int_var, axis=0)/len(int_flux)
            final_err = np.sqrt(final_var)
            final_mask = 1 * (np.nansum(int_mask, axis=0) > 0)

        else:
            msg.append("          - Combination method: inverse variance weighting")
            int_flux = np.array(int_flux)
            M = ~np.array(int_mask, dtype=bool)
            weight = np.array(weight)
            final_flux = np.nansum(M*int_flux*weight, 0) / np.nansum(M*weight, 0)
            final_var = 1. / np.nansum(M*weight, 0)
            final_err = np.sqrt(final_var)
            final_mask = 1 * (np.nansum(int_mask, axis=0) > 0)

    if output is None:
        obj_name = hdr.get('OBJECT')
        if not obj_name:
            fname = 'combined_spectrum_1d.fits'
        fname = '%s_combined.fits' % obj_name
    elif not output.endswith('.fits'):
        fname = output + '.fits'
    else:
        fname = output

    if table_output:
        save_fitstable_spectrum(fname, final_wl, final_flux, final_err, hdr, mask=final_mask)
        msg.append(" [OUTPUT] - Saving the combined spectrum as a FITS table: %s" % output)
    else:
        save_fits_spectrum(fname, final_wl, final_flux, final_err, hdr, mask=final_mask)
        msg.append(" [OUTPUT] - Saving the combined spectrum as a MultiExtension FITS file: %s" % output)

    output_msg = "\n".join(msg)

    return (final_wl, final_flux, final_err, final_mask, output_msg)
