from __future__ import print_function, division, absolute_import

import numpy as np


def linear_interp_resampler(wave_new, wave, flux, err=None):
    """
    Call interpolation, repackage new spectra


    Parameters
    ----------
    orig_spectrum : `~specutils.Spectrum1D`
        The original 1D spectrum.
    fin_spec_axis : ndarray
        The desired spectral axis array.

    Returns
    -------
    resample_spectrum : `~specutils.Spectrum1D`
        An output spectrum containing the resampled `~specutils.Spectrum1D`
    """

    fill_val = np.nan
    flux_new = np.interp(wave_new, wave, flux, left=fill_val, right=fill_val)

    new_unc = None
    if err is not None:
        err_new = np.interp(
            wave_new,
            wave,
            err,
            left=fill_val,
            right=fill_val,
        )
        return flux_new, err_new

    return flux_new


""" 
Vectorized implementation of the Flux Conserving resampler.
The implementation in `desispec` is not vectorized along the sample axis.
Obtained from: https://github.com/ACCarnall/SpectRes/blob/master/spectres/spectral_resampling.py
Implements the equations from: https://ui.adsabs.harvard.edu/abs/2017arXiv170505165C/abstract
Coed still involves a bunch of loops which can be eliminated using more vectorization or JIT
Look at the code from specutils for inspiration
"""


def make_bins(wavs):
    """Given a series of wavelength points, find the edges and widths
    of corresponding wavelength bins."""
    edges = np.zeros(wavs.shape[0] + 1)
    widths = np.zeros(wavs.shape[0])
    edges[0] = wavs[0] - (wavs[1] - wavs[0]) / 2
    widths[-1] = wavs[-1] - wavs[-2]
    edges[-1] = wavs[-1] + (wavs[-1] - wavs[-2]) / 2
    edges[1:-1] = (wavs[1:] + wavs[:-1]) / 2
    widths[:-1] = edges[1:-1] - edges[:-2]

    return edges, widths


def _spectres(
    new_wavs, spec_wavs, spec_fluxes, spec_errs=None, fill=None, verbose=True
):

    """
    Function for resampling spectra (and optionally associated
    uncertainties) onto a new wavelength basis.

    Parameters
    ----------

    new_wavs : numpy.ndarray
        Array containing the new wavelength sampling desired for the
        spectrum or spectra.

    spec_wavs : numpy.ndarray
        1D array containing the current wavelength sampling of the
        spectrum or spectra.

    spec_fluxes : numpy.ndarray
        Array containing spectral fluxes at the wavelengths specified in
        spec_wavs, last dimension must correspond to the shape of
        spec_wavs. Extra dimensions before this may be used to include
        multiple spectra.

    spec_errs : numpy.ndarray (optional)
        Array of the same shape as spec_fluxes containing uncertainties
        associated with each spectral flux value.

    fill : float (optional)
        Where new_wavs extends outside the wavelength range in spec_wavs
        this value will be used as a filler in new_fluxes and new_errs.

    verbose : bool (optional)
        Setting verbose to False will suppress the default warning about
        new_wavs extending outside spec_wavs and "fill" being used.

    Returns
    -------

    new_fluxes : numpy.ndarray
        Array of resampled flux values, first dimension is the same
        length as new_wavs, other dimensions are the same as
        spec_fluxes.

    new_errs : numpy.ndarray
        Array of uncertainties associated with fluxes in new_fluxes.
        Only returned if spec_errs was specified.
    """

    # Rename the input variables for clarity within the function.
    old_wavs = spec_wavs
    old_fluxes = spec_fluxes
    old_errs = spec_errs

    # Make arrays of edge positions and widths for the old and new bins

    old_edges, old_widths = make_bins(old_wavs)
    new_edges, new_widths = make_bins(new_wavs)

    # Generate output arrays to be populated
    new_fluxes = np.zeros(old_fluxes[..., 0].shape + new_wavs.shape)

    if old_errs is not None:
        if old_errs.shape != old_fluxes.shape:
            raise ValueError(
                "If specified, spec_errs must be the same shape " "as spec_fluxes."
            )
        else:
            new_errs = np.copy(new_fluxes)

    start = 0
    stop = 0
    warned = False

    # Calculate new flux and uncertainty values, looping over new bins
    for j in range(new_wavs.shape[0]):

        # Add filler values if new_wavs extends outside of spec_wavs
        if (new_edges[j] < old_edges[0]) or (new_edges[j + 1] > old_edges[-1]):
            new_fluxes[..., j] = fill

            if spec_errs is not None:
                new_errs[..., j] = fill

            if (j == 0 or j == new_wavs.shape[0] - 1) and verbose and not warned:
                warned = True
                print(
                    "\nSpectres: new_wavs contains values outside the range "
                    "in spec_wavs, new_fluxes and new_errs will be filled "
                    "with the value set in the 'fill' keyword argument. \n"
                )
            continue

        # Find first old bin which is partially covered by the new bin
        while old_edges[start + 1] <= new_edges[j]:
            start += 1

        # Find last old bin which is partially covered by the new bin
        while old_edges[stop + 1] < new_edges[j + 1]:
            stop += 1

        # If new bin is fully inside an old bin start and stop are equal
        if stop == start:
            new_fluxes[..., j] = old_fluxes[..., start]
            if old_errs is not None:
                new_errs[..., j] = old_errs[..., start]

        # Otherwise multiply the first and last old bin widths by P_ij
        else:
            start_factor = (old_edges[start + 1] - new_edges[j]) / (
                old_edges[start + 1] - old_edges[start]
            )

            end_factor = (new_edges[j + 1] - old_edges[stop]) / (
                old_edges[stop + 1] - old_edges[stop]
            )

            old_widths[start] *= start_factor
            old_widths[stop] *= end_factor

            # Populate new_fluxes spectrum and uncertainty arrays
            f_widths = old_widths[start : stop + 1] * old_fluxes[..., start : stop + 1]
            new_fluxes[..., j] = np.sum(f_widths, axis=-1)
            new_fluxes[..., j] /= np.sum(old_widths[start : stop + 1])

            if old_errs is not None:
                e_wid = old_widths[start : stop + 1] * old_errs[..., start : stop + 1]

                new_errs[..., j] = np.sqrt(np.sum(e_wid ** 2, axis=-1))
                new_errs[..., j] /= np.sum(old_widths[start : stop + 1])

            # Put back the old bin widths to their initial values
            old_widths[start] /= start_factor
            old_widths[stop] /= end_factor

    # If errors were supplied return both new_fluxes and new_errs.
    if old_errs is not None:
        return new_fluxes, new_errs

    # Otherwise just return the new_fluxes spectrum array
    else:
        return new_fluxes


def resample(
    wave_new, wave, flux, ivar=None, fill=None, method="linear", verbose=False
):

    """
    Function for resampling spectra (and optionally associated
    uncertainties) onto a new wavelength basis.

    Parameters
    ----------

    wave_new : numpy.ndarray
        Array containing the new wavelength sampling desired for the
        spectrum or spectra.

    wave : numpy.ndarray
        array containing the current wavelength sampling of the
        spectrum or spectra.

    flux : numpy.ndarray
        Array containing spectral fluxes at the wavelengths specified in
        spec_wavs, last dimension must correspond to the shape of
        spec_wavs. Extra dimensions before this may be used to include
        multiple spectra.

    ivar : numpy.ndarray (optional)
        Array of the same shape as spec_fluxes containing uncertainties
        associated with each spectral flux value.

    fill : float (optional)
        Where new_wavs extends outside the wavelength range in spec_wavs
        this value will be used as a filler in new_fluxes and new_errs.
    method : str (optional)
        choose from `linear` or `flux_cons`.
        Use a linear resampled spectra or flux conserving resampler.
    verbose : bool (optional)
        Setting verbose to False will suppress the default warning about
        new_wavs extending outside spec_wavs and "fill" being used.

    Returns
    -------

    flux_new : numpy.ndarray
        Array of resampled flux values, first dimension is the same
        length as new_wavs, other dimensions are the same as
        spec_fluxes.

    ivar_new : numpy.ndarray
        Array of uncertainties associated with fluxes in new_fluxes.
        Only returned if spec_errs was specified.
    New bins which include missing data are denoted by nan.
    """
    num_spec = len(flux)
    # TODO depending on num_spec choose whether to JIT or not
    if wave_new.ndim == 1:
        wave_new = np.repeat(wave_new[np.newaxis, :], num_spec, axis=0)
    if wave.ndim == 1:
        wave = np.repeat(wave[np.newaxis, :], num_spec, axis=0)

    if ivar is not None:
        ivar = ivar.copy()
        flux = flux.copy()
        flux[ivar == 0] = np.nan
        ivar[ivar == 0] = np.nan
        flux_new = np.zeros_like(wave_new)
        err_new = np.zeros_like(wave_new)
        err = np.sqrt(1 / ivar)
        # for i in range(num_spec):
        #     flux_new[i], err_new[i] = _spectres(
        #         wave_new[i], wave[i], flux[i], err[i], fill=fill, verbose=verbose
        #     )
        for i in range(num_spec):
            flux_new[i], err_new[i] = linear_interp_resampler(
                wave_new[i], wave[i], flux[i], err[i]
            )
        ivar_new = 1 / err_new ** 2
        # ivar_new[np.isnan(ivar_new)] == 0
        return flux_new, ivar_new

    # Otherwise just return the new_fluxes spectrum array
    else:
        flux_new = np.zeros_like(wave_new)
        for i in range(num_spec):
            flux_new[i] = _spectres(
                wave_new[i],
                wave[i],
                flux[i],
                spec_errs=None,
                fill=fill,
                verbose=verbose,
            )
        return flux_new
