"""
Band-averaged spectral estimation for use in tidal error analysis.

There is only one function in this module that is imported by
other modules: `band_psd` (``ut_pdgm`` in the Matlab version), which
uses the residuals from the tidal fit to calculate power spectral
density averaged in bands near the constituent clusters.

Because the functions here are intended for internal use only,
they are stripped down to their essentials, with little or
no argument checking.  Some argument flexibility is included
to facilitate testing.
"""

from __future__ import (absolute_import, division, print_function)

import numpy as np

from scipy import signal

from utide.utilities import Bunch


# __all__ = ['freq_bands', 'band_psd']

# Frequency bands in cycles per hour for
# estimating background noise levels.
freq_bands = np.array([[.00010, .00417],   # fortnightly etc.
                       [.03192, .04859],   # diurnal
                       [.07218, .08884],   # semidiurnal
                       [.11243, .12910],   # mk3, 2mk3
                       [.15269, .16936],   # about 4 cpd
                       [.19295, .20961],
                       [.23320, .25100],   # s6
                       [.26000, .29000],
                       [.30000, .50000]])  # m8 is 0.322; highest is 0.487


def fbndavg(P, freq, cfreq=None, fbands=None):
    """
    Band-average periodogram, excluding constituent frequencies

    Parameters
    ----------
    P : ndarray (nfreq,)
        periodogram in units^2 per CPH
    freq: ndarray (nfreq,)
        frequencies in CPH
    cfreq: None or array_like
        constituent frequencies (CPH) to be excluded
    fbands: array_like (nbands, 2)
        limits of frequency bands; defaults to 9 bands in `freq_bands`

    Returns
    -------
    ndarray (nbands,)
        mean of `P` within `fbands`

    Notes
    -----
    From ut_fbndavg, which was based on residual_spectrum.m of
    t_tide, Pawlowicz et al. (2002).
    """

    if fbands is None:
        fbands = freq_bands

    nfreq = len(freq)

    if cfreq is not None:
        # Exclude periodogram values that are nearest neighbors of
        # input constituent frequencies.
        P = np.ma.array(P, copy=True)
        indices = np.arange(nfreq)
        i_constit = np.round(np.interp(cfreq, freq, indices)).astype(int)
        i_constit = np.clip(i_constit, 0, nfreq)
        P[i_constit] = np.ma.masked

    nbands = fbands.shape[0]
    avP = np.ma.zeros((nbands,), dtype=P.dtype)

    for k, (f0, f1) in enumerate(fbands):
        i0, i1 = np.searchsorted(freq, [f0, f1])
        i1 = min(nfreq, i1+1)
        inband = slice(i0, i1)
        avP[k] = P[inband].mean()

    # Better to leave it as a masked array?
    avP = avP.filled(np.nan)

    return avP


def _lomb_freqs(t, fbands=None, ofac=1, max_per_band=500):
    """
    Calculate uniformly spaced frequencies in each band, with
    no more than `max_per_band` in any band.
    """

    if fbands is None:
        fbands = freq_bands

    # Estimated record length as n delta-t, where delta-t
    # is the *average* time per sample.
    n = len(t)
    delta_t = (t[-1] - t[0]) / (n-1)
    reclen = n * delta_t

    nf = n * ofac  # number of "Fourier frequencies" based on oversampling
    # Simplify by ignoring the 0 and Nyquist frequencies.
    # Divide by reclen to convert cycles/record to cycles/time unit.
    freq = np.arange(1, nf//2) / reclen
    nfreq = len(freq)

    freqs = []
    for k, (f0, f1) in enumerate(fbands):
        i0, i1 = np.searchsorted(freq, [f0, f1])
        i1 = min(nfreq, i1+1)
        inband = slice(i0, i1)
        if i1 - i0 > max_per_band:
            band = np.linspace(freq[i0], freq[i1-1], max_per_band)
        else:
            band = freq[inband]
        freqs.append(band)

    return np.hstack(freqs)