from joblib import dump, load
from keras.models import load_model
from tslearn.shapelets.shapelets import LocalSquaredDistanceLayer, GlobalMinPooling1D
from tslearn.shapelets import LearningShapelets
from tslearn.datasets import UCR_UEA_datasets

# Load a dataset
X_train, y_train, X_test, y_test = UCR_UEA_datasets().load_dataset('Coffee')

# Define and fit an instance of LearningShapelets
clf = LearningShapelets(random_state=42)
clf.fit(X_train, y_train)

# Save the model by accessing the 'model_' attribute
clf.model_.save('shapelet_model.h5')

# Save all the attributes
clf.model_ = None
clf.transformer_model_ = None
clf.locator_model_ = None
dump(clf.__dict__, 'attributes.gz')

# Delete the instance
del clf

# Create a new instance
clf = LearningShapelets()

# Define the 'model_' attribute of this instance
model = load_model(
    'shapelet_model.h5',
    custom_objects={'LocalSquaredDistanceLayer': LocalSquaredDistanceLayer,
                    'GlobalMinPooling1D': GlobalMinPooling1D}
)
clf.model_ = model

# Define the other necessary attributes
attributes = load('attributes.gz')
for i in attributes:
    clf.i = attributes[i]

# Predict on the test set
clf.score(X_test, y_test)