"""
:mod:`operalib.quantile` implements joint quantile regression with
Operator-Valued Kernels.
"""
# Authors: Maxime Sangnier <maxime.sangnier@gmail.com>
#          Romain Brault <romain.brault@telecom-paristech.fr> with help from
#          the scikit-learn community.
# License: MIT
from numpy import (asarray, eye, reshape, argsort, kron, ones, diag, delete,
                   zeros, r_, c_, percentile, fmax)
from numpy import inf as npinf
from numpy import sum as npsum

from cvxopt import matrix, solvers

from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.utils import check_X_y, check_array
from sklearn.utils.validation import check_is_fitted, check_consistent_length
from sklearn.metrics.pairwise import rbf_kernel

# from operalib. .kernels import DecomposableKernel
from operalib.kernels import DecomposableKernel

# When adding a new kernel, update this table and the _get_kernel_map method
PAIRWISE_KERNEL_FUNCTIONS = {
    'DGauss': DecomposableKernel,
    'DPeriodic': DecomposableKernel, }

# pylint: disable=C0103,W0201,R0902,R0913


class Quantile(BaseEstimator, RegressorMixin):
    """Joint quantile regression with operator-valued kernels.
    Joint quantile regression enables to learn and predict simultaneously
    several conditional quantiles (for prescribed quantile levels). Following
    the seminal work by Koenker and Basset (1978), the method minimizes an
    empirical risk based on the pinball loss and a squared-norm regularization.
    The hypothesis space considered here is the reproducing kernel Hilbert
    space generated by an operator-valued kernel chosen by the practitioner.

    Attributes
    ----------
    model : dict
        coefs : {array-like}, shape [n_samples*n_probs, 1]
            Dual coefficients of the kernel machine
        intercept : {array-like}, shape [n_probs]
            Vector of intercepts

    References
    ----------
    * R. Koenker and G. Bassett
      "Regression Quantiles", Econometrica, 46(1):33-50, 1978.
    * R. Koenker
      "Quantile Regression", Cambridge University Press, Cambridge, New York,
      2005.
    * M. Sangnier, O. Fercoq and F. d'Alche-Buc
      "Joint quantile regression in vector-valued RKHSs." In Advances in Neural
    See also
    --------

    operalib.Ridge
        Operator-Valued kernel ridge regression
    sklearn.Ridge
        Linear ridge regression.
    sklearn.KernelRidge
        Kernel ridge regression.
    sklearn.SVR
        Support Vector Regression implemented using libsvm.

    Examples
    --------
    >>> import operalib as ovk
    >>> import numpy as np
    >>> n_samples, n_features = 10, 5
    >>> quantile_levels = [0.1, 0.5, 0.9]
    >>> rng = np.random.RandomState(0)
    >>> y = rng.randn(n_samples)
    >>> X = rng.randn(n_samples, n_features)
    >>> reg = ovk.Quantile('DGauss', lbda=1.0)
    >>> reg.fit(X, y)  # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS
    Quantile(gamma=None, gamma_quantile=0.0, kernel='DGauss',
        kernel_params=None, lbda=1.0, nc_const=False, probs=0.5, tol=None,
        verbose=False)
    """

    def __init__(self, kernel='DGauss', probs=0.5, lbda=1e-5, gamma=None,
                 gamma_quantile=0., eps=0., tol=None, nc_const=False,
                 kernel_params=None, verbose=False):
        """Initialize quantile regression model.

        Parameters
        ----------
        probs : {float, list}, default=0.5
            Probabilities (quantiles levels).
        kernel : {string, callable}, default='DGauss'
            Kernel mapping used internally. A callable should accept two
            arguments and the keyword arguments passed to this object as
            kernel_params, and should return a LinearOperator.
        lbda : {float}, default=1e-5
            Regularization parameter.  Lbda corresponds to ``C^-1`` in other
            large margin models such as SVC.
        gamma : {float}, default=None.
            Gamma parameter for the Decomposable Gaussian kernel.
            Ignored by other kernels.
        gamma_quantile : {float}, default=None.
            Gamma parameter for the output Gaussian kernel.
        eps : {float}, default=0
            Threshold for the epsilon-loss (data sparse regression, only
            without non-crossing constraints)
        tol : {float}, default=None
            Optimization tolerance (None leads to the default value of the
            solver).
        nc_const : {Boolean}, default=False
            Whether to enforce non-crossing constraints when optimizing
            (default is no).
        kernel_params : {mapping of string to any}, optional
            Additional parameters (keyword arguments) for kernel function
            passed as callable object.
        verbose : {Boolean}, default=False
            Verbosity
        """
        self.probs = probs
        self.lbda = lbda
        self.kernel = kernel
        self.gamma = gamma
        self.gamma_quantile = gamma_quantile
        self.eps = eps
        self.tol = tol
        self.nc_const = nc_const
        self.kernel_params = kernel_params
        self.verbose = verbose

        # self.linop_ = None
        # self.reg_c_ = None
        # self.sol_ = None
        # self.model_ = None

    def _validate_params(self):
        # check on self.kernel is performed in method __get_kernel
        if self.lbda < 0:
            raise ValueError('lbda must be positive')
        # if self.A < 0: # Check whether A is S PD would be really expensive
        #     raise ValueError('A must be a symmetric positive operator')
        if self.gamma is not None:
            if self.gamma < 0:
                raise ValueError('sigma must be positive or default (None)')
        # probs = asarray(self.probs).reshape(-1, 1)
        probs = asarray(self.probs).flatten()
        if (probs < 0).any() or (probs > 1).any():
            raise ValueError('Probabilities must be in [0., 1.]')

    def _default_decomposable_op(self):
        probs = asarray(self.probs).reshape((1, -1))  # 2D array
        return (rbf_kernel(probs.T, gamma=self.gamma_quantile)
                if self.gamma_quantile != npinf else eye(len(self.probs)))

    def _get_kernel_map(self, inputs):
        # When adding a new kernel, update this table and the _get_kernel_map
        # method
        if callable(self.kernel):
            kernel_params = self.kernel_params or {}
            ov_kernel = self.kernel(**kernel_params)
        elif isinstance(self.kernel, str):
            # 1) check string and assign the right parameters
            if self.kernel == 'DGauss':
                kernel_params = {'A': self._default_decomposable_op(),
                                 'scalar_kernel': rbf_kernel,
                                 'scalar_kernel_params': {'gamma': self.gamma}}
            else:
                raise NotImplementedError('unsupported kernel')
            # 2) Uses lookup table to select the right kernel from string
            ov_kernel = PAIRWISE_KERNEL_FUNCTIONS[self.kernel](**kernel_params)
        else:
            raise NotImplementedError('unsupported kernel')
        return ov_kernel(inputs)

    def _decision_function(self, X):
        n_samples = X.shape[0]
        n_quantiles = asarray(self.probs).reshape((-1, 1)).size

        pred = reshape(self.linop_(X) * self.model_['coefs'],
                       (n_samples, n_quantiles))
        pred += asarray(self.model_['intercept']).squeeze()

        return pred.T if self.linop_.p > 1 else pred.T.ravel()

    def predict(self, X):
        """Predict conditional quantiles.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = [n_samples, n_features]
            Samples.

        Returns
        -------
        y : {array}, shape = [n_samples, n_quantiles]
            Returns predicted values for each prescribed quantile level.
        """
        check_is_fitted(self, ['model_', 'linop_'], all_or_any=all)
        X = check_array(X)
        y = self._decision_function(X)
        y[y < 1e-9] = 0 # Stability hack
        return y

    def fit(self, X, y):
        """Fit joint quantile regression model.

        Parameters
        ----------
        inputs : {array-like, sparse matrix}, shape = [n_samples, n_features]
            Training data.
        targets : {array-like}, shape = [n_samples]
            Target values.

        Returns
        -------
        self : returns an instance of self.
        """
        if self.eps > 0 and self.nc_const:
            raise UserWarning("eps is considered null because you chose to "
                              "enfoce non-crossing constraints.")
        X, y = check_X_y(X, y, ['csr', 'csc', 'coo'], y_numeric=True)
        y = asarray(y).flatten()
        self._validate_params()

        self.linop_ = self._get_kernel_map(X)
        gram = self.linop_.Gram_dense(X)
        self.reg_c_ = 1. / self.lbda

        # Solve the optimization problem
        # probs = asarray(self.probs).reshape((-1, 1))
        probs = asarray(self.probs).flatten()
        if self.nc_const:
            self._qp_nc(gram, y, probs)
        else:
            self._coneqp(gram, y, probs)
        return self

    def _qp_nc(self, gram, targets, probs):
        # Needed to sort constraints on quantiles levels
        ind = argsort(probs.T).squeeze()

        n_quantiles = probs.size  # Number of quantiles to predict
        n_coefs = gram.shape[0]  # Number of coefficients
        n_train = int(n_coefs / n_quantiles)  # Number of training instances
        # Number of non-crossing dual variables
        n_no_cross = n_train * (n_quantiles - 1)
        probs = kron(ones(n_train), probs.squeeze())  # Quantiles levels

        # Difference matrix
        diff_mat = -eye(n_quantiles) + diag(ones(n_quantiles - 1), 1)
        diff_mat = delete(diff_mat, -1, 0)
        diff_mat = diff_mat.T[ind].T
        # Quad. part of the objective function
        gram = matrix(r_[c_[gram, zeros((n_coefs, n_no_cross))],
                         zeros((n_no_cross, n_coefs + n_no_cross))])
        # Linear part of the objective
        q_lin = matrix(r_[-kron(targets, ones(n_quantiles)),
                          zeros(n_no_cross)])
        # LHS of the inequality constraint
        g_lhs = matrix(r_[c_[eye(n_coefs), -kron(eye(n_train), diff_mat.T)],
                          c_[-eye(n_coefs), kron(eye(n_train), diff_mat.T)],
                          c_[zeros((n_no_cross, n_coefs)),
                             -eye(n_no_cross)]])
        # RHS of the inequality constraint
        h_rhs = matrix(r_[self.reg_c_ * probs,
                          self.reg_c_ * (1 - probs),
                          zeros(n_train * (n_quantiles - 1))])
        # LHS of the equality constraint
        lhs_eqc = matrix(c_[kron(ones(n_train), eye(n_quantiles)),
                            zeros((n_quantiles, n_no_cross))])

        solvers.options['show_progress'] = self.verbose
        if self.tol:
            solvers.options['reltol'] = self.tol
        # Solve the dual optimization problem
        self.sol_ = solvers.qp(gram, q_lin, g_lhs, h_rhs, lhs_eqc,
                               matrix(zeros(n_quantiles)))

        self.model_ = {'coefs': asarray(self.sol_['x'][:n_coefs]),
                       'intercept': asarray(self.sol_['y']).squeeze()}

    def _qp(self, gram, targets, probs):
        n_quantiles = probs.size  # Number of quantiles to predict
        n_coefs = gram.shape[0]  # Number of variables
        # Quantiles levels
        kronprobs = kron(ones(int(n_coefs / n_quantiles)), probs.squeeze())

        # Quadratic part of the objective
        gram = matrix(gram)
        # Linear part of the objective
        q_lin = matrix(-kron(targets, ones(n_quantiles)))

        # LHS of the inequality constraint
        g_lhs = matrix(r_[eye(n_coefs), -eye(n_coefs)])
        # RHS of the inequality
        h_rhs = matrix(r_[self.reg_c_ * kronprobs,
                          self.reg_c_ * (1 - kronprobs)])
        # LHS of the equality constraint
        lhs_eqc = matrix(kron(ones(int(n_coefs / n_quantiles)),
                              eye(n_quantiles)))

        solvers.options['show_progress'] = self.verbose
        if self.tol:
            solvers.options['reltol'] = self.tol
        # Solve the dual optimization problem
        self.sol_ = solvers.qp(gram, q_lin, g_lhs, h_rhs, lhs_eqc,
                               matrix(zeros(n_quantiles)))

        # Set coefs
        coefs = asarray(self.sol_['x'])

        # Set the intercept

        # Erase the previous intercept before prediction
        self.model_ = {'coefs': coefs, 'intercept': 0}
        predictions = self.predict(self.linop_.X)
        if predictions.ndim < 2:
            predictions = predictions.reshape(1, -1)  # 2D array
        intercept = [percentile(targets - pred, 100. * prob) for
                     (pred, prob) in zip(predictions, probs)]
        intercept = asarray(intercept).squeeze()
        self.model_ = {'coefs': coefs, 'intercept': intercept}

    def _coneqp(self, gram, targets, probs):
        n_quantiles = probs.size  # Number of quantiles to predict
        n_coefs = gram.shape[0]  # Number of variables
        n_samples = n_coefs // n_quantiles
        # Quantiles levels
        kronprobs = kron(ones(int(n_coefs / n_quantiles)), probs.squeeze())

        solvers.options['show_progress'] = self.verbose
        if self.tol:
            solvers.options['reltol'] = self.tol

        if self.eps == 0:
            # Quadratic part of the objective
            gram = matrix(gram)
            # Linear part of the objective
            q_lin = matrix(-kron(targets, ones(n_quantiles)))

            # LHS of the inequality constraint
            g_lhs = matrix(r_[eye(n_coefs), -eye(n_coefs)])
            # RHS of the inequality
            h_rhs = matrix(r_[self.reg_c_ * kronprobs,
                              self.reg_c_ * (1 - kronprobs)])
            # LHS of the equality constraint
            lhs_eqc = matrix(kron(ones(n_samples), eye(n_quantiles)))

            # Solve the dual optimization problem
            self.sol_ = solvers.qp(gram, q_lin, g_lhs, h_rhs, lhs_eqc,
                                   matrix(zeros(n_quantiles)))

            # Set coefs
            coefs = asarray(self.sol_['x'])
        else:
            def build_lhs(m, p):
                # m: n_samples
                # p: n_quantiles
                n = m*p  # n_variables

                # Get the norm bounds (m last variables)
                A = zeros(p+1)
                A[0] = -1
                A = kron(eye(m), A).T
                # Get the m p-long vectors
                B = kron(eye(m), c_[zeros(p), eye(p)].T)
                # Box constraint
                C = c_[r_[eye(n), -eye(n)], zeros((2*n, m))]
                # Set everything together
                C = r_[C, c_[B, A]]
                return C

            # Quadratic part of the objective
            gram = matrix(r_[c_[gram, zeros((n_coefs, n_samples))],
                             zeros((n_samples, n_coefs+n_samples))])
            # Linear part of the objective
            q_lin = matrix(r_[-kron(targets, ones(n_quantiles)),
                              ones(n_samples)*self.eps])

            # LHS of the inequality constraint
            g_lhs = matrix(build_lhs(n_samples, n_quantiles))
            # RHS of the inequality
            h_rhs = matrix(r_[self.reg_c_ * kronprobs,
                              self.reg_c_ * (1-kronprobs),
                              zeros(n_samples * (n_quantiles+1))])
            # LHS of the equality constraint
            lhs_eqc = matrix(c_[kron(ones(n_samples),
                                eye(n_quantiles)),
                                zeros((n_quantiles, n_samples))])
            # Parameters of the optimization problem
            dims = {'l': 2*n_coefs, 'q': [n_quantiles+1]*n_samples, 's': []}

            # Solve the dual optimization problem
            self.sol_ = solvers.coneqp(gram, q_lin, g_lhs, h_rhs, dims,
                                       lhs_eqc, matrix(zeros(n_quantiles)))

            # Set coefs
            coefs = asarray(self.sol_['x'][:n_coefs])

        # Set the intercept

        # Erase the previous intercept before prediction
        self.model_ = {'coefs': coefs, 'intercept': 0}
        predictions = self.predict(self.linop_.X)
        if predictions.ndim < 2:
            predictions = predictions.reshape(1, -1)  # 2D array
        intercept = [percentile(targets - pred, 100. * prob) for
                     (pred, prob) in zip(predictions, probs)]
        intercept = asarray(intercept).squeeze()
        self.model_ = {'coefs': coefs, 'intercept': intercept}

    @staticmethod
    def pinball_loss(y_true, y_pred, probs):
        """Compute the pinball loss.

        Parameters
        ----------
        pred : {array-like}, shape = [n_quantiles, n_samples] or [n_samples]
            Predictions.
        y : {array-like}, shape = [n_samples]
            Targets.

        Returns
        -------
        l : {array}, shape = [n_quantiles]
            Average loss for each quantile level.
        """
        probs = asarray(probs).reshape(-1)
        check_consistent_length(y_true, y_pred.T)
        y_true = check_array(y_true.reshape((-1, 1)),
                             ensure_2d=True)
        y_pred = check_array(y_pred.T.reshape((y_true.shape[0], -1)),
                             ensure_2d=True)
        residual = y_true - y_pred
        loss = npsum([fmax(prob * res, (prob - 1) * res) for (res, prob) in
                      zip(residual.T, probs)], axis=1)
        return loss / y_true.size

    def score(self, X, y, sample_weight=None):
        """Compute the pinball score for the given dataset.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = [n_samples, n_features]
            Training data.
        y : {array-like}, shape = [n_samples]
            Target values.

        Returns
        -------
        l : {float}
            Average pinball score (the higher, the better).
        """
        check_is_fitted(self, ['model_', 'linop_'], all_or_any=all)
        X, y = check_X_y(X, y)
        return 1 - Quantile.pinball_loss(y, self.predict(X), self.probs).mean()
