# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/111b_models.MINIROCKET.ipynb (unless otherwise specified).

__all__ = ['MiniRocketFeatures', 'MiniRocket', 'MiniRocketClassifier', 'load_minirocket', 'MiniRocketRegressor',
           'load_minirocket', 'MiniRocketVotingClassifier', 'get_minirocket_preds', 'MiniRocketVotingRegressor']

# Cell
from ..imports import *
from ..utils import *
from ..data.external import *
from .layers import *

# Cell
from sktime.transformations.panel.rocket._minirocket import _fit as minirocket_fit
from sktime.transformations.panel.rocket._minirocket import _transform as minirocket_transform
from sktime.transformations.panel.rocket._minirocket_multivariate import _fit_multi as minirocket_fit_multi
from sktime.transformations.panel.rocket._minirocket_multivariate import _transform_multi as minirocket_transform_multi
from sktime.transformations.panel.rocket import MiniRocketMultivariate
from sklearn.linear_model import RidgeCV, RidgeClassifierCV

# Cell
# This is a wrapper used to extract features in the format required by the tsai library.

class MiniRocketFeatures:
    def __init__(self, standardize=False, by_sample=False, by_var=False, by_step=False):
        self.standardize = standardize or by_sample or by_var or by_step
        drop_axes = []
        if by_sample: drop_axes.append(0)
        if by_var: drop_axes.append(1)
        if by_step: drop_axes.append(2)
        self.axes = tuple([ax for ax in (0, 1, 2) if ax not in drop_axes])
        self.by_sample = by_sample

    def fit(self, o, num_features=10_000, max_dilations_per_kernel=32):
        if o.dtype != 'float32': o = o.astype('float32')
        if o.ndim == 2:
            o = o[:, np.newaxis]
        if self.standardize:
            mean = o.mean(axis=self.axes, keepdims=True)
            std = o.std(axis=self.axes, keepdims=True) + 1e-8
            if not self.by_sample:
                self.mean, self.std = mean, std
        if o.shape[1] == 1:
            parameters = minirocket_fit(o[0, 0][np.newaxis], num_features=num_features, max_dilations_per_kernel=max_dilations_per_kernel)
        else:
            parameters = minirocket_fit_multi(o[0][np.newaxis], num_features=num_features, max_dilations_per_kernel=max_dilations_per_kernel)
        self.parameters = parameters

    def transform(self, o, fname='X_tfm', path='./data/MiniRocketFeatures', on_disk=True, mode='c', chunksize=10_000):
        if o.dtype != 'float32': o = o.astype('float32')
        if o.ndim == 2:
            o = o[:, np.newaxis]
        if self.standardize:
            if self.by_sample:
                mean = o.mean(axis=self.axes, keepdims=True)
                std = o.std(axis=self.axes, keepdims=True) + 1e-8
            else:
                mean, std = self.mean, self.std
            o = (o - mean) / std
        if chunksize is None:
            if o.shape[1] == 1:
                o_tfm = minirocket_transform(o[:, 0], self.parameters)[..., np.newaxis]
            else:
                o_tfm = minirocket_transform_multi(o, self.parameters)[..., np.newaxis]
            return o_tfm
        else:
            start = 0
            pb = progress_bar(range(math.ceil(len(o) / chunksize)), leave=False, comment="creating MiniRocket features")
            for i in pb:
                end = start + chunksize
                if o.shape[1] == 1:
                    _o_tfm = minirocket_transform(o[start:end, 0], self.parameters)[..., np.newaxis]
                else:
                    _o_tfm = minirocket_transform_multi(o[start:end], self.parameters)[..., np.newaxis]
                if i == 0:
                    shape = (o.shape[0], _o_tfm.shape[1], _o_tfm.shape[2])
                    o_tfm = create_empty_array(shape, fname=fname, path=path, on_disk=on_disk, mode=mode)
                o_tfm[start:end] = _o_tfm
                start = end
                del _o_tfm
                gc.collect()
            return o_tfm

    def fit_transform(self, o, num_features=10_000, max_dilations_per_kernel=32,
                        fname='X_tfm', path='./data/MiniRocketFeatures', on_disk=True, mode='c', chunksize=10_000):
        self.fit(o, num_features=num_features, max_dilations_per_kernel=max_dilations_per_kernel)
        return self.transform(o, fname=fname, path=path, on_disk=on_disk, mode=mode, chunksize=chunksize)

    def save(self, fname, path='./models/MiniRocketFeatures'):
        path = Path(path)
        if not fname.endswith('pkl'): fname = f'{fname}.pkl'
        filename = path/fname
        filename.parent.mkdir(parents=True, exist_ok=True)
        with open(filename, 'wb') as output:
            pickle.dump(self, output, pickle.HIGHEST_PROTOCOL)

    def load(self, fname, path='./models/MiniRocketFeatures'):
        path = Path(path)
        if not fname.endswith('pkl'): fname = f'{fname}.pkl'
        filename = path/fname
        filename.parent.mkdir(parents=True, exist_ok=True)
        with open(filename, 'rb') as input:
            output = pickle.load(input)
        return output

# Cell
# This is an unofficial MINIROCKET implementation in Pytorch developed by Ignacio Oguiza - timeseriesAI@gmail.com based on:
# Dempster, A., Schmidt, D. F., & Webb, G. I. (2020). MINIROCKET: A Very Fast (Almost) Deterministic Transform for Time Series Classification.
# arXiv preprint arXiv:2012.08791.
# Official repo: https://github.com/angus924/minirocket

class MiniRocket(nn.Sequential):
    def __init__(self, c_in, c_out, seq_len=1, fc_dropout=0., **kwargs):
        """
        MINIROCKET implementation where features are previously calculated.
        Args:
            c_in: number of features per sample. For 10_000 kernels iw will be 9996.
            c_out: number of classes.
            seq_len: For MINIROCKET this is always 1 as features are previously calculated. Included for compatibility.
            fc_dropout: indicates whether dropout should be added to the last fully connected layer

        Input shape: [batch_size x c_in x 1]
        """
        backbone = nn.Sequential()
        self.head_nf = c_in
        layers = [Squeeze()]
        if fc_dropout:
            layers += [nn.Dropout(fc_dropout)]
        linear = nn.Linear(c_in, c_out)
        nn.init.constant_(linear.weight.data, 0)
        nn.init.constant_(linear.bias.data, 0)
        layers += [linear]
        head = nn.Sequential(*layers)
        super().__init__(OrderedDict([('backbone', backbone), ('head', head)]))

# Cell
class MiniRocketClassifier(sklearn.pipeline.Pipeline):
    def __init__(self, num_features=10_000, max_dilations_per_kernel=32, random_state=None,
                 alphas=np.logspace(-3, 3, 7), normalize_features=True, memory=None, verbose=False, scoring=None, class_weight=None, **kwargs):
        """
        MiniRocketClassifier is recommended for up to 10k time series.
        For a larger dataset, you can use MINIROCKET (in Pytorch).
        scoring = None --> defaults to accuracy.
        """
        self.steps = [('minirocketmultivariate', MiniRocketMultivariate(num_features=num_features,
                                                                        max_dilations_per_kernel=max_dilations_per_kernel,
                                                                        random_state=random_state)),
                      ('ridgeclassifiercv', RidgeClassifierCV(alphas=alphas,
                                                              normalize=normalize_features,
                                                              scoring=scoring,
                                                              class_weight=class_weight,
                                                              **kwargs))]
        store_attr()
        self._validate_steps()

    def __repr__(self):
        return f'Pipeline(steps={self.steps.copy()})'

    def save(self, fname=None, path='./models'):
        fname = ifnone(fname, 'MiniRocketClassifier')
        path = Path(path)
        filename = path/fname
        with open(f'{filename}.pkl', 'wb') as output:
            pickle.dump(self, output, pickle.HIGHEST_PROTOCOL)


def load_minirocket(fname, path='./models'):
    path = Path(path)
    filename = path/fname
    with open(f'{filename}.pkl', 'rb') as input:
        output = pickle.load(input)
    return output

# Cell
class MiniRocketRegressor(sklearn.pipeline.Pipeline):
    def __init__(self, num_features=10000, max_dilations_per_kernel=32, random_state=None,
                 alphas=np.logspace(-3, 3, 7), *, normalize_features=True, memory=None, verbose=False, scoring=None, **kwargs):
        """
        MiniRocketRegressor is recommended for up to 10k time series.
        For a larger dataset, you can use MINIROCKET (in Pytorch).
        scoring = None --> defaults to r2.
        """
        self.steps = [('minirocketmultivariate', MiniRocketMultivariate(num_features=num_features,
                                                                        max_dilations_per_kernel=max_dilations_per_kernel,
                                                                        random_state=random_state)),
                      ('ridgecv', RidgeCV(alphas=alphas, normalize=normalize_features, scoring=scoring, **kwargs))]
        store_attr()
        self._validate_steps()

    def __repr__(self):
        return f'Pipeline(steps={self.steps.copy()})'

    def save(self, fname=None, path='./models'):
        fname = ifnone(fname, 'MiniRocketRegressor')
        path = Path(path)
        filename = path/fname
        with open(f'{filename}.pkl', 'wb') as output:
            pickle.dump(self, output, pickle.HIGHEST_PROTOCOL)


def load_minirocket(fname, path='./models'):
    path = Path(path)
    filename = path/fname
    with open(f'{filename}.pkl', 'rb') as input:
        output = pickle.load(input)
    return output

# Cell
from sklearn.ensemble import VotingClassifier
class MiniRocketVotingClassifier(VotingClassifier):
    def __init__(self, n_estimators=5, weights=None, n_jobs=-1, num_features=10_000, max_dilations_per_kernel=32, random_state=None,
                 alphas=np.logspace(-3, 3, 7), normalize_features=True, memory=None, verbose=False, scoring=None, class_weight=None, **kwargs):
        store_attr()
        estimators = [(f'clf{i}', MiniRocketClassifier(num_features=num_features, max_dilations_per_kernel=max_dilations_per_kernel,
                                                       random_state=random_state, alphas=alphas, normalize_features=normalize_features, memory=memory,
                                                       verbose=verbose, scoring=scoring, class_weight=class_weight, **kwargs))
                    for i in range(n_estimators)]
        super().__init__(estimators, voting='hard', weights=weights, n_jobs=n_jobs, verbose=verbose)

    def __repr__(self):
        return f'MiniRocketVotingClassifier(n_estimators={self.n_estimators}, \nsteps={self.estimators[0][1].steps})'

    def save(self, fname=None, path='./models'):
        fname = ifnone(fname='MiniRocketVotingClassifier')
        path = Path(path)
        filename = path/fname
        filename.parent.mkdir(parents=True, exist_ok=True)
        with open(f'{filename}.pkl', 'wb') as output:
            pickle.dump(self, output, pickle.HIGHEST_PROTOCOL)


def get_minirocket_preds(X, fname, path='./models', model=None):
    if X.ndim == 1: X = X[np.newaxis][np.newaxis]
    elif X.ndim == 2: X = X[np.newaxis]
    if model is None:
        model = load_minirocket(fname=fname, path=path)
    return model.predict(X)

# Cell
from sklearn.ensemble import VotingRegressor

class MiniRocketVotingRegressor(VotingRegressor):
    def __init__(self, n_estimators=5, weights=None, n_jobs=-1, num_features=10_000, max_dilations_per_kernel=32, random_state=None,
                 alphas=np.logspace(-3, 3, 7), normalize_features=True, memory=None, verbose=False, scoring=None, **kwargs):
        store_attr()
        estimators = [(f'clf{i}', MiniRocketRegressor(num_features=num_features, max_dilations_per_kernel=max_dilations_per_kernel,
                                                      random_state=random_state, alphas=alphas, normalize_features=normalize_features, memory=memory,
                                                      verbose=verbose, scoring=scoring, **kwargs))
                      for i in range(n_estimators)]
        super().__init__(estimators, weights=weights, n_jobs=n_jobs, verbose=verbose)

    def __repr__(self):
        return f'MiniRocketVotingRegressor(n_estimators={self.n_estimators}, \nsteps={self.estimators[0][1].steps})'

    def save(self, fname=None, path='./models'):
        fname = ifnone(fname='MiniRocketVotingRegressor')
        path = Path(path)
        filename = path/fname
        filename.parent.mkdir(parents=True, exist_ok=True)
        with open(f'{filename}.pkl', 'wb') as output:
            pickle.dump(self, output, pickle.HIGHEST_PROTOCOL)