"""
Regular Transfer
"""

import numpy as np
from sklearn.preprocessing import LabelBinarizer
from scipy.sparse.linalg import lsqr
import tensorflow as tf
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Flatten, Dense

from adapt.base import BaseAdaptEstimator, BaseAdaptDeep, make_insert_doc
from adapt.utils import (check_arrays,
                         set_random_seed,
                         check_estimator,
                         check_network,
                         check_fitted_estimator)


@make_insert_doc(supervised=True)
class RegularTransferLR(BaseAdaptEstimator):
    """
    Regular Transfer with Linear Regression
    
    RegularTransferLR is a parameter-based domain adaptation method.
    
    The method is based on the assumption that a good target estimator
    can be obtained by adapting the parameters of a pre-trained source
    estimator using a few labeled target data.
    
    The approach consist in fitting a linear estimator on target data
    according to an objective function regularized by the euclidean
    distance between source and target parameters:
    
    .. math::
    
        \\beta_T = \\underset{\\beta \in \\mathbb{R}^p}{\\text{argmin}}
        \\, ||X_T\\beta - y_T||^2 + \\lambda ||\\beta - \\beta_S||^2
        
    Where:
    
    - :math:`\\beta_T` are the target model parameters.
    - :math:`\\beta_S = \\underset{\\beta \\in \\mathbb{R}^p}{\\text{argmin}}
      \\, ||X_S\\beta - y_S||^2` are the source model parameters.
    - :math:`(X_S, y_S), (X_T, y_T)` are respectively the source and
      the target labeled data.
    - :math:`p` is the number of features in :math:`X_T`
      (:math:`+1` if ``intercept`` is True).
    - :math:`\\lambda` is a trade-off parameter.

    Parameters
    ----------        
    lambda_ : float (default=1.0)
        Trade-Off parameter.

    Attributes
    ----------
    estimator_ : Same class as estimator
        Fitted Estimator.
        
    Examples
    --------
    >>> import numpy as np
    >>> from adapt.parameter_based import RegularTransferLR
    >>> from sklearn.linear_model import LinearRegression
    >>> np.random.seed(0)
    >>> Xs = np.random.randn(50) * 0.1
    >>> Xs = np.concatenate((Xs, Xs + 1.))
    >>> Xt = np.random.randn(100) * 0.1
    >>> ys = (np.array([-0.2 * x if x<0.5 else 1. for x in Xs])
    ...       + 0.1 * np.random.randn(100))
    >>> yt = 0.75 * Xt + 0.1 * np.random.randn(100)
    >>> lr = LinearRegression()
    >>> lr.fit(Xs.reshape(-1, 1), ys)
    >>> lr.score(Xt.reshape(-1, 1), yt)
    0.2912...
    >>> rt = RegularTransferLR(lr, lambda_=0.01, random_state=0)
    >>> rt.fit(Xt[:10].reshape(-1, 1), yt[:10])
    >>> rt.estimator_.score(Xt.reshape(-1, 1), yt)
    0.3276...
        
    See also
    --------
    RegularTransferLC, RegularTransferNN

    References
    ----------
    .. [1] `[1] <https://www.microsoft.com/en-us/research/wp-\
content/uploads/2004/07/2004-chelba-emnlp.pdf>`_ C. Chelba and \
A. Acero. "Adaptation of maximum entropy classifier: Little data \
can help a lot". In EMNLP, 2004.
    """
    def __init__(self,
                 estimator=None,
                 Xt=None,
                 yt=None,
                 lambda_=1.,
                 copy=True,
                 verbose=1,
                 random_state=None,
                 **params):
                
        if not hasattr(estimator, "coef_"):
            raise ValueError("`estimator` argument has no ``coef_`` attribute, "
                             "please call `fit` on `estimator` or use "
                             "another estimator as `LinearRegression` or "
                             "`RidgeClassifier`.")
            
        estimator = check_fitted_estimator(estimator)

        names = self._get_param_names()
        kwargs = {k: v for k, v in locals().items() if k in names}
        kwargs.update(params)
        super().__init__(**kwargs)

    
    def fit(self, Xt=None, yt=None, **fit_params):
        """
        Fit RegularTransferLR.

        Parameters
        ----------
        Xt : numpy array (default=None)
            Target input data.

        yt : numpy array (default=None)
            Target output data.
            
        fit_params : key, value arguments
            Not used. Here for sklearn compatibility.

        Returns
        -------
        self : returns an instance of self
        """        
        Xt, yt = self._get_target_data(Xt, yt)
        Xt, yt = check_arrays(Xt, yt)
        set_random_seed(self.random_state)
        
        self.estimator_ = check_estimator(self.estimator,
                                          copy=self.copy,
                                          force_copy=True)

        if self.estimator_.fit_intercept:
            intercept_ = np.reshape(
                self.estimator_.intercept_,
                np.ones(self.estimator_.coef_.shape).mean(-1, keepdims=True).shape)
            beta_src = np.concatenate((
                intercept_,
                self.estimator_.coef_
            ), axis=-1)
            Xt = np.concatenate(
                (np.ones((len(Xt), 1)), Xt),
                axis=-1)
        else:
            beta_src = self.estimator_.coef_
        
        yt_ndim_below_one_ = False
        if yt.ndim <= 1:
            yt = yt.reshape(-1, 1)
            yt_ndim_below_one_ = True
        
        if beta_src.ndim <= 1:
            beta_src = beta_src.reshape(1, -1)
            
        if beta_src.shape[0] != yt.shape[1]:
            raise ValueError("The number of features of `yt`"
                             " does not match the number of coefs in 'estimator', "
                             "expected %i, got %i"%(beta_src.shape[0], yt.shape[1]))
        
        if beta_src.shape[1] != Xt.shape[1]:            
            beta_shape = beta_src.shape[1]; Xt_shape = Xt.shape[1]
            if self.estimator_.fit_intercept:
                beta_shape -= 1; Xt_shape -= 1
            raise ValueError("The number of features of `Xt`"
                             " does not match the number of coefs in 'estimator', "
                             "expected %i, got %i"%(beta_shape, Xt_shape))
            
        beta_tgt = []
        for i in range(yt.shape[1]):
            sol = lsqr(A=Xt, b=yt[:, i], damp=self.lambda_, x0=beta_src[i, :])
            beta_tgt.append(sol[0])
            
        beta_tgt = np.stack(beta_tgt, axis=0)
        
        if self.estimator_.fit_intercept:
            self.coef_ = beta_tgt[:, 1:]
            self.intercept_ = beta_tgt[:, 0]
        else:
            self.coef_ = beta_tgt
            
        if yt_ndim_below_one_:
            self.coef_ = self.coef_.reshape(-1)
            self.intercept_ = self.intercept_[0]
            
        self.estimator_.coef_ = self.coef_
        if self.estimator_.fit_intercept:
            self.estimator_.intercept_ = self.intercept_
        return self

    

@make_insert_doc(supervised=True)
class RegularTransferLC(RegularTransferLR):
    """
    Regular Transfer for Linear Classification
    
    RegularTransferLC is a parameter-based domain adaptation method.
        
    This classifier first converts the target values into ``{-1, 1}``
    and then treats the problem as a regression task
    (multi-output regression in the multiclass case). It then fits
    the target data as a ``RegularTransferLR`` regressor, i.e it
    performs the following optimization:
    
    .. math::
    
        \\beta_T = \\underset{\\beta \in \\mathbb{R}^p}{\\text{argmin}}
        \\, ||X_T\\beta - y_T||^2 + \\lambda ||\\beta - \\beta_S||^2
        
    Where:
    
    - :math:`\\beta_T` are the target model parameters.
    - :math:`\\beta_S = \\underset{\\beta \\in \\mathbb{R}^p}{\\text{argmin}}
      \\, ||X_S\\beta - y_S||^2` are the source model parameters.
    - :math:`(X_S, y_S), (X_T, y_T)` are respectively the source and
      the target labeled data.
    - :math:`p` is the number of features in :math:`X_T`
      (:math:`+1` if ``intercept`` is True).
    - :math:`\\lambda` is a trade-off parameter.

    Parameters
    ----------        
    lambda_ : float (default=1.0)
        Trade-Off parameter.

    Attributes
    ----------
    estimator_ : Same class as estimator
        Fitted Estimator.
            
    Examples
    --------
    >>> import numpy as np
    >>> from adapt.parameter_based import RegularTransferLC
    >>> from sklearn.linear_model import LogisticRegression
    >>> np.random.seed(0)
    >>> Xs = np.random.randn(50) * 0.1
    >>> Xs = np.concatenate((Xs, Xs + 1.))
    >>> Xt = np.random.randn(100) * 0.1
    >>> ys = (Xs < 0.1).astype(int)
    >>> yt = (Xt < 0.05).astype(int)
    >>> lc = LogisticRegression()
    >>> lc.fit(Xs.reshape(-1, 1), ys)
    >>> lc.score(Xt.reshape(-1, 1), yt)
    0.67
    >>> rt = RegularTransferLC(lc, lambda_=0.01, random_state=0)
    >>> rt.fit(Xt[:10].reshape(-1, 1), yt[:10].reshape(-1, 1))
    >>> rt.estimator_.score(Xt.reshape(-1, 1), yt)
    0.67

    See also
    --------
    RegularTransferLR, RegularTransferNN

    References
    ----------
    .. [1] `[1] <https://www.microsoft.com/en-us/research/wp-\
content/uploads/2004/07/2004-chelba-emnlp.pdf>`_ C. Chelba and \
A. Acero. "Adaptation of maximum entropy classifier: Little data \
can help a lot". In EMNLP, 2004.
    """
    ### TODO reshape yt for multiclass.
    
    def fit(self, Xt=None, yt=None, **fit_params):       
        Xt, yt = self._get_target_data(Xt, yt)
        Xt, yt = check_arrays(Xt, yt)
        
        _label_binarizer = LabelBinarizer(pos_label=1, neg_label=-1)
        yt = _label_binarizer.fit_transform(yt)
        return super().fit(Xt, yt, **fit_params)


@make_insert_doc(["task"], supervised=True)
class RegularTransferNN(BaseAdaptDeep):
    """
    Regular Transfer with Neural Network
    
    RegularTransferNN is a parameter-based domain adaptation method.
    
    The method is based on the assumption that a good target estimator
    can be obtained by adapting the parameters of a pre-trained source
    estimator using a few labeled target data.
    
    The approach consist in fitting a neural network on target data
    according to an objective function regularized by the euclidean
    distance between source and target parameters:
    
    .. math::
    
        \\beta_T = \\underset{\\beta=(\\beta_1, ... , \\beta_D)}{\\text{argmin}}
        \\, ||f(X_T, \\beta) - y_T||^2 + \sum_{i=1}^{D}
        \\lambda_i ||\\beta_i - {\\beta_S}_i||^2
        
    Where:
    
    - :math:`f` is a neural network with :math:`D` layers.
    - :math:`\\beta_T` are the parameters of the target neural network.
    - :math:`\\beta_S = \\underset{\\beta}{\\text{argmin}}
      \\, ||f(X_S,\\beta) - y_S||^2` are the source neural network parameters.
    - :math:`(X_S, y_S), (X_T, y_T)` are respectively the source and
      the target labeled data.
    - :math:`\\lambda_i` is the trade-off parameter of layer :math:`i`.
    
    Different trade-off can be given to the layer of the 
    neural network through the ``lambdas`` parameter.

    Parameters
    ----------        
    lambdas : float or list of float, optional (default=1.0)
        Trade-off parameters.
        If a list is given, values from ``lambdas`` are assigned
        successively to the list of ``network`` layers with 
        weights parameters going from the last layer to the first one.
        If the length of ``lambdas`` is smaller than the length of
        ``network`` layers list, the last trade-off value will be
        asigned to the remaining layers.

    Attributes
    ----------
    task_ : tensorflow Model
        Network.
        
    history_ : dict
        history of the losses and metrics across the epochs
        of the network training.
        
    Examples
    --------
    >>> import numpy as np
    >>> import tensorflow as tf
    >>> from adapt.parameter_based import RegularTransferNN
    >>> np.random.seed(0)
    >>> tf.random.set_seed(0)
    >>> Xs = np.random.randn(50) * 0.1
    >>> Xs = np.concatenate((Xs, Xs + 1.))
    >>> Xt = np.random.randn(100) * 0.1
    >>> ys = (np.array([-0.2 * x if x<0.5 else 1. for x in Xs])
    ...       + 0.1 * np.random.randn(100))
    >>> yt = 0.75 * Xt + 0.1 * np.random.randn(100)
    >>> model = tf.keras.Sequential()
    >>> model.add(tf.keras.layers.Dense(1))
    >>> model.compile(optimizer="adam", loss="mse")
    >>> model.predict(Xt.reshape(-1,1))
    >>> model.fit(Xs.reshape(-1, 1), ys, epochs=300, verbose=0)
    >>> np.abs(model.predict(Xt).ravel() - yt).mean()
    0.48265...
    >>> rt = RegularTransferNN(model, lambdas=0.01, random_state=0)
    >>> rt.fit(Xt[:10].reshape(-1,1), yt[:10], epochs=300, verbose=0)
    >>> np.abs(rt.predict(Xt.reshape(-1,1)).ravel() - yt).mean()
    0.1900...
        
    See also
    --------
    RegularTransferLR, RegularTransferLC

    References
    ----------
    .. [1] `[1] <https://www.microsoft.com/en-us/research/wp-\
content/uploads/2004/07/2004-chelba-emnlp.pdf>`_ C. Chelba and \
A. Acero. "Adaptation of maximum entropy classifier: Little data \
can help a lot". In EMNLP, 2004.
    """
    def __init__(self,
                 task=None,
                 Xt=None,
                 yt=None,
                 lambdas=1.0,
                 regularizer="l2",
                 verbose=1,
                 copy=True,
                 random_state=None,
                 **params):
        
        if not regularizer in ["l1", "l2"]:
            raise ValueError("`regularizer` argument should be "
                             "'l1' or 'l2', got, %s"%str(regularizer))
        
        names = self._get_param_names()
        kwargs = {k: v for k, v in locals().items() if k in names}
        kwargs.update(params)
        super().__init__(**kwargs)


    def fit(self, Xt=None, yt=None, **fit_params):
        """
        Fit RegularTransferNN.

        Parameters
        ----------
        Xt : numpy array (default=None)
            Target input data.

        yt : numpy array (default=None)
            Target output data.
            
        fit_params : key, value arguments
            Arguments given to the fit method of the model
            (epochs, batch_size, callbacks...).

        Returns
        -------
        self : returns an instance of self
        """        
        Xt, yt = self._get_target_data(Xt, yt)
        Xs = Xt
        ys = yt
        return super().fit(Xs, ys, Xt=Xt, yt=yt, **fit_params)
    
    
    def _initialize_networks(self):
        if self.task is None:
            self.task_ = get_zeros_network(name="task")
        else:
            self.task_ = check_network(self.task,
                                       copy=self.copy,
                                       name="task")
        self._add_regularization()
    
    
    def _get_regularizer(self, old_weight, weight, lambda_=1.):
        if self.regularizer == "l2":
            def regularizer():
                return lambda_ * tf.reduce_mean(tf.square(old_weight - weight))
        if self.regularizer == "l1":
            def regularizer():
                return lambda_ * tf.reduce_mean(tf.abs(old_weight - weight))
        return regularizer


    def _add_regularization(self):
        i = 0
        if not hasattr(self.lambdas, "__iter__"):
            lambdas = [self.lambdas]
        else:
            lambdas = self.lambdas
        
        for layer in reversed(self.task_.layers):
            if (hasattr(layer, "weights") and 
            layer.weights is not None and
            len(layer.weights) != 0):
                if i >= len(lambdas):
                    lambda_ = lambdas[-1]
                else:
                    lambda_ = lambdas[i]
                for weight in reversed(layer.weights):
                    old_weight = tf.identity(weight)
                    old_weight.trainable = False
                    self.add_loss(self._get_regularizer(
                        old_weight, weight, lambda_))
                i += 1
        
        
    def call(self, inputs):
        return self.task_(inputs)
    
    
    def transform(self, X):
        """
        Return X
        
        Parameters
        ----------
        X : array
            input data
            
        Returns
        -------
        X_enc : array
            predictions of encoder network
        """
        return X
    
    
    def predict_disc(self, X):
        """
        Not used.
        """     
        pass