"""
Classic TRF.
"""

import numpy as np
from sklearn.model_selection import KFold
import matplotlib.pyplot as plt
from mne.decoding import BaseEstimator
from ..utils import lag_matrix, lag_span, lag_sparse, mem_check
from scipy import linalg
import mne
from ._methods import _ridge_fit_SVD, _get_covmat, _corr_multifeat, _rmse_multifeat, _r2_multifeat, _rankcorr_multifeat, _ezr2_multifeat, _adjr2_multifeat, _fourier_fit
from matplotlib import colormaps as cmaps

# Memory cap (i.e. max usage).
# By default set to 90% to prevent bricking machines in corner cases.
MEM_CAP = 0.9


class TRFEstimator(BaseEstimator):
    """
    This implements a TRF to relate a continuous stimuli and electrophysiological data. The model functions under the assumption that the signal behaves like a linear time invariant system i.e. its response is a linear function of the input, that does not change over time. It lags the input in time and combine each lag to reconstruct the electrophysiological signal using a Ridge regression. Since this operation is equivalent to a convolution in time, the TRF can alternatively be fitted in the Fourier domain as a multiplication.
        
    Parameters
    ----------
    times : tuple
        mismatch a -> b, where a - dependent, b - predicted. Negative timelags indicate a lagging behind b. Positive timelags indicate b lagging behind a
    tmin : float
        Minimum time lag (in seconds). Can be negative to check for lags in the past (~null model).
    tmax : float
        Maximum time lag (in seconds). Can be large to check for lags in the future (~null model).
    srate : float
        Sampling rate of the data.
    alpha : list
        Regularization parameter(s) for the model. If a list is provided, the model will be fitted for each alpha.
        The fit is looping over alpha AFTER most matrix operations to optimize compute time. Careful, we weight matrices according to the
        eigenvalues and this value of alpha is not directly equivalent to the one given in (XtX + alpha*I)XtY. Ideally, alpha should be computed for many values.
    fit_intercept : bool
        Whether to fit an intercept term in the model.
    mtype : str
        'forward' or 'backward'. Required for formatting coefficients in get_coef (convention: forward - stimulus -> eeg, backward - eeg - stimulus)
    alpha_feat : bool
        Whether to compute alpha for each feature separately and fit them. If True, alpha will be modified to be a list of all possible combinations. This increases computation time exponentially, only use if dealing with few very different regressors, or to check differences are minimal.
    fit_domain : str
        'time' or 'frequency', whether to treat TRF fitting as a regression (fit in time via Ridge regression) or as a convolution (fit in frequency as a division). The time-fit is the usual method but the frequency-fit offers comparable performance while being generally less time and space consumming. 
    """

    def __init__(self, times=(0.,), tmin=None, tmax=None, srate=1., alpha=[0.], 
                 fit_intercept=False, mtype='forward', alpha_feat = False, fit_domain = 'time'):
        """
        Initialize the class instance.
        """

        self.tmin = tmin
        self.tmax = tmax
        self.times = times
        self.srate = srate
        self.alpha = alpha
        self.alpha_feat = alpha_feat
        # Forward or backward. Required for formatting coefficients in get_coef (convention: forward - stimulus -> eeg, backward - eeg - stimulus)
        self.mtype = mtype
        self.fit_intercept = fit_intercept
        self.fitted = False
        self.lags = None
        self.fit_domain = fit_domain

        # All following attributes are only defined once fitted (hence the "_" suffix)
        self.intercept_ = None
        self.coef_ = None
        self.n_feats_ = None
        self.n_chans_ = None
        self.feat_names_ = None
        self.valid_samples_ = None
        # Autocorrelation matrix of feature X (thus XtX) -> used for computing model using fit_from_cov
        self.XtX_ = None
        # Covariance matrix of features X and Y (thus XtX) -> used for computing model using fit_from_cov
        self.XtY_ = None
        # Scores when computed
        self.scores = None

    def fill_lags(self):
        """
        Fill the lags attributes, with number of samples and times in seconds.
        """
        if (self.tmin != None) and (self.tmax != None):
            # LOGGER.info("Will use lags spanning form tmin to tmax.\nTo use individual lags, use the `times` argument...")
            self.lags = lag_span(self.tmin, self.tmax, srate=self.srate)[
                ::-1]  # pylint: disable=invalid-unary-operand-type
            # self.lags = lag_span(-tmax, -tmin, srate=srate) #pylint: disable=invalid-unary-operand-type
            self.times = self.lags[::-1] / self.srate
        else:
            self.times = np.asarray(self.times)
            self.lags = lag_sparse(self.times, self.srate)[::-1]

    def get_XY(self, X, y, lagged=False, drop=True, feat_names=()):
        """
        Preprocess X and y before fitting (finding mapping between X -> y).
        
        Parameters
        ----------
        X : ndarray 
            input, of shape (T, nfeat)
        y : ndarray 
            output, of shape (T, nchan)
        lagged : bool
            Whether the X matrix has been previously 'lagged'.
        drop : bool
            Whether to drop non valid samples (if False, non valid sample are filled with 0.)
        feat_names : list
            Names of features being fitted. Must be of length ``nfeats``.
            
        Returns
        -------
        X : ndarray 
            Preprocessed input, of shape (T, nlags * nfeats)
        y : ndarray
            Preprocessed output, of shape (T, nlags * nfeats)
        """
        self.fill_lags()

        X = np.asarray(X)
        y = np.asarray(y)

        # Estimate the necessary size to compute stuff
        y_memory = sum([yy.nbytes for yy in y]) if np.ndim(
            y) == 3 else y.nbytes
        estimated_mem_usage = X.nbytes * \
            (len(self.lags) if not lagged else 1) + y_memory
        if estimated_mem_usage/1024.**3 > MEM_CAP*mem_check():
            raise MemoryError("Not enough RAM available! (needed %.1fGB, but only %.1fGB available)" % (
                estimated_mem_usage/1024.**3, mem_check()))

        # Fill n_feat and n_chan attributes
        # if X has been lagged, divide by number of lags
        self.n_feats_ = X.shape[1] if not lagged else X.shape[1] // len(
            self.lags)
        self.n_chans_ = y.shape[1] if y.ndim == 2 else y.shape[2]

        # Assess if feat names corresponds to feat number
        if feat_names:
            err_msg = "Length of feature names does not match number of columns from feature matrix"
            if lagged:
                assert len(feat_names) == X.shape[1] // len(self.lags), err_msg
            else:
                assert len(feat_names) == X.shape[1], err_msg
            self.feat_names_ = feat_names

        # this include non-valid samples for now
        n_samples_all = y.shape[0] if y.ndim == 2 else y.shape[1]

        # drop samples that can't be reconstructed because on the edge, all is true otherwise
        if drop:
            self.valid_samples_ = np.logical_not(np.logical_or(np.arange(n_samples_all) < abs(max(self.lags)),
                                                               np.arange(n_samples_all)[::-1] < abs(min(self.lags))))
        else:
            self.valid_samples_ = np.ones((n_samples_all,), dtype=bool)

        # Creating lag-matrix droping NaN values if necessary
        y = y[self.valid_samples_, :] if y.ndim == 2 else y[:,
                                                            self.valid_samples_, :]
        if not lagged:
            X = lag_matrix(X, lag_samples=self.lags,
                           drop_missing=drop, filling=np.nan if drop else 0.)

        return X, y

    def fit(self, X, y, lagged=False, drop=True, feat_names=()):
        """
        Fit the TRF model in either the time of frequency domain. The convention is to map X -> y. In order to properly retrieve the coefficients, respect the convention using the 'mtype' argument. This fills the coefficients attributes with shape (alphas, nlags, nfeats)
        
        Parameters
        ----------
        X : ndarray 
            input of shape (T, nfeats)
        y : ndarray
            output of shape (T, nfeats)
        lagged : bool
            Whether the X matrix has been previously 'lagged'.
        drop : bool
            Whether to drop non valid samples (if False, non valid sample are filled with 0.)
        feat_names : list
            Names of features being fitted. Must be of length ``nfeats``.
            
        Returns
        -------
        coef_ : ndarray 
            coefficients of shape (alphas, nlags, nfeats)
        """

        if self.fit_domain == 'time':
            # Preprocess and lag inputs
            X, y = self.get_XY(X, y, lagged, drop, feat_names)
    
            # Regress with Ridge to obtain coef for the input alpha
            self.coef_ = _ridge_fit_SVD(X, y, self.alpha, alpha_feat = self.alpha_feat, n_feat=self.n_feats_)
            
        elif self.fit_domain == 'frequency':
            # Compute TRF in the Frequency Domain
            self.coef_ = _fourier_fit(X, y, self.alpha, self.lags)
        self.fitted = True

        return self.coef_.copy()

    def get_coef(self):
        """
        Format and return coefficients. Note mtype attribute needs to be declared in the __init__.

        Returns
        -------
        coef_ : ndarray (nlags x nfeats x nchans x regularization params)
        """
        if np.ndim(self.alpha) == 0:
            betas = np.reshape(self.coef_, (len(self.lags),
                                            self.n_feats_, self.n_chans_))
        elif self.alpha_feat:
            betas = np.reshape(self.coef_, (len(self.lags),
                                            self.n_feats_, self.n_chans_, np.power(len(self.alpha), self.n_feats_)))
        else:
            betas = np.reshape(self.coef_, (len(self.lags),
                                            self.n_feats_, self.n_chans_, len(self.alpha)))

        if self.mtype == 'forward':
            betas = betas[::-1, :]

        return betas

    def add_cov(self, X, y, lagged=False, drop=True, n_parts=1):
        """
        Compute and add (with normalization factor) covariance matrices XtX, XtY.
        For v. large population models when it's not possible to load all the data to memory at once.
        
        Parameters
        ----------
        X : ndarray 
            input of shape (T, nfeats)
        y : ndarray 
            output of shape (T, nchans)
        lagged : bool
            Whether the X matrix has been previously 'lagged' (intercept still to be added).
        drop : bool
            Whether to drop non valid samples (if False, non valid sample are filled with 0.)
        n_parts : number of parts from which the covariance matrix are computed (required for normalization)
            Default: 1
        """
        if isinstance(X, (list, tuple)) and n_parts > 1:

            assert len(X) == len(y)

            for part in range(len(X)):
                assert len(X[part]) == len(y[part])
                X_part, y_part = self.get_XY(X[part], y[part], lagged, drop)
                XtX = _get_covmat(X_part, X_part)
                XtY = _get_covmat(X_part, y_part)
                norm_pool_factor = np.sqrt(
                    (n_parts*X_part.shape[0] - 1)/(n_parts*(X_part.shape[0] - 1)))

                if self.XtX_ is None:
                    self.XtX_ = XtX*norm_pool_factor
                else:
                    self.XtX_ += XtX*norm_pool_factor

                if self.XtY_ is None:
                    self.XtY_ = XtY*norm_pool_factor
                else:
                    self.XtY_ += XtY*norm_pool_factor

        else:

            X_part, y_part = self.get_XY(X, y, lagged, drop)

            self.XtX_ = _get_covmat(X_part, X_part)
            self.XtY_ = _get_covmat(X_part, y_part)

        return self

    def fit_direct_cov(self, XXcov=None, XYcov=None, clear_after=True):
        """
        Directly fit the TRF by manually giving the preprocessed covariance matrices. This automatically uses the time domain fit as it uses the lagged covariance matrices. Mostly useful to work with specific data shape to test scenarios.
        
        Parameters
        ----------
        XXcov : ndarray
            Covariance matrix of the input, should be lagged beforehand.
        XYcov : ndarray
            Covariance matrix and X and Y
        clear_after : bool
            whether to keep in memory the variables, default to True
        
        """
        self.XtX_ = XXcov
        self.XtY_ = XYcov

        self.fit_intercept = False
        self.intercept_ = None
        self.coef_ = _ridge_fit_SVD(
            self.XtX_, self.XtY_, self.alpha, from_cov=True, alpha_feat = self.alpha_feat, n_feat = self.n_feats_)
        self.fitted = True

        if clear_after:
            self.clear_cov()

        return self

    def fit_from_cov(self, X=None, y=None, lagged=False, drop=True, overwrite=True, part_length=150., clear_after=True):
        """
        Fit model from covariance matrices, handy for very large data, but slower.
        
        Parameters
        ----------
        X : ndarray 
            if None, model will be fitted from accumulated XtX & XtY, for size otherwise (T x nfeats).
        y : ndarray (T x nchans), if None, model will be fitted from accumulated XtX & XtY
            if None, model will be fitted from accumulated XtX & XtY, for size otherwise (T x nfeats).
        lagged : bool
            Whether the X matrix has been previously 'lagged'.
        drop : bool
            Whether to drop non valid samples (if False, non valid sample are filled with 0.)
        overwrite : bool
            Whether to reset the accumulated covariance matrices (when X and Y are not None)
        part_length : integer | float
            Size of the parts in which the data will be chopped for fitting the model (when X and Y are provided), Default to 150 (seconds) ~ 2.5 minutes. Estimate what will fit in RAM.
            
        Returns
        -------
        coef_ : ndarray (alphas x nlags x nfeats)
        """

        # If X and y are not none, chop them into pieces, compute cov matrices and fit the model (memory efficient)
        if (X is not None) and (y is not None):
            if overwrite:
                self.clear_cov()  # If overwrite, wipe clean

            part_length = int(part_length*self.srate)  # seconds to samples

            assert X.shape[0] == y.shape[0]

            # ToDo -> add padding to the segment to spare some data...
            segments = [(part_length * i) + np.arange(part_length)
                        for i in range(X.shape[0] // part_length)]  # Indices making up segments

            if X.shape[0] % part_length > part_length/2:
                # Create new segment from leftover data (if more than 1/2 part length)
                segments = segments + [np.arange(segments[-1][-1], X.shape[0])]
            else:
                # Add letover data to the last segment
                segments[-1] = np.concatenate((segments[-1],
                                               np.arange(segments[-1][-1], X.shape[0])))

            X = [X[segments[i], :] for i in range(len(segments))]
            y = [y[segments[i], :] for i in range(len(segments))]

            self.add_cov(X, y, lagged=False, drop=True, n_parts=len(segments))

        self.fit_intercept = False
        self.intercept_ = None
        self.coef_ = _ridge_fit_SVD(
            self.XtX_, self.XtY_, self.alpha, from_cov=True,
            alpha_feat = self.alpha_feat, n_feat=self.n_feats_)
        self.fitted = True

        if clear_after:
            self.clear_cov()

        return self

    def clear_cov(self):
        """
        Wipe clean / reset covariance matrices.
        """
        # print("Clearing saved covariance matrices...")
        self.XtX_ = None
        self.XtY_ = None
        return self

    def predict(self, X):
        """
        Compute output based on fitted coefficients and feature matrix X.
        
        Parameters
        ----------
        X : ndarray
            Matrix of features (can be already lagged or not).
            
        Returns
        -------
        ndarray
            Reconstruction of target with current beta estimates
        """
        
        assert self.fitted, "Fit model first!"

        if self.fit_intercept:
            betas = np.concatenate((self.intercept_, self.coef_), axis=0)
        else:
            #betas = self.get_coef()[:]
            betas = self.coef_[:]

        # Check if input has been lagged already, if not, do it:
        if X.shape[1] != int(self.fit_intercept) + len(self.lags) * self.n_feats_:
            # LOGGER.info("Creating lagged feature matrix...")
            X = lag_matrix(X, lag_samples=self.lags, filling=0.)
            if self.fit_intercept:  # Adding intercept feature:
                X = np.hstack([np.ones((len(X), 1)), X])

        # Do it for every alpha
        pred = np.stack([X.dot(betas[..., i])
                         for i in range(betas.shape[-1])], axis=-1)

        return pred  # Shape T x Nchan x Alpha

        
    def score(self, Xtest, ytrue, Xtrain = None, scoring="R2"):
        """Compute a score of the model given true target and estimated target from Xtest.
        
        Parameters
        ----------
        Xtest : ndarray
            Array used to get "yhat" estimate from model
        ytrue : ndarray
            True target
        scoring : str
            Scoring function to be used ("corr", "rankcorr", "rmse", "R2", "ezekiel", 'adj_R2')
            
        Returns
        -------
        scores: ndarray
            Scores computed on whole segment.
        """
        yhat = self.predict(Xtest)
        window_length = self.times.shape[0]
        if self.alpha_feat:
            reg_len = np.power(len(self.alpha), self.n_feats_)
        else:
            reg_len = len(self.alpha)
        alpha = self.alpha
        lags = self.lags
        if scoring == 'corr':
            scores = np.stack([_corr_multifeat(yhat[..., a], ytrue, nchans=self.n_chans_) for a in range(reg_len)], axis=-1)
            self.scores = scores
            return scores
        elif scoring == 'rmse':
            scores = np.stack([_rmse_multifeat(yhat[..., a], ytrue) for a in range(reg_len)], axis=-1)
            self.scores = scores
            return scores
        elif scoring == 'R2':
            scores = np.stack([_r2_multifeat(yhat[..., a], ytrue) for a in range(reg_len)], axis=-1)
            self.scores = scores
            return scores
        elif scoring == 'rankcorr':
            scores = np.stack([_rankcorr_multifeat(yhat[..., a], ytrue, nchans=self.n_chans_) for a in range(reg_len)], axis=-1)
            self.scores = scores
            return scores
        elif scoring == 'ezekiel':
            scores = np.stack([_ezr2_multifeat(yhat[..., a], ytrue, Xtest, window_length) for a in range(reg_len)], axis=-1)
            self.scores = scores
            return scores
        elif scoring == 'adj_R2':
            scores = np.stack([_adjr2_multifeat(yhat[..., a], ytrue, Xtrain, Xtest, alpha[a], lags) for a in range(reg_len)], axis=-1)
            self.scores = scores
            return scores
        else:
            raise NotImplementedError(
                "Only correlation & RMSE scores are valid for now...")

    def xval_eval(self, X, y, n_splits=5, lagged=False, drop=True, train_full=True, scoring="R2", segment_length=None, fit_mode='direct', verbose=True):
        """
        Standard cross-validation. Scoring. 
        
        Parameters
        ----------
        X : ndarray 
            input of size (T x nfeats)
        y : ndarray 
            output of size (T x nchans)
        n_splits : integer 
            Number of folds, default to 5.
        lagged : bool
            Whether the X matrix has been previously 'lagged', default to False.
        drop : bool
            Whether to drop non valid samples (if False, non valid sample are filled with 0.).
        train_full : bool 
            Train model using all the available data after the end of x-val
        scoring : string
            Scoring method (see scoring()), default to "R2".
        segment_length: integer, float 
            Length of a testing segments (that testing data will be chopped into). If None, use all the available data. Default to None.
        fit_mode : string {'direct' | 'from_cov_xxx'} (default: 'direct')
            Model training mode. 'direct' - fit using all the avaiable data at once (i.e. fit()), 'from_cov_xxx' - fit using all the avaiable data from covariance matrices. The routine will chop data into pieces, compute piece-wise cov matrices and fit the model. 'xxx' portion of the string indicates the lenght of the segments that the data will be chopped into. 
        verbose : bool
        
        Returns
        -------
        scores : ndarray 
            scores across each fold (n_splits x segments x nchans x alpha)
        """

        #if np.ndim(self.alpha) < 1 or len(self.alpha) <= 1:
        #    raise ValueError(
        #        "Supply several alphas to TRF constructor to use this method.")

        if segment_length:
            segment_length = segment_length*self.srate  # time to samples

        self.fill_lags()

        self.n_feats_ = X.shape[1] if not lagged else X.shape[1] // len(
            self.lags)
        self.n_chans_ = y.shape[1] if y.ndim == 2 else y.shape[2]

        if self.alpha_feat:
            reg_len = np.power(len(self.alpha), self.n_feats_)
        else:
            reg_len = len(self.alpha)

        kf = KFold(n_splits=n_splits)
        if segment_length:
            scores = []
        else:
            scores = np.zeros((n_splits, self.n_chans_, reg_len))

        for kfold, (train, test) in enumerate(kf.split(X)):
            if verbose:
                print("Training/Evaluating fold %d/%d" % (kfold+1, n_splits))

            # Fit using trick with adding covariance matrices -> saves RAM
            if fit_mode.find('from_cov') > -1:
                # Format 'from_cov_xxx' -> xxx - duration of a single part.
                # The main idea is to chunk the data into bite-sized parts that will fit in the RAM
                # Careful with the chunking of data
                if len(fit_mode.split('_')) == 2:
                    part_lenght = 150
                elif len(fit_mode.split('_')) == 3:
                    part_lenght = int(fit_mode.split('_')[-1])
                self.fit_from_cov(X[train, :], y[train, :],
                                  overwrite=True, part_length=part_lenght)
            else:  # Fit directly -> slightly faster, but uses more RAM
                self.fit(X[train, :], y[train, :])

            if segment_length:  # Chop testing data into smaller pieces

                if (len(test) % segment_length) > 0:  # Crop if there are some odd samples
                    test_crop = test[:-int(len(test) % segment_length)]
                else:
                    test_crop = test[:]

                # Reshape to # segments x segment duration
                test_segments = test_crop.reshape(
                    int(len(test_crop) / segment_length), -1)

                ccs = [self.score(X[test_segments[i], :], y[test_segments[i], :], scoring=scoring, Xtrain = X[train, :]) for i in range(
                    test_segments.shape[0])]  # Evaluate each segment

                scores.append(ccs)
            else:  # Evaluate using the entire testing data
                scores[kfold, :] = self.score(X[test, :], y[test, :], scoring=scoring, Xtrain = X[train, :])

        if segment_length:
            scores = np.asarray(scores)

        if train_full:
            if verbose:
                print("Fitting full model...")
            # Fit using trick with adding covariance matrices -> saves RAM
            if fit_mode.find('from_cov') > -1:
                self.fit_from_cov(X, y, overwrite=True,
                                  part_length=part_lenght)
            else:
                self.fit(X, y)
        self.scores = scores

        return scores

    def __getitem__(self, feats):
        """Extract a sub-part of TRF instance as a new TRF instance (useful for plotting only some features...)"""
        # Argument check
        if self.feat_names_ is None:
            if np.ndim(feats) > 0:
                assert isinstance(
                    feats[0], int), "Type not understood, feat_names are ot defined, can only index with int"
                indices = feats

            else:
                assert isinstance(
                    feats, int), "Type not understood, feat_names are ot defined, can only index with int"
                indices = [feats]
                feats = [feats]
        else:
            if np.ndim(feats) > 0:
                assert all([f in self.feat_names_ for f in feats]
                           ), "an element in argument %s in not present in %s" % (feats, self.feat_names_)
                indices = [self.feat_names_.index(f) for f in feats]
            else:
                assert feats in self.feat_names_, "argument %s not present in %s" % (
                    feats, self.feat_names_)
                indices = [self.feat_names_.index(feats)]
                feats = [feats]

        trf = TRFEstimator(tmin=self.tmin, tmax=self.tmax,
                           srate=self.srate, alpha=self.alpha)
        trf.coef_ = self.coef_[:, indices]
        trf.feat_names_ = feats
        trf.n_feats_ = len(feats)
        trf.n_chans_ = self.n_chans_
        trf.fitted = True
        trf.times = self.times
        trf.lags = self.lags
        trf.intercept_ = self.intercept_

        return trf

    def __add__(self, other_trf):
        "Make available the '+' operator. Will simply add coefficients. Be mindful of dividing by the number of elements later if you want the true mean."
        assert (other_trf.n_feats_ == self.n_feats_ and other_trf.n_chans_ ==
                self.n_chans_), "Both TRF objects must have the same number of features and channels"
        trf = TRFEstimator(tmin=self.tmin, tmax=self.tmax,
                           srate=self.srate, alpha=self.alpha)
        trf.coef_ = np.sum([self.coef_, other_trf.coef_], 0)
        trf.intercept_ = np.sum([self.intercept_, other_trf.intercept_], 0)
        trf.feat_names_ = self.feat_names_
        trf.n_feats_ = self.n_feats_
        trf.n_chans_ = self.n_chans_
        trf.fitted = True
        trf.times = self.times
        trf.lags = self.lags

        return trf

    def get_best_alpha(self):
        """
        Return the best regularization parameters over the cross-fold evaluation

        Returns
        -------
        alphas : ndarray
            Best alpha per channel
        """
        best_alpha = np.zeros(self.n_chans_)
        for chan in range(self.n_chans_):
            if len(self.scores.shape) == 3:
                best_alpha[chan] = np.argmax(np.mean(self.scores[:,chan,:],axis=0))
            else:
                best_alpha[chan] = np.argmax(self.scores[:,chan,:],axis=0)
        alphas = best_alpha.astype(int)
        return alphas


    def plot_score(self, figax = None, figsize = (5,5), color_type = 'rainbow', 
                   channels = [], title = 'R2 sumary', minscore = -np.inf):
        """
        Plot the score according to the regularization parameters.

        Parameters
        ----------
        figax : tuple
            contains (fig,ax) matplotlib object, if existing, the dimensions should fit. If None, create a new figure.
        figsize : tuple
            (x,y) size of figure
        color_type : str
            cmap to use
        channels : list
            Select a list of channels indices to plot. If empty, all channels are taken into account.
        title : str
            title of figure
        minscore : float
            Only plot channels that have a score above this value

        Returns
        -------
        fig : Figure
            Matplotlib figure object
        ax : Axes
            Matplotlib axis/axes
        """
        if figax == None:
            fig,ax = plt.subplots(figsize = figsize)
        else:
            fig,ax = figax
        if len(channels) == 0:
            channels = np.arange(self.scores.shape[1])

        #Extract Coef
        color_map = dict()
        for index_channel in range(self.scores.shape[1]):
            color_map[index_channel] = cmaps[color_type](index_channel/self.scores.shape[1])

        for index_channel in range(self.scores.shape[1]):
            score_chan = np.mean(self.scores[:,index_channel,:],axis = 0)
            if np.max(score_chan > minscore):
                ax.plot(self.alpha, score_chan, color = color_map[index_channel], linewidth = 1.5, label = channels[index_channel])
        ax.set_title(title)
        ax.set_xlabel('Alpha')
        ax.set_ylabel('R2')
        ax.set_xticks(self.alpha)
        ax.set_xscale('log')
        ax.plot(self.alpha, np.mean(self.scores[:,:,:],axis = (0,1)), color = 'k', linewidth = 3, linestyle = '--')
        ax.legend()


        return fig, ax

    def plot_kernel(self, figax = None, figsize = False, color_type = 'rainbow', center_line = False,
                    channels = None, features = None, title = 'kernel sumary', minR2 = -np.inf):
        """
        Plot the TRF of the feature requested as a butterfly plot.
        
        Parameters
        ----------
        figax : tuple
            contains (fig,ax) matplotlib object, if existing, the dimensions should fit. If None, create a new figure.
        figsize : tuple
            (x,y) size of figure
        color_type : str
            cmap to use
        center_line : bool
            Whether to plot a line marking the time 0 (i.e. when input and output are in sync)
        channels : list
            Select a list of channels indices to plot. If empty, all channels are taken into account.
        features : list
            Select a list of features indices to plot. If empty, all features are taken into account.
        title : str
            title of figure
        minscore : float
            Only plot channels that have a score above this value

        Returns
        -------
        fig : Figure
            Matplotlib figure object
        ax : Axes
            Matplotlib axis/axes
        """
        
        if not figsize:
            figsize = (15, (self.n_feats_) * 4)
        if figax is None:
            fig,ax = plt.subplots(self.n_feats_,figsize = figsize, sharex = True)
        else:
            fig,ax = figax
        if channels == None:
            channels = np.arange(self.n_chans_)
        if features == None:
            features = np.arange(self.n_feats_)


        color_map = dict()
        for index_channel in range(self.n_chans_):
            color_map[index_channel] = cmaps[color_type](index_channel/self.n_chans_)

        best_alpha = self.get_best_alpha()
        for feat_index in range(self.n_feats_):
            feat = features[feat_index]
            if self.n_feats_ > 1:
                axfeat = ax[feat_index]
            else:
                axfeat = ax
            for chan_index in range(self.n_chans_):
                alpha_index = best_alpha[chan_index]
                chan = channels[chan_index]
                score_chan = np.mean(self.scores[:,chan_index,:],axis = 0)
                if np.max(score_chan > minR2):

                    axfeat.plot(self.times, self.get_coef()[:,feat_index,chan_index, alpha_index], color = color_map[chan_index], linewidth = 1.5, label = chan)
                    axfeat.set_xlabel('Time (s)')
                    axfeat.set_ylabel(feat)
            if center_line:
                axfeat.plot([0,0],[np.min(self.get_coef()[:,feat_index,:, alpha_index]),np.max(self.get_coef()[:,feat_index,:, alpha_index])], color = 'k', linewidth = 1.5, linestyle = '--')
        handles, labels = axfeat.get_legend_handles_labels()
        fig.legend(handles, labels, bbox_to_anchor=(1.15, 0.8),loc='right')
        if self.n_feats_ > 1:
            ax[0].set_title(title)
        else:
            ax.set_title(title)
        return fig,ax

    def __repr__(self):
        obj = """TRFEstimator(
            alpha=%s,
            srate=%s,
            tmin=%s,
            tmax=%s,
            n_feats=%s,
            n_chans=%s,
            n_lags=%s,
            features : %s,
            fit_domain : %s,
            direction: %s
        )
        """ % (self.alpha, self.srate, self.tmin, self.tmax,
               self.n_feats_, self.n_chans_, 
               len(self.lags) if self.lags is not None else None, str(self.feat_names_),
              self.fit_domain, self.mtype)
        return obj

