"""
Utilities for computing periodogram statistics.

This is an internal module; users should access this functionality via the
``false_alarm_probability`` and ``false_alarm_level`` methods of the
``astropy.timeseries.LombScargle`` API.
"""

from functools import wraps

import numpy as np

from astropy import units as u


def _weighted_sum(val, dy):
    if dy is not None:
        return (val / dy ** 2).sum()
    else:
        return val.sum()


def _weighted_mean(val, dy):
    if dy is None:
        return val.mean()
    else:
        return _weighted_sum(val, dy) / _weighted_sum(np.ones(val.shape), dy)


def _weighted_var(val, dy):
    return _weighted_mean(val ** 2, dy) - _weighted_mean(val, dy) ** 2


def _gamma(N):
    from scipy.special import gammaln
    # Note: this is closely approximated by (1 - 0.75 / N) for large N
    return np.sqrt(2 / N) * np.exp(gammaln(N / 2) - gammaln((N - 1) / 2))


def vectorize_first_argument(func):
    @wraps(func)
    def new_func(x, *args, **kwargs):
        x = np.asarray(x)
        return np.array([func(xi, *args, **kwargs)
                         for xi in x.flat]).reshape(x.shape)
    return new_func


def pdf_single(z, N, normalization, dH=1, dK=3):
    """Probability density function for Lomb-Scargle periodogram

    Compute the expected probability density function of the periodogram
    for the null hypothesis - i.e. data consisting of Gaussian noise.

    Parameters
    ----------
    z : array-like
        The periodogram value.
    N : int
        The number of data points from which the periodogram was computed.
    normalization : {'standard', 'model', 'log', 'psd'}
        The periodogram normalization.
    dH, dK : int, optional
        The number of parameters in the null hypothesis and the model.

    Returns
    -------
    pdf : np.ndarray
        The expected probability density function.

    Notes
    -----
    For normalization='psd', the distribution can only be computed for
    periodograms constructed with errors specified.
    All expressions used here are adapted from Table 1 of Baluev 2008 [1]_.

    References
    ----------
    .. [1] Baluev, R.V. MNRAS 385, 1279 (2008)
    """
    z = np.asarray(z)
    if dK - dH != 2:
        raise NotImplementedError("Degrees of freedom != 2")
    Nk = N - dK

    if normalization == 'psd':
        return np.exp(-z)
    elif normalization == 'standard':
        return 0.5 * Nk * (1 - z) ** (0.5 * Nk - 1)
    elif normalization == 'model':
        return 0.5 * Nk * (1 + z) ** (-0.5 * Nk - 1)
    elif normalization == 'log':
        return 0.5 * Nk * np.exp(-0.5 * Nk * z)
    else:
        raise ValueError(f"normalization='{normalization}' is not recognized")


def fap_single(z, N, normalization, dH=1, dK=3):
    """Single-frequency false alarm probability for the Lomb-Scargle periodogram

    This is equal to 1 - cdf, where cdf is the cumulative distribution.
    The single-frequency false alarm probability should not be confused with
    the false alarm probability for the largest peak.

    Parameters
    ----------
    z : array-like
        The periodogram value.
    N : int
        The number of data points from which the periodogram was computed.
    normalization : {'standard', 'model', 'log', 'psd'}
        The periodogram normalization.
    dH, dK : int, optional
        The number of parameters in the null hypothesis and the model.

    Returns
    -------
    false_alarm_probability : np.ndarray
        The single-frequency false alarm probability.

    Notes
    -----
    For normalization='psd', the distribution can only be computed for
    periodograms constructed with errors specified.
    All expressions used here are adapted from Table 1 of Baluev 2008 [1]_.

    References
    ----------
    .. [1] Baluev, R.V. MNRAS 385, 1279 (2008)
    """
    z = np.asarray(z)
    if dK - dH != 2:
        raise NotImplementedError("Degrees of freedom != 2")
    Nk = N - dK

    if normalization == 'psd':
        return np.exp(-z)
    elif normalization == 'standard':
        return (1 - z) ** (0.5 * Nk)
    elif normalization == 'model':
        return (1 + z) ** (-0.5 * Nk)
    elif normalization == 'log':
        return np.exp(-0.5 * Nk * z)
    else:
        raise ValueError(f"normalization='{normalization}' is not recognized")


def inv_fap_single(fap, N, normalization, dH=1, dK=3):
    """Single-frequency inverse false alarm probability

    This function computes the periodogram value associated with the specified
    single-frequency false alarm probability. This should not be confused with
    the false alarm level of the largest peak.

    Parameters
    ----------
    fap : array-like
        The false alarm probability.
    N : int
        The number of data points from which the periodogram was computed.
    normalization : {'standard', 'model', 'log', 'psd'}
        The periodogram normalization.
    dH, dK : int, optional
        The number of parameters in the null hypothesis and the model.

    Returns
    -------
    z : np.ndarray
        The periodogram power corresponding to the single-peak false alarm
        probability.

    Notes
    -----
    For normalization='psd', the distribution can only be computed for
    periodograms constructed with errors specified.
    All expressions used here are adapted from Table 1 of Baluev 2008 [1]_.

    References
    ----------
    .. [1] Baluev, R.V. MNRAS 385, 1279 (2008)
    """
    fap = np.asarray(fap)
    if dK - dH != 2:
        raise NotImplementedError("Degrees of freedom != 2")
    Nk = N - dK

    # No warnings for fap = 0; rather, just let it give the right infinity.
    with np.errstate(divide='ignore'):
        if normalization == 'psd':
            return -np.log(fap)
        elif normalization == 'standard':
            return 1 - fap ** (2 / Nk)
        elif normalization == 'model':
            return -1 + fap ** (-2 / Nk)
        elif normalization == 'log':
            return -2 / Nk * np.log(fap)
        else:
            raise ValueError(f"normalization='{normalization}' is not recognized")


def cdf_single(z, N, normalization, dH=1, dK=3):
    """Cumulative distribution for the Lomb-Scargle periodogram

    Compute the expected cumulative distribution of the periodogram
    for the null hypothesis - i.e. data consisting of Gaussian noise.

    Parameters
    ----------
    z : array-like
        The periodogram value.
    N : int
        The number of data points from which the periodogram was computed.
    normalization : {'standard', 'model', 'log', 'psd'}
        The periodogram normalization.
    dH, dK : int, optional
        The number of parameters in the null hypothesis and the model.

    Returns
    -------
    cdf : np.ndarray
        The expected cumulative distribution function.

    Notes
    -----
    For normalization='psd', the distribution can only be computed for
    periodograms constructed with errors specified.
    All expressions used here are adapted from Table 1 of Baluev 2008 [1]_.

    References
    ----------
    .. [1] Baluev, R.V. MNRAS 385, 1279 (2008)
    """
    return 1 - fap_single(z, N, normalization=normalization, dH=dH, dK=dK)


def tau_davies(Z, fmax, t, y, dy, normalization='standard', dH=1, dK=3):
    """tau factor for estimating Davies bound (Baluev 2008, Table 1)"""
    N = len(t)
    NH = N - dH  # DOF for null hypothesis
    NK = N - dK  # DOF for periodic hypothesis
    Dt = _weighted_var(t, dy)
    Teff = np.sqrt(4 * np.pi * Dt)  # Effective baseline
    W = fmax * Teff
    Z = np.asarray(Z)
    if normalization == 'psd':
        # 'psd' normalization is same as Baluev's z
        return W * np.exp(-Z) * np.sqrt(Z)
    elif normalization == 'standard':
        # 'standard' normalization is Z = 2/NH * z_1
        return (_gamma(NH) * W * (1 - Z) ** (0.5 * (NK - 1))
                * np.sqrt(0.5 * NH * Z))
    elif normalization == 'model':
        # 'model' normalization is Z = 2/NK * z_2
        return (_gamma(NK) * W * (1 + Z) ** (-0.5 * NK)
                * np.sqrt(0.5 * NK * Z))
    elif normalization == 'log':
        # 'log' normalization is Z = 2/NK * z_3
        return (_gamma(NK) * W * np.exp(-0.5 * Z * (NK - 0.5))
                * np.sqrt(NK * np.sinh(0.5 * Z)))
    else:
        raise NotImplementedError(f"normalization={normalization}")


def fap_naive(Z, fmax, t, y, dy, normalization='standard'):
    """False Alarm Probability based on estimated number of indep frequencies"""
    N = len(t)
    T = max(t) - min(t)
    N_eff = fmax * T
    fap_s = fap_single(Z, N, normalization=normalization)
    # result is 1 - (1 - fap_s) ** N_eff
    # this is much more precise for small Z / large N
    # Ignore divide by zero no np.log1p - fine to let it return -inf.
    with np.errstate(divide='ignore'):
        return -np.expm1(N_eff * np.log1p(-fap_s))


def inv_fap_naive(fap, fmax, t, y, dy, normalization='standard'):
    """Inverse FAP based on estimated number of indep frequencies"""
    fap = np.asarray(fap)
    N = len(t)
    T = max(t) - min(t)
    N_eff = fmax * T
    # fap_s = 1 - (1 - fap) ** (1 / N_eff)
    # Ignore divide by zero no np.log - fine to let it return -inf.
    with np.errstate(divide='ignore'):
        fap_s = -np.expm1(np.log(1 - fap) / N_eff)
    return inv_fap_single(fap_s, N, normalization)


def fap_davies(Z, fmax, t, y, dy, normalization='standard'):
    """Davies upper-bound to the false alarm probability

    (Eqn 5 of Baluev 2008)
    """
    N = len(t)
    fap_s = fap_single(Z, N, normalization=normalization)
    tau = tau_davies(Z, fmax, t, y, dy, normalization=normalization)
    return fap_s + tau


@vectorize_first_argument
def inv_fap_davies(p, fmax, t, y, dy, normalization='standard'):
    """Inverse of the davies upper-bound"""
    from scipy import optimize
    args = (fmax, t, y, dy, normalization)
    z0 = inv_fap_naive(p, *args)
    func = lambda z, *args: fap_davies(z, *args) - p
    res = optimize.root(func, z0, args=args, method='lm')
    if not res.success:
        raise ValueError(f'inv_fap_baluev did not converge for p={p}')
    return res.x


def fap_baluev(Z, fmax, t, y, dy, normalization='standard'):
    """Alias-free approximation to false alarm probability

    (Eqn 6 of Baluev 2008)
    """
    fap_s = fap_single(Z, len(t), normalization)
    tau = tau_davies(Z, fmax, t, y, dy, normalization=normalization)
    # result is 1 - (1 - fap_s) * np.exp(-tau)
    # this is much more precise for small numbers
    return -np.expm1(-tau) + fap_s * np.exp(-tau)


@vectorize_first_argument
def inv_fap_baluev(p, fmax, t, y, dy, normalization='standard'):
    """Inverse of the Baluev alias-free approximation"""
    from scipy import optimize
    args = (fmax, t, y, dy, normalization)
    z0 = inv_fap_naive(p, *args)
    func = lambda z, *args: fap_baluev(z, *args) - p
    res = optimize.root(func, z0, args=args, method='lm')
    if not res.success:
        raise ValueError(f'inv_fap_baluev did not converge for p={p}')
    return res.x


def _bootstrap_max(t, y, dy, fmax, normalization, random_seed, n_bootstrap=1000):
    """Generate a sequence of bootstrap estimates of the max"""
    from .core import LombScargle
    rng = np.random.default_rng(random_seed)
    power_max = []
    for _ in range(n_bootstrap):
        s = rng.integers(0, len(y), len(y))  # sample with replacement
        ls_boot = LombScargle(t, y[s], dy if dy is None else dy[s],
                              normalization=normalization)
        freq, power = ls_boot.autopower(maximum_frequency=fmax)
        power_max.append(power.max())

    power_max = u.Quantity(power_max)
    power_max.sort()

    return power_max


def fap_bootstrap(Z, fmax, t, y, dy, normalization='standard',
                  n_bootstraps=1000, random_seed=None):
    """Bootstrap estimate of the false alarm probability"""
    pmax = _bootstrap_max(t, y, dy, fmax, normalization, random_seed,
                          n_bootstraps)

    return 1 - np.searchsorted(pmax, Z) / len(pmax)


def inv_fap_bootstrap(fap, fmax, t, y, dy, normalization='standard',
                      n_bootstraps=1000, random_seed=None):
    """Bootstrap estimate of the inverse false alarm probability"""
    fap = np.asarray(fap)
    pmax = _bootstrap_max(t, y, dy, fmax, normalization, random_seed,
                          n_bootstraps)

    return pmax[np.clip(np.floor((1 - fap) * len(pmax)).astype(int),
                        0, len(pmax) - 1)]


METHODS = {'single': fap_single,
           'naive': fap_naive,
           'davies': fap_davies,
           'baluev': fap_baluev,
           'bootstrap': fap_bootstrap}


def false_alarm_probability(Z, fmax, t, y, dy, normalization='standard',
                            method='baluev', method_kwds=None):
    """Compute the approximate false alarm probability for periodogram peaks Z

    This gives an estimate of the false alarm probability for the largest value
    in a periodogram, based on the null hypothesis of non-varying data with
    Gaussian noise. The true probability cannot be computed analytically, so
    each method available here is an approximation to the true value.

    Parameters
    ----------
    Z : array-like
        The periodogram value.
    fmax : float
        The maximum frequency of the periodogram.
    t, y, dy : array-like
        The data times, values, and errors.
    normalization : {'standard', 'model', 'log', 'psd'}, optional
        The periodogram normalization.
    method : {'baluev', 'davies', 'naive', 'bootstrap'}, optional
        The approximation method to use.
    method_kwds : dict, optional
        Additional method-specific keywords.

    Returns
    -------
    false_alarm_probability : np.ndarray
        The false alarm probability.

    Notes
    -----
    For normalization='psd', the distribution can only be computed for
    periodograms constructed with errors specified.

    See Also
    --------
    false_alarm_level : compute the periodogram level for a particular fap

    References
    ----------
    .. [1] Baluev, R.V. MNRAS 385, 1279 (2008)
    """
    if method == 'single':
        return fap_single(Z, len(t), normalization)
    elif method not in METHODS:
        raise ValueError(f"Unrecognized method: {method}")
    method = METHODS[method]
    method_kwds = method_kwds or {}

    return method(Z, fmax, t, y, dy, normalization, **method_kwds)


INV_METHODS = {'single': inv_fap_single,
               'naive': inv_fap_naive,
               'davies': inv_fap_davies,
               'baluev': inv_fap_baluev,
               'bootstrap': inv_fap_bootstrap}


def false_alarm_level(p, fmax, t, y, dy, normalization,
                      method='baluev', method_kwds=None):
    """Compute the approximate periodogram level given a false alarm probability

    This gives an estimate of the periodogram level corresponding to a specified
    false alarm probability for the largest peak, assuming a null hypothesis
    of non-varying data with Gaussian noise. The true level cannot be computed
    analytically, so each method available here is an approximation to the true
    value.

    Parameters
    ----------
    p : array-like
        The false alarm probability (0 < p < 1).
    fmax : float
        The maximum frequency of the periodogram.
    t, y, dy : arrays
        The data times, values, and errors.
    normalization : {'standard', 'model', 'log', 'psd'}, optional
        The periodogram normalization.
    method : {'baluev', 'davies', 'naive', 'bootstrap'}, optional
        The approximation method to use.
    method_kwds : dict, optional
        Additional method-specific keywords.

    Returns
    -------
    z : np.ndarray
        The periodogram level.

    Notes
    -----
    For normalization='psd', the distribution can only be computed for
    periodograms constructed with errors specified.

    See Also
    --------
    false_alarm_probability : compute the fap for a given periodogram level

    References
    ----------
    .. [1] Baluev, R.V. MNRAS 385, 1279 (2008)
    """
    if method == 'single':
        return inv_fap_single(p, len(t), normalization)
    elif method not in INV_METHODS:
        raise ValueError(f"Unrecognized method: {method}")
    method = INV_METHODS[method]
    method_kwds = method_kwds or {}

    return method(p, fmax, t, y, dy, normalization, **method_kwds)