# -----------------------------------------------------------------------------
# Copyright (C) 2008-2022 Yunshun Chen, Aaron TL Lun, Davis J McCarthy, Matthew E Ritchie, Belinda Phipson, Yifang Hu, Xiaobei Zhou, Mark D Robinson, Gordon K Smyth
# Copyright (C) 2022-2023 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/mglmOneGroup.R' of the Bioconductor edgeR package (version 3.38.4).


import numpy as np

from ..utils import LOGGER
from .glm_one_group import fit_one_group
from .makeCompressedMatrix import (
    _compressDispersions,
    _compressOffsets,
    _compressWeights,
)
from .utils import _isAllZero


def mglmOneGroup(
    y,
    dispersion=0,
    offset=0,
    weights=None,
    coef_start=None,
    maxit=50,
    tol=1e-10,
    verbose=False,
):
    """
    Fit single-group negative-binomial GLMs genewise.

    This is a low-level work-horse used by higher-level functions, especially
    :func:`glmFit`. Most users will not need to call this function directly.

    This function fits a negative binomial GLM to each row of :code:`y`. The
    row-wise GLMs all have the same design matrix but possibly different
    dispersions, offsets and weights. It is low-level in that it
    operates on atomic objects (matrices and vectors).

    This function fits an intercept only model to each response vector. In other
    words, it treats all the libraries as belonging to one group. It implements
    Fisher scoring with a score-statistic stopping criterion for each gene.  It
    treats the dispersion parameter of the negative binomial distribution as a
    known input.  Excellent starting values are available for the null model so
    this function seldom has any problems with convergence. It is used by other
    functions to compute the overall abundance for each gene.

    Arguments
    ---------
    y : array_like
        matrix of negative binomial counts. Rows for genes and columns for
        libraries.
    dispersion : float or array_like
        scalar or vector giving the dispersion parameter for each GLM. Can be a
        scalar giving one value for all genes, or a vector of length equal to
        the number of genes giving genewise dispersions.
    offset : array_like
        vector or matrix giving the offset that is to be included in the log
        linear model predictor. Can be a scalar, a vector of length equal to the
        number of libraries, or a matrix of the same shape as :code:`y`.
    weights : matrix, optional
        vector or matrix of non-negative quantitative weights. Can be a vector
        of length equal to the number of libraries, or a matrix of the same
        shape as :code:`y`.
    coef_start : array_like, optional
        matrix of starting values for the linear model coefficient. Number of
        rows should agree with :code:`y` and a single column. This argument does
        not usually need to be set as the automatic starting values perform well.
    maxit : int
        the maximum number of iterations for the Fisher scoring algorithm. The
        iteration will be stopped when this limit is reached even if the
        convergence criterion has not been satisfied.
    tol : float
        the convergence tolerance. Convergence if judged successful when the
        step size falls below :code:`tol` in absolute size.

    Returns
    -------
    ndarray
        vector of coefficients
    """
    # Check y
    # TODO check that y is a matrix and numeric
    _isAllZero(y)

    # Check dispersion
    dispersion = _compressDispersions(y, dispersion)

    # Check offset
    offset = _compressOffsets(y, offset=offset)

    # Check starting values
    if coef_start is None:
        coef_start = np.nan
    coef_start = np.full((y.shape[0],), coef_start, dtype="double")

    # Check weights
    weights = _compressWeights(y, weights)

    # Fisher scoring iteration
    output = fit_one_group(y, offset, dispersion, weights, maxit, tol, coef_start)

    # Convergence achieved for all tags?
    if np.count_nonzero(output[1]) > 0:
        LOGGER.debug(f"max iterations exceeded for {np.count_nonzero(output[1])} tags")

    return output[0]