'''
Author: Dr. Mohamed A. Bouhlel <mbouhlel@umich.edu>

This package is distributed under New BSD license.
'''

import numpy as np
from sklearn.cross_decomposition.pls_ import PLSRegression as pls
from pyDOE2 import bbdesign

def standardization(X,y,copy=False):

    """
    We substract the mean from each variable. Then, we divide the values of each
    variable by its standard deviation.

    Parameters
    ----------

    X: np.ndarray [n_obs, dim]
            - The input variables.

    y: np.ndarray [n_obs, 1]
            - The output variable.

    copy: bool
            - A copy of matrices X and y will be used (copy = True).
            - Matrices X and y will be used. The matrices X and y will be
              normalized (copy = False).
            - (copy = False by default).

    Returns
    -------

    X: np.ndarray [n_obs, dim]
          The standardized input matrix.

    y: np.ndarray [n_obs, 1]
          The standardized output vector.

    X_mean: list(dim)
            The mean of each input variable.

    y_mean: list(1)
            The mean of the output variable.

    X_std:  list(dim)
            The standard deviation of each input variable.

    y_std:  list(1)
            The standard deviation of the output variable.

    """
    X_mean = np.mean(X, axis=0)
    X_std = X.std(axis=0,ddof=1)
    y_mean = np.mean(y, axis=0)
    y_std = y.std(axis=0,ddof=1)
    X_std[X_std == 0.] = 1.
    y_std[y_std == 0.] = 1.

    # center and scale X
    if copy:
        Xr = (X.copy() - X_mean) / X_std
        yr = (y.copy() - y_mean) / y_std
        return Xr, yr, X_mean, y_mean, X_std, y_std

    else:
        X = (X - X_mean) / X_std
        y = (y - y_mean) / y_std
        return X, y, X_mean, y_mean, X_std, y_std

def l1_cross_distances(X):

    """
    Computes the nonzero componentwise L1 cross-distances between the vectors
    in X.

    Parameters
    ----------

    X: np.ndarray [n_obs, dim]
            - The input variables.

    Returns
    -------

    D: np.ndarray [n_obs * (n_obs - 1) / 2, dim]
            - The L1 cross-distances between the vectors in X.

    ij: np.ndarray [n_obs * (n_obs - 1) / 2, 2]
            - The indices i and j of the vectors in X associated to the cross-
              distances in D.
    """

    n_samples, n_features = X.shape
    n_nonzero_cross_dist = n_samples * (n_samples - 1) // 2
    ij = np.zeros((n_nonzero_cross_dist, 2), dtype=np.int)
    D = np.zeros((n_nonzero_cross_dist, n_features))
    ll_1 = 0

    for k in range(n_samples - 1):
        ll_0 = ll_1
        ll_1 = ll_0 + n_samples - k - 1
        ij[ll_0:ll_1, 0] = k
        ij[ll_0:ll_1, 1] = np.arange(k + 1, n_samples)
        D[ll_0:ll_1] = np.abs(X[k] - X[(k + 1):n_samples])

    return D, ij.astype(np.int)


def abs_exp(theta, d):

    """
    Absolute exponential autocorrelation model.
    (Ornstein-Uhlenbeck stochastic process)::

    Parameters
    ----------
    theta : list[ncomp]
        the autocorrelation parameter(s).

    d: np.ndarray[n_obs * (n_obs - 1) / 2, n_comp]
        |d_i * coeff_pls_i| if PLS is used, |d_i| otherwise

    Returns
    -------
    r : np.ndarray[n_obs * (n_obs - 1) / 2,1]
        An array containing the values of the autocorrelation model.
    """

    r = np.zeros((d.shape[0],1))
    n_components = d.shape[1]

    # Construct/split the correlation matrix
    i,nb_limit  = 0,int(1e4)
    while True:
        if i * nb_limit > d.shape[0]:
            return r
        else:
            r[i*nb_limit:(i+1)*nb_limit,0] = np.exp(-np.sum(theta.reshape(1,
                    n_components) * d[i*nb_limit:(i+1)*nb_limit,:], axis=1))
            i+=1


def squar_exp(theta, d):

    """
    Squared exponential correlation model.

    Parameters
    ----------
    theta : list[ncomp]
        the autocorrelation parameter(s).

    d: np.ndarray[n_obs * (n_obs - 1) / 2, n_comp]
        |d_i * coeff_pls_i| if PLS is used, |d_i| otherwise

    Returns
    -------
    r: np.ndarray[n_obs * (n_obs - 1) / 2,1]
        An array containing the values of the autocorrelation model.
    """

    r = np.zeros((d.shape[0],1))
    n_components = d.shape[1]

    # Construct/split the correlation matrix
    i,nb_limit  = 0,int(1e4)

    while True:
        if i * nb_limit > d.shape[0]:
            return r
        else:
            r[i*nb_limit:(i+1)*nb_limit,0] = np.exp(-np.sum(theta.reshape(1,
                    n_components) * d[i*nb_limit:(i+1)*nb_limit,:], axis=1))
            i+=1


def ge_compute_pls(X,y,n_comp,pts,delta_x,xlimits,extra_points):

    """
    Gradient-enhanced PLS-coefficients.

    Parameters
    ----------

    X: np.ndarray [n_obs,dim]
            - - The input variables.

    y: np.ndarray [n_obs,1]
            - The output variable

    n_comp: int
            - Number of principal components used.

    pts: dict()
            - The gradient values.

    delta_x: real
            - The step used in the FOTA.

    xlimits: np.ndarray[dim, 2]
            - The upper and lower var bounds.

    extra_points: int
            - The number of extra points per each training point.

    Returns
    -------

    Coeff_pls: np.ndarray[dim, n_comp]
            - The PLS-coefficients.

    XX: np.ndarray[extra_points*nt, dim]
            - Extra points added (when extra_points > 0)

    yy: np.ndarray[extra_points*nt, 1]
            - Extra points added (when extra_points > 0)

    """
    nt,dim = X.shape
    XX = np.empty(shape = (0,dim))
    yy = np.empty(shape = (0,1))
    _pls = pls(n_comp)

    coeff_pls = np.zeros((nt,dim,n_comp))
    for i in range(nt):
        if dim >= 3:
            sign = np.roll(bbdesign(int(dim),center=1),1,axis=0)
            _X = np.zeros((sign.shape[0],dim))
            _y = np.zeros((sign.shape[0],1))
            sign = sign * delta_x*(xlimits[:,1]-xlimits[:,0])
            _X = X[i,:]+ sign
            for j in range(1,dim+1):
                sign[:,j-1] = sign[:,j-1]*pts[None][j][1][i,0]
            _y = y[i,:]+ np.sum(sign,axis=1).reshape((sign.shape[0],1))
        else:
            _X = np.zeros((9,dim))
            _y = np.zeros((9,1))
            # center
            _X[:,:] = X[i,:].copy()
            _y[0,0] = y[i,0].copy()
            # right
            _X[1,0] +=delta_x*(xlimits[0,1]-xlimits[0,0])
            _y[1,0] = _y[0,0].copy()+ pts[None][1][1][i,0]*delta_x*(
                xlimits[0,1]-xlimits[0,0])
            # up
            _X[2,1] +=delta_x*(xlimits[1,1]-xlimits[1,0])
            _y[2,0] = _y[0,0].copy()+ pts[None][2][1][i,0]*delta_x*(
                xlimits[1,1]-xlimits[1,0])
            # left
            _X[3,0] -=delta_x*(xlimits[0,1]-xlimits[0,0])
            _y[3,0] = _y[0,0].copy()- pts[None][1][1][i,0]*delta_x*(
                xlimits[0,1]-xlimits[0,0])
            # down
            _X[4,1] -=delta_x*(xlimits[1,1]-xlimits[1,0])
            _y[4,0] = _y[0,0].copy()-pts[None][2][1][i,0]*delta_x*(
                xlimits[1,1]-xlimits[1,0])
            # right up
            _X[5,0] +=delta_x*(xlimits[0,1]-xlimits[0,0])
            _X[5,1] +=delta_x*(xlimits[1,1]-xlimits[1,0])
            _y[5,0] = _y[0,0].copy()+ pts[None][1][1][i,0]*delta_x*(
                xlimits[0,1]-xlimits[0,0])+pts[None][2][1][i,0]*delta_x*(
                xlimits[1,1]-xlimits[1,0])
            # left up
            _X[6,0] -=delta_x*(xlimits[0,1]-xlimits[0,0])
            _X[6,1] +=delta_x*(xlimits[1,1]-xlimits[1,0])
            _y[6,0] = _y[0,0].copy()- pts[None][1][1][i,0]*delta_x*(
                xlimits[0,1]-xlimits[0,0])+pts[None][2][1][i,0]*delta_x*(
                xlimits[1,1]-xlimits[1,0])
            # left down
            _X[7,0] -=delta_x*(xlimits[0,1]-xlimits[0,0])
            _X[7,1] -=delta_x*(xlimits[1,1]-xlimits[1,0])
            _y[7,0] = _y[0,0].copy()- pts[None][1][1][i,0]*delta_x*(
                xlimits[0,1]-xlimits[0,0])-pts[None][2][1][i,0]*delta_x*(
                xlimits[1,1]-xlimits[1,0])
            # right down
            _X[8,0] +=delta_x*(xlimits[0,1]-xlimits[0,0])
            _X[8,1] -=delta_x*(xlimits[1,1]-xlimits[1,0])
            _y[8,0] = _y[0,0].copy()+ pts[None][1][1][i,0]*delta_x*(
                xlimits[0,1]-xlimits[0,0])-pts[None][2][1][i,0]*delta_x*(
                xlimits[1,1]-xlimits[1,0])

        _pls.fit(_X.copy(),_y.copy())
        coeff_pls[i,:,:] = _pls.x_rotations_
        #Add additional points
        if extra_points != 0:
            max_coeff = np.argsort(np.abs(coeff_pls[i,:,0]))[-extra_points:]
            for ii in max_coeff:
                XX = np.vstack((XX,X[i,:]))
                XX[-1,ii] += delta_x*(xlimits[ii,1]-xlimits[ii,0])
                yy = np.vstack((yy,y[i,0]))
                yy[-1,0] += pts[None][1+ii][1][i,0]*delta_x*(
                    xlimits[ii,1]-xlimits[ii,0])
    return np.abs(coeff_pls).mean(axis=0), XX, yy

def componentwise_distance(D,corr,dim):

    """
    Computes the nonzero componentwise cross-spatial-correlation-distance
    between the vectors in X.

    Parameters
    ----------

    D: np.ndarray [n_obs * (n_obs - 1) / 2, dim]
            - The cross-distances between the vectors in X depending of the correlation function.

    corr: str
            - Name of the correlation function used.
              squar_exp or abs_exp.

    dim: int
            - Number of dimension.

    Returns
    -------

    D_corr: np.ndarray [n_obs * (n_obs - 1) / 2, dim]
            - The componentwise cross-spatial-correlation-distance between the
              vectors in X.

    """
    # Fit the matrix iteratively: avoid some memory troubles .
    limit=int(1e4)

    D_corr = np.zeros((D.shape[0],dim))
    i,nb_limit  = 0,int(limit)

    while True:
        if i * nb_limit > D_corr.shape[0]:
            return D_corr
        else:
            if corr == 'squar_exp':
                D_corr[i*nb_limit:(i+1)*nb_limit,:] = D[i*nb_limit:(i+1)*
                                                      nb_limit,:]**2
            else:
                # abs_exp
                D_corr[i*nb_limit:(i+1)*nb_limit,:] = np.abs(D[i*nb_limit:
                                                    (i+1)*nb_limit,:])
            i+=1

def componentwise_distance_PLS(D,corr,n_comp,coeff_pls):

    """
    Computes the nonzero componentwise cross-spatial-correlation-distance
    between the vectors in X.

    Parameters
    ----------

    D: np.ndarray [n_obs * (n_obs - 1) / 2, dim]
            - The L1 cross-distances between the vectors in X.

    corr: str
            - Name of the correlation function used.
              squar_exp or abs_exp.

    n_comp: int
            - Number of principal components used.

    coeff_pls: np.ndarray [dim, n_comp]
            - The PLS-coefficients.

    Returns
    -------

    D_corr: np.ndarray [n_obs * (n_obs - 1) / 2, n_comp]
            - The componentwise cross-spatial-correlation-distance between the
              vectors in X.

    """
    # Fit the matrix iteratively: avoid some memory troubles .
    limit=int(1e4)

    D_corr = np.zeros((D.shape[0],n_comp))
    i,nb_limit  = 0,int(limit)

    while True:
        if i * nb_limit > D_corr.shape[0]:
            return D_corr
        else:
            if corr == 'squar_exp':
                D_corr[i*nb_limit:(i+1)*nb_limit,:] = np.dot(D[i*nb_limit:
                                (i+1)*nb_limit,:]** 2,coeff_pls**2)
            else:
                # abs_exp
                D_corr[i*nb_limit:(i+1)*nb_limit,:] = np.dot(np.abs(D[i*
                                nb_limit:(i+1)*nb_limit,:]),np.abs(coeff_pls))
            i+=1
