# -----------------------------------------------------------------------------
# Copyright (C) 2004-2022 Gordon Smyth, Yifang Hu, Matthew Ritchie, Jeremy Silver, James Wettenhall, Davis McCarthy, Di Wu, Wei Shi, Belinda Phipson, Aaron Lun, Natalie Thorne, Alicia Oshlack, Carolyn de Graaf, Yunshun Chen, Mette Langaas, Egil Ferkingstad, Marcus Davy, Francois Pepin, Dongseok Choi
# Copyright (C) 2024 Maximilien Colange

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
# -----------------------------------------------------------------------------

# This file is based on the file 'R/squeezeVar.R' of the Bioconductor limma package (version 3.55.1).


import numpy as np

from .fitFDist import fitFDist


def squeezeVar(var, df, covariate=None, robust=False, winsor_tail_p=(0.05, 0.1)):
    """
    Squeeze a set of sample variances together by computing empirical Bayes posterior means
    This function implements empirical Bayes algorithms proposed by
    [Smyth2004]_ and [Phipson2016]_.

    A conjugate Bayesian hierarchical model is assumed for a set of sample
    variances. The hyperparameters are estimated by fitting a scaled
    F-distribution to the sample variances. The function returns the posterior
    variances and the estimated hyperparameters.

    Specifically, the sample variances :code:`var` are assumed to follow scaled
    chi-squared distributions, conditional on the true variances, and a scaled
    inverse chi-squared prior is assumed for the true variances.  The scale and
    degrees of freedom of this prior distribution are estimated from the values
    of :code:`var`.

    The effect of this function is to squeeze the variances towards a common
    value, or to a global trend if a :code:`covariate` is provided. The
    squeezed variances have a smaller expected mean square error to the true
    variances than do the sample variances themselves.

    If :code:`covariate` is not :code:`None`, then the scale parameter of the
    prior distribution is assumed to depend on the covariate. If the covariate
    is average log-expression, then the effect is an intensity-dependent trend
    similar to that in [Sartor2006]_.

    :code:`robust=True` implements the robust empirical Bayes procedure of
    [Phipson2016]_ which allows some of the :code:`var` values to be outliers.

    Arguments
    ---------
    var : array_like
        1-D array of independent sample variances
    df : array_like
        1-D array of degrees of freedom for the sample variances
    covariate : ??
        if not :code:`None`, :code:`var_prior` will depend on this numeric
        covariate. Otherwise, :code:`var_prior` is constant.
    robust : bool
        whether the estimation of :code:`df_prior` and :code:`var_prior` be
        robustified against outlier sample variances
    winsor_tail_p : float or pair of floats
        left and right tail proportions of :code:`x` to Winsorize. Only used
        when :code:`robust=True`

    Returns
    -------
    dict
        a dictionnary with keys:

        - :code:`"var_post"`, 1-D array of posterior variances. Same length as
          :code:`var`.
        - :code:`"var_prior"`, location or scale of prior distribution. 1-D
          array of same length as :code:`var` if :code:`covariate` is not
          :code:`None`, otherwise a single value.
        - :code:`"df_prior"`, degrees of freedom of prior distribution. 1-D
          array of same length as :code:`var` if :code:`robust=True`, otherwise
          a single value.
    """
    n = len(var)

    # Degenerate special cases
    if n == 0:
        raise ValueError("var is empty")
    if n == 1:
        return {"var_post": var, "var_prior": var, "df_prior": 0}

    # When df==0, guard against missing or infinite values in var
    df = np.broadcast_to(df, var.shape)
    var[df == 0] = 0

    # Estimate hyperparameters
    if robust:
        raise NotImplementedError("Robust estimation in squeezeVar is not implemented")
    else:
        fit = fitFDist(var, df1=df, covariate=covariate)
        df_prior = fit["df2"]
    if np.isnan(df_prior).any():
        raise RuntimeError("Could not estimate prior df")

    # Posterior variances
    var_post = _squeezeVar(var=var, df=df, var_prior=fit["scale"], df_prior=df_prior)

    return {"var_post": var_post, "var_prior": fit["scale"], "df_prior": df_prior}


def _squeezeVar(var, df, var_prior, df_prior):
    """
    Squeeze posterior variances given hyperparameters

    NAs not allowed in :code:`df_prior`

    Arguments
    ---------
    var : array_like
        1-D array of independent sample variances
    df : array_like
        1-D array of degrees of freedom for the sample variances
    var_prior : array_like
        array of prior variances
    df_prior :
        array of degrees of freedom for the prior variances

    Returns
    -------
    ndarray
        array of posterior variances
    """
    df = np.broadcast_to(df, var.shape)
    var_prior = np.broadcast_to(var_prior, var.shape)
    df_prior = np.broadcast_to(df_prior, var.shape)

    isfin = np.isfinite(df_prior)
    var_post = np.zeros(var.shape)

    # For infinite df_prior, return var_prior
    var_post[~isfin] = var_prior[~isfin]

    var_post[isfin] = (df[isfin] * var[isfin] + df_prior[isfin] * var_prior[isfin]) / (
        df[isfin] + df_prior[isfin]
    )

    return var_post