"""Python implementation of elastic-net regularized GLMs."""

from copy import deepcopy

import numpy as np
from scipy.special import expit
from scipy.stats import norm
from .utils import logger, set_log_level
from .base import BaseEstimator, is_classifier, check_version


def _lmb(distr, beta0, beta, X, eta):
    """Conditional intensity function."""
    z = beta0 + np.dot(X, beta)
    return _mu(distr, z, eta)


def _mu(distr, z, eta):
    """The non-linearity (inverse link)."""
    if distr in ['softplus', 'gamma']:
        mu = np.log1p(np.exp(z))
    elif distr == 'poisson':
        mu = z.copy()
        intercept = (1 - eta) * np.exp(eta)
        mu[z > eta] = z[z > eta] * np.exp(eta) + intercept
        mu[z <= eta] = np.exp(z[z <= eta])
    elif distr == 'gaussian':
        mu = z
    elif distr == 'binomial':
        mu = expit(z)
    elif distr == 'probit':
        mu = norm.cdf(z)
    return mu


def _grad_mu(distr, z, eta):
    """Derivative of the non-linearity."""
    if distr in ['softplus', 'gamma']:
        grad_mu = expit(z)
    elif distr == 'poisson':
        grad_mu = z.copy()
        grad_mu[z > eta] = np.ones_like(z)[z > eta] * np.exp(eta)
        grad_mu[z <= eta] = np.exp(z[z <= eta])
    elif distr == 'gaussian':
        grad_mu = np.ones_like(z)
    elif distr == 'binomial':
        grad_mu = expit(z) * (1 - expit(z))
    elif distr in 'probit':
        grad_mu = norm.pdf(z)
    return grad_mu


def _logL(distr, y, y_hat):
    """The log likelihood."""
    if distr in ['softplus', 'poisson']:
        eps = np.spacing(1)
        logL = np.sum(y * np.log(y_hat + eps) - y_hat)
    elif distr == 'gaussian':
        logL = -0.5 * np.sum((y - y_hat)**2)
    elif distr == 'binomial':
        # analytical formula
        logL = np.sum(y * np.log(y_hat) + (1 - y) * np.log(1 - y_hat))

        # but this prevents underflow
        # z = beta0 + np.dot(X, beta)
        # logL = np.sum(y * z - np.log(1 + np.exp(z)))
    elif distr == 'probit':
        logL = np.sum(y * np.log(y_hat) + (1 - y) * np.log(1 - y_hat))
    elif distr == 'gamma':
        # see
        # https://www.statistics.ma.tum.de/fileadmin/w00bdb/www/czado/lec8.pdf
        nu = 1.  # shape parameter, exponential for now
        logL = np.sum(nu * (-y / y_hat - np.log(y_hat)))
    return logL


def _penalty(alpha, beta, Tau, group):
    """The penalty."""
    # Combine L1 and L2 penalty terms
    P = 0.5 * (1 - alpha) * _L2penalty(beta, Tau) + \
        alpha * _L1penalty(beta, group)
    return P


def _L2penalty(beta, Tau):
    """The L2 penalty."""
    # Compute the L2 penalty
    if Tau is None:
        # Ridge=like penalty
        L2penalty = np.linalg.norm(beta, 2) ** 2
    else:
        # Tikhonov penalty
        if (Tau.shape[0] != beta.shape[0] or
           Tau.shape[1] != beta.shape[0]):
            raise ValueError('Tau should be (n_features x n_features)')
        else:
            L2penalty = np.linalg.norm(np.dot(Tau, beta), 2) ** 2
    return L2penalty


def _L1penalty(beta, group=None):
    """The L1 penalty."""
    # Compute the L1 penalty
    if group is None:
        # Lasso-like penalty
        L1penalty = np.linalg.norm(beta, 1)
    else:
        # Group sparsity case: apply group sparsity operator
        group_ids = np.unique(group)
        L1penalty = 0.0
        for group_id in group_ids:
            if group_id != 0:
                L1penalty += \
                    np.linalg.norm(beta[group == group_id], 2)
        L1penalty += np.linalg.norm(beta[group == 0], 1)
    return L1penalty


def _loss(distr, alpha, Tau, reg_lambda, X, y, eta, group, beta):
    """Define the objective function for elastic net."""
    n_samples = X.shape[0]
    y_hat = _mu(distr, beta[0] + np.dot(X, beta[1:]), eta)
    L = 1. / n_samples * _logL(distr, y, y_hat)
    P = _penalty(alpha, beta[1:], Tau, group)
    J = -L + reg_lambda * P
    return J


def _L2loss(distr, alpha, Tau, reg_lambda, X, y, eta, group, beta):
    """Define the objective function for elastic net."""
    n_samples = X.shape[0]
    y_hat = _mu(distr, beta[0] + np.dot(X, beta[1:]), eta)
    L = 1. / n_samples * _logL(distr, y, y_hat)
    P = 0.5 * (1 - alpha) * _L2penalty(beta[1:], Tau)
    J = -L + reg_lambda * P
    return J


def _grad_L2loss(distr, alpha, Tau, reg_lambda, X, y, eta, beta):
    """The gradient."""
    n_samples = np.float(X.shape[0])

    if Tau is None:
        Tau = np.eye(beta[1:].shape[0])
    InvCov = np.dot(Tau.T, Tau)

    z = beta[0] + np.dot(X, beta[1:])
    mu = _mu(distr, z, eta)
    grad_mu = _grad_mu(distr, z, eta)

    if distr in ['poisson', 'softplus']:
        grad_beta0 = np.sum(grad_mu) - np.sum(y * grad_mu / mu)
        grad_beta = ((np.dot(grad_mu.T, X) -
                     np.dot((y * grad_mu / mu).T, X)).T)

    elif distr == 'gaussian':
        grad_beta0 = np.sum((mu - y) * grad_mu)
        grad_beta = np.dot((mu - y).T, X * grad_mu[:, None]).T

    elif distr == 'binomial':
        grad_beta0 = np.sum(mu - y)
        grad_beta = np.dot((mu - y).T, X).T

    elif distr == 'probit':
        grad_beta0 = -np.sum((y * (grad_mu / mu)) -
                             ((1 - y) * (grad_mu / (1 - mu))))
        grad_logl = ((y * (grad_mu / mu)) -
                     ((1 - y) * (grad_mu / (1 - mu))))
        grad_beta = -np.dot(grad_logl.T, X).T

    elif distr == 'gamma':
        nu = 1.
        grad_logl = (y / mu ** 2 - 1 / mu) * grad_mu
        grad_beta0 = -nu * np.sum(grad_logl)
        grad_beta = -nu * np.dot(grad_logl.T, X).T

    grad_beta0 *= 1. / n_samples
    grad_beta *= 1. / n_samples
    grad_beta += reg_lambda * (1 - alpha) * np.dot(InvCov, beta[1:])
    n_features = X.shape[1]
    g = np.zeros((n_features + 1, ))
    g[0] = grad_beta0
    g[1:] = grad_beta
    return g


def _gradhess_logloss_1d(distr, xk, y, z, eta):
    """
    Compute gradient (1st derivative)
    and Hessian (2nd derivative)
    of log likelihood for a single coordinate.

    Parameters
    ----------
    xk: float
        (n_samples,)
    y: float
        (n_samples,)
    z: float
        (n_samples,)

    Returns
    -------
    gk: gradient, float:
        (n_features + 1,)
    hk: float:
        (n_features + 1,)
    """
    n_samples = xk.shape[0]

    if distr == 'softplus':
        mu = _mu(distr, z, eta)
        s = expit(z)
        gk = np.sum(s * xk) - np.sum(y * s / mu * xk)

        grad_s = s * (1 - s)
        grad_s_by_mu = grad_s / mu - s / (mu ** 2)
        hk = np.sum(grad_s * xk ** 2) - np.sum(y * grad_s_by_mu * xk ** 2)

    elif distr == 'poisson':
        mu = _mu(distr, z, eta)
        s = expit(z)
        gk = np.sum((mu[z <= eta] - y[z <= eta]) *
                    xk[z <= eta]) + \
            np.exp(eta) * \
            np.sum((1 - y[z > eta] / mu[z > eta]) *
                   xk[z > eta])
        hk = np.sum(mu[z <= eta] * xk[z <= eta] ** 2) + \
            np.exp(eta) ** 2 * \
            np.sum(y[z > eta] / (mu[z > eta] ** 2) *
                   (xk[z > eta] ** 2))

    elif distr == 'gaussian':
        gk = np.sum((z - y) * xk)
        hk = np.sum(xk * xk)

    elif distr == 'binomial':
        mu = _mu(distr, z, eta)
        gk = np.sum((mu - y) * xk)
        hk = np.sum(mu * (1.0 - mu) * xk * xk)

    elif distr == 'probit':
        prob = norm.cdf(z)
        grad_prob = norm.pdf(z)
        gk = np.sum(y * (grad_prob / prob) -
                    (1 - y) * (grad_prob / (1 - prob)) * xk)

        mid_l = y * (z * prob + grad_prob) / (prob ** 2)
        mid_r = (1 - y) * (-z * (1 - prob) + grad_prob) / ((1 - prob) ** 2)
        hk = np.sum(prob * (mid_l + mid_r) * (xk * xk))

    elif distr == 'gamma':
        raise NotImplementedError('cdfast is not implemented for Gamma '
                                  'distribution')

    return 1. / n_samples * gk, 1. / n_samples * hk


def simulate_glm(distr, beta0, beta, X, eta=2.0, random_state=None):
    """Simulate target data under a generative model.

    Parameters
    ----------
    distr: str
        distribution
    beta0: float
        intercept coefficient
    beta: array
        coefficients of shape (n_features,)
    X: array
        design matrix of shape (n_samples, n_features)
    eta: float
        parameter for poisson non-linearity
    random_state: float
        random state

    Returns
    -------
    y: array
        simulated target data of shape (n_samples,)
    """
    if random_state is not None:
        np.random.RandomState(random_state)
    if distr == 'softplus' or distr == 'poisson':
        y = np.random.poisson(_lmb(distr, beta0, beta, X, eta))
    if distr == 'gaussian':
        y = np.random.normal(_lmb(distr, beta0, beta, X, eta))
    if distr == 'binomial' or distr == 'probit':
        y = np.random.binomial(1, _lmb(distr, beta0, beta, X, eta))
    if distr == 'gamma':
        mu = _lmb(distr, beta0, beta, X, eta)
        y = np.exp(mu)
    return y


class GLM(BaseEstimator):
    """Class for estimating regularized generalized linear models (GLM).
    The regularized GLM minimizes the penalized negative log likelihood:

    .. math::

        \min_{\\beta_0, \\beta} \\frac{1}{N}
        \sum_{i = 1}^N \mathcal{L} (y_i, \\beta_0 + \\beta^T x_i)
        + \lambda [ \\frac{1}{2}(1 - \\alpha) \mathcal{P}_2 +
                    \\alpha \mathcal{P}_1 ]

    where :math:`\mathcal{P}_2` and :math:`\mathcal{P}_1` are the generalized
    L2 (Tikhonov) and generalized L1 (Group Lasso) penalties, given by:

    .. math::

        \mathcal{P}_2 = \|\Gamma \\beta \|_2^2 \\
        \mathcal{P}_1 = \sum_g \|\\beta_{j,g}\|_2

    where :math:`\Gamma` is the Tikhonov matrix: a square factorization
    of the inverse covariance matrix and :math:`\\beta_{j,g}` is the
    :math:`j` th coefficient of group :math:`g`.

    The generalized L2 penalty defaults to the ridge penalty when
    :math:`\Gamma` is identity.

    The generalized L1 penalty defaults to the lasso penalty when each
    :math:`\\beta` belongs to its own group.

    Parameters
    ----------
    distr : str
        distribution family can be one of the following
        'gaussian' | 'binomial' | 'poisson' | 'softplus'
        default: 'poisson'.
    alpha : float
        the weighting between L1 penalty and L2 penalty term
        of the loss function.
        default: 0.5
    Tau : array | None
        the (n_features, n_features) Tikhonov matrix.
        default: None, in which case Tau is identity
        and the L2 penalty is ridge-like
    group : array | list | None
        the (n_features, )
        list or array of group identities for each parameter :math:`\\beta`.
        Each entry of the list/ array should contain an int from 1 to n_groups
        that specify group membership for each parameter
        (except :math:`\\beta_0`).
        If you do not want to specify a group for a specific parameter,
        set it to zero.
        default: None, in which case it defaults to L1 regularization
    reg_lambda : float
        regularization parameter :math:`\\lambda` of penalty term.
        default: 0.1
    solver : str
        optimization method, can be one of the following
        'batch-gradient' (vanilla batch gradient descent)
        'cdfast' (Newton coordinate gradient descent).
        default: 'batch-gradient'
    learning_rate : float
        learning rate for gradient descent.
        default: 2e-1
    max_iter : int
        maximum iterations for the model.
        default: 1000
    tol : float
        convergence threshold or stopping criteria.
        Optimization loop will stop below setting threshold.
        default: 1e-3
    eta : float
        a threshold parameter that linearizes the exp() function above eta.
        default: 2.0
    score_metric : str
        specifies the scoring metric.
        one of either 'deviance' or 'pseudo_R2'.
        default: 'deviance'
    random_state : int
        seed of the random number generator used to initialize the solution.
        default: 0
    verbose : boolean or int
        default: False

    Reference
    ---------
    Friedman, Hastie, Tibshirani (2010). Regularization Paths for Generalized
        Linear Models via Coordinate Descent, J Statistical Software.
        https://core.ac.uk/download/files/153/6287975.pdf
    """

    def __init__(self, distr='poisson', alpha=0.5,
                 Tau=None, group=None,
                 reg_lambda=0.1,
                 solver='batch-gradient',
                 learning_rate=2e-1, max_iter=1000,
                 tol=1e-3, eta=2.0, score_metric='deviance',
                 random_state=0, verbose=False):

        if not isinstance(max_iter, int):
            raise ValueError('max_iter must be of type int')

        self.distr = distr
        self.alpha = alpha
        self.reg_lambda = reg_lambda
        self.Tau = Tau
        self.group = group
        self.solver = solver
        self.learning_rate = learning_rate
        self.max_iter = max_iter
        self.beta0_ = None
        self.beta_ = None
        self.ynull_ = None
        self.tol = tol
        self.eta = eta
        self.score_metric = score_metric
        self.random_state = random_state
        self.verbose = verbose
        set_log_level(verbose)

    def _set_cv(cv, estimator=None, X=None, y=None):
        """Set the default CV depending on whether clf
           is classifier/regressor."""
        # Detect whether classification or regression
        if estimator in ['classifier', 'regressor']:
            est_is_classifier = estimator == 'classifier'
        else:
            est_is_classifier = is_classifier(estimator)
        # Setup CV
        if check_version('sklearn', '0.18'):
            from sklearn import model_selection as models
            from sklearn.model_selection import (check_cv,
                                                 StratifiedKFold, KFold)
            if isinstance(cv, (int, np.int)):
                XFold = StratifiedKFold if est_is_classifier else KFold
                cv = XFold(n_splits=cv)
            elif isinstance(cv, str):
                if not hasattr(models, cv):
                    raise ValueError('Unknown cross-validation')
                cv = getattr(models, cv)
                cv = cv()
            cv = check_cv(cv=cv, y=y, classifier=est_is_classifier)
        else:
            from sklearn import cross_validation as models
            from sklearn.cross_validation import (check_cv,
                                                  StratifiedKFold, KFold)
            if isinstance(cv, (int, np.int)):
                if est_is_classifier:
                    cv = StratifiedKFold(y=y, n_folds=cv)
                else:
                    cv = KFold(n=len(y), n_folds=cv)
            elif isinstance(cv, str):
                if not hasattr(models, cv):
                    raise ValueError('Unknown cross-validation')
                cv = getattr(models, cv)
                if cv.__name__ not in ['KFold', 'LeaveOneOut']:
                    raise NotImplementedError('CV cannot be defined with str'
                                              ' for sklearn < .017.')
                cv = cv(len(y))
            cv = check_cv(cv=cv, X=X, y=y, classifier=est_is_classifier)

        # Extract train and test set to retrieve them at predict time
        if hasattr(cv, 'split'):
            cv_splits = [(train, test) for train, test in
                         cv.split(X=np.zeros_like(y), y=y)]
        else:
            # XXX support sklearn.cross_validation cv
            cv_splits = [(train, test) for train, test in cv]

        if not np.all([len(train) for train, _ in cv_splits]):
            raise ValueError('Some folds do not have any train epochs.')

        return cv, cv_splits

    def __repr__(self):
        """Description of the object."""
        reg_lambda = self.reg_lambda
        s = '<\nDistribution | %s' % self.distr
        s += '\nalpha | %0.2f' % self.alpha
        s += '\nmax_iter | %0.2f' % self.max_iter
        s += '\nlambda: %0.2f\n>' % reg_lambda
        return s

    def copy(self):
        """Return a copy of the object.

        Parameters
        ----------
        none:

        Returns
        -------
        self: instance of GLM
            A copy of the GLM instance.
        """
        return deepcopy(self)

    def _prox(self, beta, thresh):
        """Proximal operator."""
        if self.group is None:
            # The default case: soft thresholding
            return np.sign(beta) * (np.abs(beta) - thresh) * \
                (np.abs(beta) > thresh)
        else:
            # Group sparsity case: apply group sparsity operator
            group_ids = np.unique(self.group)
            group_norms = np.abs(beta)

            for group_id in group_ids:
                if group_id != 0:
                    group_norms[self.group == group_id] = \
                        np.linalg.norm(beta[self.group == group_id], 2)

            nzero_norms = group_norms > 0.0
            over_thresh = group_norms > thresh
            idxs_to_update = nzero_norms & over_thresh

            result = beta
            result[idxs_to_update] = (beta[idxs_to_update] -
                                      thresh * beta[idxs_to_update] /
                                      group_norms[idxs_to_update])

            return result

    def _cdfast(self, X, y, z, ActiveSet, beta, rl):
        """
        Perform one cycle of Newton updates for all coordinates.

        Parameters
        ----------
        X : array
            n_samples x n_features
            The input data
        y : array
            Labels to the data
            n_samples x 1
        z:  array
            n_samples x 1
            beta[0] + X * beta[1:]
        ActiveSet: array
            n_features + 1 x 1
            Active set storing which betas are non-zero
        beta: array
            n_features + 1 x 1
            Parameters to be updated
        rl: float
            Regularization lambda

        Returns
        -------
        beta: array
            (n_features + 1) x 1
            Updated parameters
        z: array
            beta[0] + X * beta[1:]
            (n_features + 1) x 1
        """
        n_samples = X.shape[0]
        n_features = X.shape[1]
        reg_scale = rl * (1 - self.alpha)

        for k in range(0, n_features + 1):
            # Only update parameters in active set
            if ActiveSet[k] != 0:
                if k > 0:
                    xk = X[:, k - 1]
                else:
                    xk = np.ones((n_samples, ))

                # Calculate grad and hess of log likelihood term
                gk, hk = _gradhess_logloss_1d(self.distr, xk, y, z, self.eta)

                # Add grad and hess of regularization term
                if self.Tau is None:
                    gk_reg = beta[k]
                    hk_reg = 1.0
                else:
                    InvCov = np.dot(self.Tau.T, self.Tau)
                    gk_reg = np.sum(InvCov[k - 1, :] * beta[1:])
                    hk_reg = InvCov[k - 1, k - 1]
                gk += np.ravel([reg_scale * gk_reg if k > 0 else 0.0])
                hk += np.ravel([reg_scale * hk_reg if k > 0 else 0.0])

                # Update parameters, z
                update = 1. / hk * gk
                beta[k], z = beta[k] - update, z - update * xk
        return beta, z

    def fit(self, X, y):
        """The fit function.

        Parameters
        ----------
        X : array
            The input data of shape (n_samples, n_features)

        y : array
            The target data

        Returns
        -------
        self : instance of GLM
            The fitted model.
        """
        np.random.RandomState(self.random_state)

        # checks for group
        if self.group is not None:
            self.group = np.array(self.group)
            self.group.dtype = np.int64
            # shape check
            if self.group.shape[0] != X.shape[1]:
                raise ValueError('group should be (n_features,)')
            # int check
            if not np.all([isinstance(g, np.int64) for g in self.group]):
                raise ValueError('all entries of group should be integers')

        # type check for data matrix
        if not isinstance(X, np.ndarray):
            raise ValueError('Input data should be of type ndarray (got %s).'
                             % type(X))

        n_features = X.shape[1]

        # Initialize parameters
        beta = np.zeros((n_features + 1,))
        if self.beta0_ is None and self.beta_ is None:
            beta[0] = 1 / (n_features + 1) * np.random.normal(0.0, 1.0, 1)
            beta[1:] = 1 / (n_features + 1) * \
                np.random.normal(0.0, 1.0, (n_features, ))
        else:
            beta[0] = self.beta0_
            beta[1:] = self.beta_

        logger.info('Lambda: %6.4f' % self.reg_lambda)

        tol = self.tol
        alpha = self.alpha
        reg_lambda = self.reg_lambda

        if self.solver == 'cdfast':
            ActiveSet = np.ones(n_features + 1)     # init active set
            z = beta[0] + np.dot(X, beta[1:])       # cache z

        # Initialize loss accumulators
        L, DL = list(), list()
        for t in range(0, self.max_iter):
            if self.solver == 'batch-gradient':
                grad = _grad_L2loss(self.distr,
                                    alpha, self.Tau,
                                    reg_lambda, X, y, self.eta,
                                    beta)

                beta = beta - self.learning_rate * grad
            elif self.solver == 'cdfast':
                beta, z = \
                    self._cdfast(X, y, z, ActiveSet, beta, reg_lambda)

            # Apply proximal operator
            beta[1:] = self._prox(beta[1:], reg_lambda * alpha)

            # Update active set
            if self.solver == 'cdfast':
                ActiveSet[beta == 0] = 0
                ActiveSet[0] = 1.

            # Compute and save loss
            L.append(_loss(self.distr, alpha, self.Tau, reg_lambda,
                           X, y, self.eta, self.group, beta))

            if t > 1:
                DL.append(L[-1] - L[-2])
                if np.abs(DL[-1] / L[-1]) < tol:
                    msg = ('\tConverged. Loss function:'
                           ' {0:.2f}').format(L[-1])
                    logger.info(msg)
                    msg = ('\tdL/L: {0:.6f}\n'.format(DL[-1] / L[-1]))
                    logger.info(msg)
                    break

        # Update the estimated variables
        self.beta0_ = beta[0]
        self.beta_ = beta[1:]
        self.ynull_ = np.mean(y)
        return self

    def predict(self, X):
        """Predict targets.

        Parameters
        ----------
        X : array
            Input data for prediction, of shape (n_samples, n_features)

        Returns
        -------
        yhat : array
            The predicted targets of shape (n_samples,)
        """
        if not isinstance(X, np.ndarray):
            raise ValueError('Input data should be of type ndarray (got %s).'
                             % type(X))

        yhat = _lmb(self.distr, self.beta0_,
                    self.beta_, X, self.eta)

        if self.distr == 'binomial':
            yhat = (yhat > 0.5).astype(int)
        yhat = np.asarray(yhat)
        return yhat

    def predict_proba(self, X):
        """Predict class probability for binomial.

        Parameters
        ----------
        X : array
            Input data for prediction, of shape (n_samples, n_features)

        Returns
        -------
        yhat : array
            The predicted targets of shape (n_samples,).

        Raises
        ------
        Works only for the binomial distribution.
        Raises error otherwise.

        """
        if self.distr != 'binomial':
            raise ValueError('This is only applicable for \
                              the binomial distribution.')

        if not isinstance(X, np.ndarray):
            raise ValueError('Input data should be of type ndarray (got %s).'
                             % type(X))

        yhat = _lmb(self.distr,
                    self.beta0_, self.beta_, X, self.eta)
        yhat = np.asarray(yhat)
        return yhat

    def fit_predict(self, X, y):
        """Fit the model and predict on the same data.

        Parameters
        ----------
        X : array
            The input data to fit and predict,
            of shape (n_samples, n_features)


        Returns
        -------
        yhat : array
            The predicted targets of shape (n_samples,).
        """
        return self.fit(X, y).predict(X)

    def score(self, X, y):
        """Score the model.

        Parameters
        ----------
        X : array
            The input data whose prediction will be scored,
            of shape (n_samples, n_features).
        y : array
            The true targets against which to score the predicted targets,
            of shape (n_samples,).

        Returns
        -------
        score: float
            The score metric
        """
        from . import metrics
        if self.score_metric not in ['deviance', 'pseudo_R2', 'accuracy']:
            raise ValueError('score_metric has to be one of' +
                             ' deviance or pseudo_R2')

        # If the model has not been fit it cannot be scored
        if self.ynull_ is None:
            raise ValueError('Model must be fit before ' +
                             'prediction can be scored')

        # For f1 as well
        if self.score_metric in ['accuracy']:
            if self.distr not in ['binomial', 'multinomial']:
                raise ValueError(self.score_metric +
                                 ' is only defined for binomial ' +
                                 'or multinomial distributions')

        y = y.ravel()

        if self.distr == 'binomial' and self.score_metric != 'accuracy':
            yhat = self.predict_proba(X)
        else:
            yhat = self.predict(X)

        # Check whether we have a list of estimators or a single estimator
        if self.score_metric == 'deviance':
            return metrics.deviance(y, yhat, self.ynull_, self.distr)
        elif self.score_metric == 'pseudo_R2':
            return metrics.pseudo_R2(X, y, yhat, self.ynull_, self.distr)
        if self.score_metric == 'accuracy':
            return metrics.accuracy(y, yhat)


class GLMCV(object):
    """Class for estimating regularized generalized linear models (GLM)
    along a regularization path with warm restarts.

    The regularized GLM minimizes the penalized negative log likelihood:

    .. math::

        \min_{\\beta_0, \\beta} \\frac{1}{N}
        \sum_{i = 1}^N \mathcal{L} (y_i, \\beta_0 + \\beta^T x_i)
        + \lambda [ \\frac{1}{2}(1 - \\alpha) \mathcal{P}_2 +
                    \\alpha \mathcal{P}_1 ]

    where :math:`\mathcal{P}_2` and :math:`\mathcal{P}_1` are the generalized
    L2 (Tikhonov) and generalized L1 (Group Lasso) penalties, given by:

    .. math::

        \mathcal{P}_2 = \|\Gamma \\beta \|_2^2 \\
        \mathcal{P}_1 = \sum_g \|\\beta_{j,g}\|_2

    where :math:`\Gamma` is the Tikhonov matrix: a square factorization
    of the inverse covariance matrix and :math:`\\beta_{j,g}` is the
    :math:`j` th coefficient of group :math:`g`.

    The generalized L2 penalty defaults to the ridge penalty when
    :math:`\Gamma` is identity.

    The generalized L1 penalty defaults to the lasso penalty when each
    :math:`\\beta` belongs to its own group.

    Parameters
    ----------
    distr : str
        distribution family can be one of the following
        'gaussian' | 'binomial' | 'poisson' | 'softplus'
        default: 'poisson'.
    alpha : float
        the weighting between L1 penalty and L2 penalty term
        of the loss function.
        default: 0.5
    Tau : array | None
        the (n_features, n_features) Tikhonov matrix.
        default: None, in which case Tau is identity
        and the L2 penalty is ridge-like
    group : array | list | None
        the (n_features, )
        list or array of group identities for each parameter :math:`\\beta`.
        Each entry of the list/ array should contain an int from 1 to n_groups
        that specify group membership for each parameter
        (except :math:`\\beta_0`).
        If you do not want to specify a group for a specific parameter,
        set it to zero.
        default: None, in which case it defaults to L1 regularization
    reg_lambda : array | list | None
        array of regularized parameters :math:`\\lambda` of penalty term.
        default: None, a list of 10 floats spaced logarithmically (base e)
        between 0.5 and 0.01.
    cv : cross validation object (default 10)
        Iterator for doing cross validation
    solver : str
        optimization method, can be one of the following
        'batch-gradient' (vanilla batch gradient descent)
        'cdfast' (Newton coordinate gradient descent).
        default: 'batch-gradient'
    learning_rate : float
        learning rate for gradient descent.
        default: 2e-1
    max_iter : int
        maximum iterations for the model.
        default: 1000
    tol : float
        convergence threshold or stopping criteria.
        Optimization loop will stop below setting threshold.
        default: 1e-3
    eta : float
        a threshold parameter that linearizes the exp() function above eta.
        default: 2.0
    score_metric : str
        specifies the scoring metric.
        one of either 'deviance' or 'pseudo_R2'.
        default: 'deviance'
    random_state : int
        seed of the random number generator used to initialize the solution.
        default: 0
    verbose : boolean or int
        default: False

    Reference
    ---------
    Friedman, Hastie, Tibshirani (2010). Regularization Paths for Generalized
        Linear Models via Coordinate Descent, J Statistical Software.
        https://core.ac.uk/download/files/153/6287975.pdf

    Notes
    -----
    To select subset of fitted glm models, you can simply do:

    >>> glm = glm[1:3]
    >>> glm[2].predict(X_test)
    """

    def __init__(self, distr='poisson', alpha=0.5,
                 Tau=None, group=None,
                 reg_lambda=None, cv=10,
                 solver='batch-gradient',
                 learning_rate=2e-1, max_iter=1000,
                 tol=1e-3, eta=2.0, score_metric='deviance',
                 random_state=0, verbose=False):

        if reg_lambda is None:
            reg_lambda = np.logspace(np.log(0.5), np.log(0.01), 10,
                                     base=np.exp(1))
        if not isinstance(reg_lambda, (list, np.ndarray)):
            reg_lambda = [reg_lambda]

        if not isinstance(max_iter, int):
            raise ValueError('max_iter must be of type int')

        self.distr = distr
        self.alpha = alpha
        self.reg_lambda = reg_lambda
        self.cv = cv
        self.Tau = Tau
        self.group = group
        self.solver = solver
        self.learning_rate = learning_rate
        self.max_iter = max_iter
        self.beta0_ = None
        self.beta_ = None
        self.reg_lambda_opt_ = None
        self.glm_ = None
        self.scores_ = None
        self.ynull_ = None
        self.tol = tol
        self.eta = eta
        self.score_metric = score_metric
        self.random_state = random_state
        self.verbose = verbose
        set_log_level(verbose)

    def __repr__(self):
        """Description of the object."""
        reg_lambda = self.reg_lambda
        s = '<\nDistribution | %s' % self.distr
        s += '\nalpha | %0.2f' % self.alpha
        s += '\nmax_iter | %0.2f' % self.max_iter
        if len(reg_lambda) > 1:
            s += ('\nlambda: %0.2f to %0.2f\n>'
                  % (reg_lambda[0], reg_lambda[-1]))
        else:
            s += '\nlambda: %0.2f\n>' % reg_lambda[0]
        return s

    def copy(self):
        """Return a copy of the object.

        Parameters
        ----------
        none:

        Returns
        -------
        self: instance of GLM
            A copy of the GLM instance.
        """
        return deepcopy(self)

    def fit(self, X, y):
        """The fit function.
        Parameters
        ----------
        X : array
            The input data of shape (n_samples, n_features)

        y : array
            The target data

        Returns
        -------
        self : instance of GLM
            The fitted model.
        """
        logger.info('Looping through the regularization path')
        glms, scores = list(), list()
        self.ynull_ = np.mean(y)

        if not type(int):
            raise ValueError('cv must be int. We do not support scikit-learn '
                             'cv objects at the moment')

        idxs = np.arange(y.shape[0])
        np.random.shuffle(idxs)
        cv_splits = np.array_split(idxs, self.cv)

        glm = GLM(distr=self.distr,
                  alpha=self.alpha,
                  Tau=self.Tau,
                  reg_lambda=0.1,
                  solver=self.solver,
                  learning_rate=self.learning_rate,
                  max_iter=self.max_iter,
                  tol=self.tol,
                  eta=self.eta,
                  score_metric=self.score_metric,
                  random_state=self.random_state,
                  verbose=self.verbose)

        for idx, rl in enumerate(self.reg_lambda):
            logger.info('Lambda: %6.4f' % rl)

            glm.reg_lambda = rl

            scores_fold = list()
            for fold in range(self.cv):
                val = cv_splits[fold]
                train = np.setdiff1d(idxs, val)
                if idx == 0:
                    glm.beta0_, glm.beta_ = self.beta0_, self.beta_
                else:
                    glm.beta0_, glm.beta_ = glms[-1].beta0_, glms[-1].beta_

                glm.fit(X[train], y[train])
                scores_fold.append(glm.score(X[val], y[val]))
            scores.append(np.mean(scores_fold))

            if idx == 0:
                glm.beta0_, glm.beta_ = self.beta0_, self.beta_
            else:
                glm.beta0_, glm.beta_ = glms[-1].beta0_, glms[-1].beta_
            glm.fit(X, y)
            glms.append(glm)
        # Update the estimated variables
        if self.score_metric == 'deviance':
            opt = np.array(scores).argmin()
        elif self.score_metric in ['pseudo_R2', 'accuracy']:
            opt = np.array(scores).argmax()
        self.beta0_, self.beta_ = glms[opt].beta0_, glms[opt].beta_
        self.reg_lambda_opt_ = self.reg_lambda[opt]
        self.glm_ = glms[opt]
        self.scores_ = scores
        return self

    def predict(self, X):
        """Predict targets.

        Parameters
        ----------
        X : array
            Input data for prediction, of shape (n_samples, n_features)

        Returns
        -------
        yhat : array
            The predicted targets of shape based on the model with optimal
            reg_lambda (n_samples,)
        """
        return self.glm_.predict(X)

    def predict_proba(self, X):
        """Predict class probability for binomial.

        Parameters
        ----------
        X : array
            Input data for prediction, of shape (n_samples, n_features)

        Returns
        -------
        yhat : array
            The predicted targets of shape (n_samples, ).

        Raises
        ------
        Works only for the binomial distribution.
        Raises error otherwise.

        """
        return self.glm_.predict_proba(X)

    def fit_predict(self, X, y):
        """Fit the model and predict on the same data.

        Parameters
        ----------
        X : array
            The input data to fit and predict,
            of shape (n_samples, n_features)

        Returns
        -------
        yhat : array
            The predicted targets of shape based on the model with optimal
            reg_lambda (n_samples,)
        """
        self.fit(X, y)
        return self.glm_.predict(X)

    def score(self, X, y):
        """Score the model.

        Parameters
        ----------
        X : array
            The input data whose prediction will be scored,
            of shape (n_samples, n_features).
        y : array
            The true targets against which to score the predicted targets,
            of shape (n_samples,).

        Returns
        -------
        score: float
            The score metric for the optimal reg_lambda
        """
        return self.glm_.score(X, y)
