import numpy as np

from scipy import signal
from pyriemann.estimation import BlockCovariances
from pyriemann.tangentspace import TangentSpace
from pyriemann.preprocessing import Whitening
from sklearn.pipeline import make_pipeline
from sklearn.base import BaseEstimator, TransformerMixin
from mne.decoding import Vectorizer


class DecimateFeature(BaseEstimator, TransformerMixin):
    def __init__(self, fs, target_fs=10, axis=-1):
        self.fs = fs
        self.target_fs = target_fs
        self.axis = axis

    def fit(self, X, y=None):
        return self
    
    def transform(self, X, y=None):
        decimate_rate = np.sqrt(self.fs / self.target_fs).astype(np.int16)
        X = signal.decimate(X, decimate_rate, axis=self.axis, zero_phase=True)
        # to 10Hz
        X = signal.decimate(X, decimate_rate, axis=self.axis, zero_phase=True)
        return X


class ChannelScaler(BaseEstimator, TransformerMixin):
    def __init__(self, norm_axis=(0, 2)):
        self.channel_mean_ = None
        self.channel_std_ = None
        self.norm_axis=norm_axis

    def fit(self, X, y=None):
        '''

        :param X: 3d array with shape (n_epochs, n_channels, n_times)
        :param y:
        :return:
        '''
        self.channel_mean_ = np.mean(X, axis=self.norm_axis, keepdims=True)
        self.channel_std_ = np.std(X, axis=self.norm_axis, keepdims=True)
        return self

    def transform(self, X, y=None):
        X = X.copy()
        X -= self.channel_mean_
        X /= self.channel_std_
        return X


def riemann_feature_embedder(feat_dim, estimator='lwf'):
    return make_pipeline(
        ChannelScaler(),
        BlockCovariances(block_size=feat_dim, estimator=estimator),
        Whitening(metric='riemann', dim_red={'expl_var': 0.99}),
        TangentSpace()
    )


def baseline_feature_embedder(fs, target_fs, axis):
    return make_pipeline(
        DecimateFeature(fs, target_fs, axis),
        ChannelScaler(),
        Vectorizer()
    )