#
# Copyright 2015-2016 Universidad Complutense de Madrid
#
# This file is part of Numina
#
# SPDX-License-Identifier: GPL-3.0+
# License-Filename: LICENSE.txt
#

from __future__ import division
from __future__ import print_function

from mpl_toolkits.axes_grid.inset_locator import inset_axes
import numpy as np
from numpy.polynomial import Polynomial

from numina.array.display.pause_debugplot import pause_debugplot
from numina.array.display.ximplotxy import ximplotxy
from numina.modeling.gaussbox import gauss_box_model


def filtmask(sp, fmin=0.02, fmax=0.15, debugplot=0):
    """Filter spectrum in Fourier space and apply cosine bell.

    Parameters
    ----------
    sp : numpy array
        Spectrum to be filtered and masked.
    fmin : float
        Minimum frequency to be employed.
    fmax : float
        Maximum frequency to be employed.
    debugplot : int
        Debugging level for messages and plots. For details see
        'numina.array.display.pause_debugplot.py'.

    Returns
    -------
    sp_filtmask : numpy array
        Filtered and masked spectrum

    """

    # Fourier filtering
    xf = np.fft.fftfreq(sp.size)
    yf = np.fft.fft(sp)
    if abs(debugplot) in (21, 22):
        ximplotxy(xf, yf.real, xlim=(0, 0.51),
                  plottype='semilog', debugplot=debugplot)

    cut = (np.abs(xf) > fmax)
    yf[cut] = 0.0
    cut = (np.abs(xf) < fmin)
    yf[cut] = 0.0
    if abs(debugplot) in (21, 22):
        ximplotxy(xf, yf.real, xlim=(0, 0.51),
                  plottype='semilog', debugplot=debugplot)

    sp_filt = np.fft.ifft(yf).real
    if abs(debugplot) in (21, 22):
        xdum = np.arange(1, sp_filt.size + 1)
        ximplotxy(xdum, sp_filt, title="filtered median spectrum",
                  debugplot=debugplot)

    sp_filtmask = sp_filt * cosinebell(sp_filt.size, 0.1)
    if abs(debugplot) in (21, 22):
        xdum = np.arange(1, sp_filt.size + 1)
        ximplotxy(xdum, sp_filtmask,
                  title="filtered and masked median spectrum",
                  debugplot=debugplot)

    return sp_filtmask


def cosinebell(n, fraction):
    """Return a cosine bell spanning n pixels, masking a fraction of pixels

    Parameters
    ----------
    n : int
        Number of pixels.
    fraction : float
        Length fraction over which the data will be masked.

    """

    mask = np.ones(n)
    nmasked = int(fraction * n)
    for i in range(nmasked):
        yval = 0.5 * (1 - np.cos(np.pi * float(i) / float(nmasked)))
        mask[i] = yval
        mask[n - i - 1] = yval

    return mask


def convolve_comb_lines(lines_wave, lines_flux, sigma,
                        crpix1, crval1, cdelt1, naxis1):
    """Convolve a set of lines of known wavelengths and flux.

    Parameters
    ----------
    lines_wave : array like
        Input array with wavelengths
    lines_flux : array like
        Input array with fluxes
    sigma : float
        Sigma of the broadening gaussian to be applied.
    crpix1 : float
        CRPIX1 of the desired wavelength calibration.
    crval1 : float
        CRVAL1 of the desired wavelength calibration.
    cdelt1 : float
        CDELT1 of the desired wavelength calibration.
    naxis1 : integer
        NAXIS1 of the output spectrum.

    Returns
    -------
    xwave : array like
        Array with wavelengths for the output spectrum.
    spectrum : array like
        Array with the expected fluxes at each pixel.

    """

    # generate wavelengths for output spectrum
    xwave = crval1 + (np.arange(naxis1) + 1 - crpix1) * cdelt1

    # initialize output spectrum
    spectrum = np.zeros(naxis1)

    # convolve each line
    for wave, flux in zip(lines_wave, lines_flux):
        sp_tmp = gauss_box_model(x=xwave, amplitude=flux, mean=wave,
                                 stddev=sigma)
        spectrum += sp_tmp

    return xwave, spectrum


def periodic_corr1d(sp_reference, sp_offset,
                    fminmax=None,
                    naround_zero=None,
                    plottitle=None,
                    pdf=None,
                    debugplot=0):
    """Periodic correlation between two spectra, implemented using FFT.

    Parameters
    ----------
    sp_reference : numpy array
        Reference spectrum.
    sp_offset : numpy array
        Spectrum which offset is going to be measured relative to the
        reference spectrum.
    fminmax : tuple of floats or None
        Minimum and maximum frequencies to be used. If None, no
        frequency filtering is employed.
    naround_zero : int
        Half width of the window (around zero offset) to look for
        the correlation peak. If None, the whole correlation
        spectrum is employed. Otherwise, the peak will be sought
        in the interval [-naround_zero, +naround_zero].
    plottitle : str
        Optional plot title.
    pdf : PdfFile object or None
        If not None, output is sent to PDF file.
    debugplot : int
        Determines whether intermediate computations and/or plots
        are displayed. The valid codes are defined in
        numina.array.display.pause_debugplot.

    Returns
    -------
    offset : float
        Offset between the two input spectra.
    fpeak : float
        Maximum of the cross-correlation function.

    """

    # protections
    if sp_reference.ndim != 1 or sp_offset.ndim != 1:
        raise ValueError("Invalid array dimensions")
    if sp_reference.shape != sp_offset.shape:
        raise ValueError("x and y shapes are different")

    naxis1 = len(sp_reference)

    xcorr = np.arange(naxis1)
    naxis1_half = int(naxis1 / 2)
    for i in range(naxis1_half):
        xcorr[i + naxis1_half] -= naxis1
    isort = xcorr.argsort()
    xcorr = xcorr[isort]

    if fminmax is not None:
        fmin, fmax = fminmax
        sp_reference_filtmask = filtmask(sp_reference, fmin=fmin, fmax=fmax)
        sp_offset_filtmask = filtmask(sp_offset, fmin=fmin, fmax=fmax)
        if abs(debugplot) % 10 != 0:
            from numina.array.display.matplotlib_qt import plt
            xdum = np.arange(naxis1) + 1
            # reference spectrum
            ax = ximplotxy(xdum, sp_reference, show=False,
                           title='reference spectrum',
                           label='original spectrum')
            ax.plot(xdum, sp_reference_filtmask,
                    label='filtered and masked spectrum')
            ax.legend()
            plt.show()
            # offset spectrum
            ax = ximplotxy(xdum, sp_offset, show=False,
                           title='offset spectrum',
                           label='original spectrum')
            ax.plot(xdum, sp_offset_filtmask,
                    label='filtered and masked spectrum')
            ax.legend()
            plt.show()
    else:
        sp_reference_filtmask = np.copy(sp_reference)
        sp_offset_filtmask = np.copy(sp_offset)

    if (abs(debugplot) % 10 != 0) or (pdf is not None):
        xdum = np.arange(naxis1) + 1
        ax = ximplotxy(xdum, sp_reference_filtmask, show=False,
                       title=plottitle,
                       label='reference spectrum')
        ax.plot(xdum, sp_offset_filtmask, label='offset spectrum')
        ax.legend()
        if pdf is not None:
            pdf.savefig()
        else:
            pause_debugplot(debugplot=debugplot, pltshow=True)

    corr = np.fft.ifft(np.fft.fft(sp_offset_filtmask) *
                       np.fft.fft(sp_reference_filtmask).conj()).real
    corr = corr[isort]

    # determine correlation peak
    if naround_zero is None:
        iminpeak = 0
        imaxpeak = naxis1
    else:
        iminpeak = max(int(naxis1 / 2 - naround_zero), 0)
        imaxpeak = min(int(naxis1 / 2 + naround_zero), naxis1)
    ixpeak = corr[iminpeak:imaxpeak].argmax() + iminpeak

    # fit correlation peak with 2nd order polynomial
    nfit = 7
    nmed = nfit // 2
    imin = ixpeak - nmed
    imax = ixpeak + nmed
    lpeak_ok = True
    if imin < 0 or imax > len(corr):
        x_refined_peak = 0
        y_refined_peak = 0
        lpeak_ok = False
        poly_peak = Polynomial([0.0])
    else:
        x_fit = np.arange(-nmed, nmed + 1, dtype=np.float)
        y_fit = corr[imin:(imax+1)]
        poly_peak = Polynomial.fit(x_fit, y_fit, 2)
        poly_peak = Polynomial.cast(poly_peak)
        coef = poly_peak.coef
        if coef[2] != 0:
            x_refined_peak = -coef[1] / (2.0 * coef[2])
        else:
            x_refined_peak = 0.0
        y_refined_peak = poly_peak(x_refined_peak)
        x_refined_peak += ixpeak

    offset = x_refined_peak - naxis1_half
    fpeak = y_refined_peak

    if (abs(debugplot) % 10 != 0) or (pdf is not None):
        ax = ximplotxy(xcorr, corr,
                       xlabel='offset (pixels)',
                       ylabel='cross-correlation function',
                       title=plottitle,
                       xlim=(-naxis1/2, naxis1/2), show=False)
        ax.axvline(offset, color='grey', linestyle='dashed')
        coffset = "(offset:{0:6.2f} pixels)".format(offset)
        ax.text(0, 1, coffset,
                horizontalalignment='left',
                verticalalignment='top',
                transform=ax.transAxes)
        if naround_zero is not None:
            cwindow = "(peak region: [{},{}] pixels)".format(-naround_zero,
                                                             naround_zero)
            ax.text(0, 0.95, cwindow,
                    horizontalalignment='left',
                    verticalalignment='top',
                    transform=ax.transAxes)
        # inset plot
        inset_ax = inset_axes(
            ax,
            width="40%",
            height="40%",
            loc=1
        )
        inset_ax.plot(xcorr, corr)
        if naround_zero is None:
            inset_ax.set_xlim([-naround_zero, naround_zero])
        else:
            inset_ax.set_xlim([-50, 50])
        if lpeak_ok:
            xplot = np.arange(-nmed, nmed, 0.5)
            yplot = poly_peak(xplot)
            xplot += ixpeak - naxis1_half
            inset_ax.plot(xplot, yplot, '-')
            inset_ax.plot([x_refined_peak - naxis1_half],
                          [y_refined_peak], 'o')
        inset_ax.axvline(offset, color='grey', linestyle='dashed')
        if pdf is not None:
            pdf.savefig()
        else:
            pause_debugplot(debugplot=debugplot, pltshow=True)

    return offset, fpeak
