# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Autoencoder

# This is the main Autoencoder notebook. In this notebook you will find all the code we use to investigate AE model. In particular, the section [Preprocessing](#preprocessing) contains all the code used for the preprocessing pipeline, including data strucure managements, transformations of data, creations of `X` and `y` vectors and so on. Section [Model (autoencoder)](#model-autoencoder) defines the final autoencoder model and trains it. Sections [KNN classifier](#knn-classifier), [KMeans classifier](#kmeans-classifier) and [NN classifier](#nn-classifier) instead define some classifiers which directly use AE features and try to classify data. The [Grid Search](#grid-search) section is responisble for the grid-search on hyperparameters in order to get the best autoencoder model.
#
# Table of contents
#
# * [Load Dataset](#load-dataset)
# * [Preprocessing](#preprocessing)
# * [Data Exploration](#data-exploration)
# * [Model (Autoencoder)](#model-autoencoder)
# * [KNN classifier](#knn-classifier)
# * [KMeans classifier](#kmeans-classifier)
# * [NN classifier](#nn-classifier)
# * [Grid Search](#grid-search)

# Set some global parameters. We use this parameters to change the notebook behaviour.

# +
# Whether to rotate data in an orientation indipendent manner. Influences the preproecssng stage.
#
# Default: True
USE_ORIENTATION_INDIPENDENT_TRANSFORMATION = True

# Whether to normalize data. Influences the preproecssng stage.
# Note: we usually do not normalize beacause the distance between points
# and reconstructed points gets reduced but the signal is not well represented
#
# Default: False
USE_NORMALIZATION = False

# Whether to center accelerometer data. Influences the preproecssng stage.
# Note: we tried centering accelerometer data to check autoencoder behaviour with CNN center data
#       but this is not useful at all for autoencoder.
#
# Default: False
USE_CENTERING_ACC = False

# Whether to apply low-pass filter on data. Influences the preproecssng stage.
#
# Default: False
USE_LOW_PASS_FILTER = False

# Whether to use Bryan dataset as validation dataset. Influences the validation stage.
#
# Default: True
USE_BRYAN_VALIDATION_DATASET = True

# Whether to use all positions (right, left, back, top, hand, pocket) from Bryan dataset.
# Influences the validation stage.
#
# Default: Flase
USE_ALL_POSITION_BRYAN_DATASET = False

# Whether to use only positions hand and pocket from Bryan dataset.
# Influences the validation stage.
#
# Default: Flase
USE_ONLY_HAND_POCKET_POSITION_BRYAN_DATASET = False

# Whether to use grid search for KNN hyperparameters. Influences the grid-search stage.
#
# Default: False
USE_KNN_GRID_SEARCH = False

# Whether to use grid search for autoencoder hyperparameters. Influences the grid-search stage.
# Note: this task is resource intensive and could take a while. Keep it disabled for prototyping.
#
# Default: True
USE_GRID_SEARCH = False

# Whether to test some of the best models obtained by grid search. Influences the grid-search stage.
# This should be set to True if you have some models to test or when `USE_GRID_SEARCH` is True.
#
# Default: True
USE_GRID_SEARCH_VALIDATION = False

# Whether to merge sit and stand classes. Influences the preprocessing stage.
#
# Default: True
USE_MERGE_SIT_STAND = False

# Whether to merge stairs up/down classes. Influences the preprocessing stage.
#
# Default: False
USE_MERGE_STAIRS_UP_DOWN = False


# -

def get_settings():
    """
    Return a string which encodes current global settings.
    Can be used as name of files.

    Note: only "relevant" configs are included in the output.
    """
    oit = "_oit" if USE_ORIENTATION_INDIPENDENT_TRANSFORMATION else ""
    merge = "_not-merge" if not USE_MERGE_SIT_STAND else ""
    norm = "_norm" if USE_NORMALIZATION else ""
    bryan = "_bryan" if USE_BRYAN_VALIDATION_DATASET else ""
    allpos = "_allpos" if USE_BRYAN_VALIDATION_DATASET and USE_ALL_POSITION_BRYAN_DATASET else ""
    return oit + norm + bryan + allpos


# +
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import random
import tensorflow as tf
from tensorflow.keras import layers

from autoencoder_utils import show_samples, show_loss, show_mse, show_reconstructed_signals, show_reconstruction_errors
from keras_utils import ModelSaveCallback
from orientation_indipendend_transformation import orientation_independent_transformation
# -

random.seed(42)
np.random.seed(42)


# ## Load dataset

# +
def load_dataset(data_filename="dataset.csv", label_filename="labels.csv", dataset_dir="./datasets"):
    """
    Load a dataset from data and label filenames.
    """
    data = pd.read_csv(os.path.join(dataset_dir, data_filename), header=None)
    labels = pd.read_csv(os.path.join(dataset_dir, label_filename), header=None, names=["user", "model", "label"])

    return data, labels


def load_dataset_position(data_filename="dataset.csv", label_filename="labels.csv", dataset_dir="./datasets"):
    """
    Load a dataset from data and label filenames.
    Label file contains user, (smartphone) model, label, and (sensor) position.
    """
    data = pd.read_csv(os.path.join(dataset_dir, data_filename), header=None)
    labels = pd.read_csv(os.path.join(dataset_dir, label_filename), header=None, names=["user", "model", "label", "position"])

    return data, labels


# -

def print_stats(ds: pd.DataFrame):
    print("Shape", ds.shape)
    print("Columns", ds.columns)


X_df_reference, y_df_reference = load_dataset(dataset_dir="./datasets/heterogeneity_f50_w2.5_o0.5")
X_df_reference_validation, y_df_reference_validation = load_dataset_position(dataset_dir="./datasets/bryan_f50_w2.5_o0.5")

print_stats(X_df_reference)
print_stats(y_df_reference)
print_stats(X_df_reference_validation)
print_stats(y_df_reference_validation)


# ## Preprocessing

def to_numpy(df):
    return df.loc[:].to_numpy()


def get_label(x):
    return x[..., 2]


def restructure(x):
    return x.reshape(-1, 6, 125)


def normalize(x):
    min_val = np.max(x)
    max_val = np.min(x)
    
    x = (x - min_val) / (max_val - min_val)
    
    return x


def center(x):
    means = np.mean(x, axis=2)  # get mean for each dimension for each window
    means[:, 3:] = 0  # do not center gyroscope
    x = (x - means[:,:,np.newaxis])


# +
from scipy.signal import butter, filtfilt, freqz
import matplotlib.pyplot as plt


def butter_lowpass(cutoff, fs, order=5):
    nyq = 0.5 * fs
    normal_cutoff = cutoff / nyq
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
    return b, a


def butter_lowpass_filter(data, cutoff, fs, order=5):
    b, a = butter_lowpass(cutoff, fs, order=order)
    y = filtfilt(b, a, data)
    return y


# -

def plot_butter_lowpass_filter(cutoff, fs, order):
    # Get the filter coefficients so we can check its frequency response.
    b, a = butter_lowpass(cutoff, fs, order)

    # Plot the frequency response.
    w, h = freqz(b, a, worN=8000)
    plt.subplot(2, 1, 1)
    plt.plot(0.5*fs*w/np.pi, np.abs(h), 'b')
    plt.plot(cutoff, 0.5*np.sqrt(2), 'ko')
    plt.axvline(cutoff, color='k')
    plt.xlim(0, 0.5*fs)
    plt.title("Lowpass Filter Frequency Response")
    plt.xlabel('Frequency [Hz]')
    plt.grid()


# +
def get_X(df):
    """
    Return a numpy array with preprocessed X data
    """
    # 1. Back to numpy
    a = to_numpy(df)

    # 2. Restructure the array
    a = restructure(a)

    # Low-pass filter
    if USE_LOW_PASS_FILTER:
        for i in range(len(a)):
            window = a[i]
            lowpassed = [butter_lowpass_filter(window[j], cutoff=5, fs=50) for j in range(6)]
            a[i] = lowpassed

    # 3. Normalize
    if USE_NORMALIZATION:
        a = normalize(a)
        
    # 4. Orientation indipendent transformation
    if USE_ORIENTATION_INDIPENDENT_TRANSFORMATION:
        a = orientation_independent_transformation(a)

    # 4. Center accelerometer data
    if USE_CENTERING_ACC:
        a = center(a)
    
    return a


def get_y(df):
    """
    Return a numpy array with labels
    """
    y = to_numpy(df)
    y = get_label(y)
    return y


def get_y_hot(df):
    """
    Return a numpy array with labels in one-hot encoding
    """
    return pd.get_dummies(df['label']).to_numpy()


# -

# ### Training Set

X_df, y_df = X_df_reference.copy(), y_df_reference.copy()

# Set up the low-pass filter

# +
cutoff = 1  # desired cutoff frequency of the filter, Hz
fs = 25     # sample rate, Hz
order = 3

plot_butter_lowpass_filter(cutoff, fs, order)
plt.show()
# -

# Prepare the dataset with pandas
# * split training and test set
# * shuffle the dataset

# +

if USE_MERGE_SIT_STAND:
    # Merge sit and stand labels
    sit_or_stand_filter = (y_df["label"] == "sit") | (y_df["label"] == "stand")
    y_df["label"].loc[sit_or_stand_filter] = "no_activity"

if USE_MERGE_STAIRS_UP_DOWN:
    # Merge stairs activity
    stairsdown_or_stairsup_filter = (y_df["label"] == "stairsdown") | (y_df["label"] == "stairsup")
    y_df["label"].loc[stairsdown_or_stairsup_filter] = "stairs"

# *** SHUFFLE
X_shuffled_df = X_df.sample(frac=1)
y_shuffled_df = y_df.reindex(X_shuffled_df.index)

# *** TRAIN AND TEST
if USE_BRYAN_VALIDATION_DATASET:
    but_last_user_indicies = ~(y_df['user'] == "z")
else:
    but_last_user_indicies = ~((y_df['user'] == "a") | (y_df['user'] == "b"))

X_train_df = X_shuffled_df.loc[but_last_user_indicies]
X_test_df = X_shuffled_df.loc[~but_last_user_indicies]

y_train_df = y_shuffled_df.loc[but_last_user_indicies]
y_test_df = y_shuffled_df.loc[~but_last_user_indicies]

print("X_train_df =", len(X_train_df))
print("X_test_df =", len(X_test_df))
print("y_train_df =", len(y_train_df))
print("y_test_df =", len(y_test_df))

assert len(X_train_df) == len(y_train_df), "X train and y train do not contain same number of samples"
assert len(X_test_df) == len(y_test_df), "X test and y test do not contain same number of samples"

# -

# Preprocess dataset

# + tags=[]
# Preprocess and prepare training and test set
X_train, y_train = get_X(X_train_df), get_y(y_train_df)
X_test, y_test = get_X(X_test_df), get_y(y_test_df)

# Retrieve al
y_train_hot = get_y_hot(y_train_df)
y_test_hot = get_y_hot(y_test_df)
# -

print(X_train.shape)
print(y_train.shape)
print(X_test.shape)
print(y_test.shape)

# Show labels

# +
classes = np.unique(y_train)
num_classes = len(np.unique(y_train))

print(f"Classes = {classes}")
print(f"Num classes = {num_classes}")
# -

# Plot some samples

show_samples(X_train, y_train, n=1, is_random=False)

if (len(X_test) > 0):
    show_samples(X_test, y_test, n=1, is_random=False)
else:
    print("WARNING: there is no data in X_test. Are you using Bryan dataset as validation set?")

# ### Validation Set

# Prepare the dataset

# + tags=[]
X_df, y_df = X_df_reference_validation.copy(), y_df_reference_validation.copy()

# Keep valid positions only
if USE_ALL_POSITION_BRYAN_DATASET:
    valid_positions_indicies = ~(y_df["position"] == "none")
elif USE_ONLY_HAND_POCKET_POSITION_BRYAN_DATASET:
    valid_positions_indicies = (y_df["position"] == "hand") | (y_df["position"] == "pocket-up") | (y_df["position"] == "pocket-down")
else:
    valid_positions_indicies = (y_df["position"] == "right") | (y_df["position"] == "left") | (y_df["position"] == "top") | (y_df["position"] == "back") | (y_df["position"] == "bottom")



# Filter out other data
X_df = X_df.loc[valid_positions_indicies]
y_df = y_df.loc[valid_positions_indicies]

# Shuffle data
X_df = X_df.sample(frac=1)
y_df = y_df.reindex(X_df.index)
# -

# Preprocess dataset

# + tags=[]
X_validation = get_X(X_df)
y_validation = get_y(y_df)
y_validation_hot = get_y_hot(y_df)

# +
X_validation_df = X_df.copy()
y_validation_df = y_df.copy()

del X_df, y_df
# -

print(X_validation.shape)
print(y_validation.shape)
print(y_validation_hot.shape)

show_samples(X_validation, y_validation, n=1, is_random=False)

if USE_BRYAN_VALIDATION_DATASET :
    X_test = X_validation
    y_test = y_validation
    y_test_hot = y_validation_hot

show_samples(X_test, y_test, n=1, is_random=False)

# Check size

assert X_train.shape[0] == y_train.shape[0], f"Invalid shape for X_train and y_train: {X_train.shape} != {y_train.shape}"
assert X_test.shape[0] == y_test.shape[0], f"Invalid shape for X_test and y_test: {X_test.shape} != {y_test.shape}"
assert X_test.shape[0] == y_test.shape[0], f"Invalid shape for X_test and y_test: {X_test.shape} != {y_test.shape}"
assert y_train_hot.shape == (y_train.shape[0],num_classes), f"Invalid shape of y_train_hot: {y_train_hot.shape}"
assert y_test_hot.shape == (y_test.shape[0],num_classes), f"Invalid shape of y_test_hot: {y_test_hot.shape}"

# ## Data Exploration

# ### Training Set

print("Users", y_train_df["user"].unique())
print("Models", y_train_df["model"].unique())
print("Classes", y_train_df["label"].unique())

# Fraction of samples per label

print(y_train_df.groupby(["label"])["label"].count() / y_train_df["label"].count())

# Fraction of samples per user

print(y_train_df.groupby(["user"])["user"].count() / y_train_df["user"].count())

# Fraction of samples per model

print(y_train_df.groupby(["model"])["model"].count() / y_train_df["model"].count())

# Number of samples per user i and fraction of samples per class for user i

# +
y_df_i = y_train_df.loc[y_train_df["user"] == "i"]

num_samples_i = y_df_i["label"].count()
fraction_of_samples_per_class_i = y_df_i.groupby(["label"])["label"].count() / y_df_i["label"].count()

print(num_samples_i)
print(fraction_of_samples_per_class_i)
# -

# ### Validation Set

print("Classes (validation)", y_validation_df["label"].unique())

print(y_validation_df.groupby(["label"])["label"].count() / y_validation_df["label"].count())

print(y_validation_df.groupby(["user"])["user"].count() / y_validation_df["user"].count())

print(y_validation_df.groupby(["model"])["model"].count() / y_validation_df["model"].count())

# ## Model (autoencoder)

DATA_SHAPE = X_train.shape[1:]
CODE_SIZE=36


def build_encoder(data_shape, code_size):
    inputs = tf.keras.Input(data_shape)

    X = inputs    
    X = layers.Flatten()(X)
    X = layers.Dense(150, activation="relu")(X)
    X = layers.Dense(code_size, activation="sigmoid")(X)
    outputs = X
    
    return tf.keras.Model(inputs=inputs, outputs=outputs)


# +

def build_decoder(data_shape, code_size):
    inputs = tf.keras.Input((code_size,))

    X = inputs
    X = layers.Dense(150, activation="relu")(X)
    X = layers.Dense(np.prod(data_shape), activation=None)(X)
    outputs = layers.Reshape(data_shape)(X)

    return tf.keras.Model(inputs=inputs, outputs=outputs)


# -

def build_autoencoder(encoder, decoder):
    inputs = tf.keras.Input(DATA_SHAPE)  # input
    codes = encoder(inputs)              # build the code with the encoder
    outputs = decoder(codes)             # reconstruction the signal with the decoder

    return tf.keras.Model(inputs=inputs, outputs=outputs)


encoder = build_encoder(DATA_SHAPE, CODE_SIZE)
decoder = build_decoder(DATA_SHAPE, CODE_SIZE)

# + tags=[]
autoencoder = build_autoencoder(encoder, decoder)

optimizer = "adam"
loss = "mse"

model_filename = 'autoencoder_network.hdf5'
last_finished_epoch = None
epochs=100
batch_size=128

early_stopping_callback = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=5)
save_model_checkpoint_callback = ModelSaveCallback(model_filename)

callbacks = [save_model_checkpoint_callback, early_stopping_callback]

autoencoder.compile(optimizer=optimizer, loss=loss)

history = autoencoder.fit(
    x=X_train, y=X_train, 
    epochs=epochs, 
    validation_data=(X_test, X_test), 
    batch_size=batch_size,
    callbacks=callbacks,
    verbose=1)
# -

encoder.save("encoder.h5")

show_loss(history)

show_mse(autoencoder, X_test)

show_reconstructed_signals(X_test, encoder, decoder, n=1)

show_reconstruction_errors(X_test, encoder, decoder, n=1)

# ## KNN classifier

# +
from sklearn.neighbors import KNeighborsClassifier

# prepare the codes
codes = encoder.predict(X_train)

assert codes.shape[1:] == (CODE_SIZE,), f"Predicted codes shape must be equal to code size, but {codes.shape[1:]} != {(CODE_SIZE,)}"

# create the k-neighbors calssifier
n_neighbors = 5
metric = "euclidean"

nbrs = KNeighborsClassifier(n_neighbors=n_neighbors, metric=metric)

# fit the model using the codes
nbrs.fit(codes, y_train)
# -

# Clustering classes

print("Classes =", nbrs.classes_)

# Show some predictions

# +
print("X_test[i] = y_true \t y_pred    \t with probs [...]")
print()

for i in range(20):
    x = X_test[i]
    y = y_test[i]

    c = encoder.predict(x[np.newaxis, :])[0]
    
    [lab] = nbrs.predict(c[np.newaxis, :])
    [probs] = nbrs.predict_proba(c[np.newaxis, :])

    print(f"X_test[{i}] = {y}\t {lab}   \t with probs {probs}")

# -

# Show performance metrics

# +
from sklearn.metrics import classification_report
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score

codes = encoder.predict(X_test)

y_true = y_test
y_pred = nbrs.predict(codes)

a = accuracy_score(y_true=y_true, y_pred=y_pred)
f1 = f1_score(y_true=y_true, y_pred=y_pred, average="weighted")
p = precision_score(y_true=y_true, y_pred=y_pred, average="weighted")
r = recall_score(y_true=y_true, y_pred=y_pred, average="weighted")

print(classification_report(y_true=y_true, y_pred=y_pred))
print(f"accuracy = {a}")
print(f"precision = {p}")
print(f"recall = {r}")
print(f"f1_score = {f1}")
# -

# Grid search for KNN

if USE_KNN_GRID_SEARCH:
    from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score

    results = pd.DataFrame()

    codes_train = encoder.predict(X_train)
    codes_test = encoder.predict(X_test)
    y_true = y_test

    for metric in ["euclidean", "manhattan", "chebyshev", "minkowski", "seuclidean", "mahalanobis"]:
        for n_neighbors in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 14, 16, 18, 20, 25, 30, 40, 100]:
            print(f"metric={metric}, n_neighbors={n_neighbors}")

            nbrs = KNeighborsClassifier(n_neighbors=n_neighbors, metric=metric)

            # fit the model using the codes
            nbrs.fit(codes_train, y_train)

            # predict
            y_pred = nbrs.predict(codes_test)

            a = accuracy_score(y_true=y_true, y_pred=y_pred)
            f1 = f1_score(y_true=y_true, y_pred=y_pred, average="weighted")
            p = precision_score(y_true=y_true, y_pred=y_pred, average="weighted")
            r = recall_score(y_true=y_true, y_pred=y_pred, average="weighted")

            data = [{"metric": metric, "n_neighbors": n_neighbors, "accuracy": a, "f1": f1, "precision": p, "recall": r }]
            results = results.append(pd.DataFrame(data), ignore_index=True)

if USE_KNN_GRID_SEARCH:
    rr = results.loc[(results["n_neighbors"] >= 3) & (results["n_neighbors"] <= 8)]
    rr = rr.groupby("metric")["accuracy"].max()
    print(rr)
    print(rr.max())

    results.to_csv("./models/knn_grid_search.csv")

# ## KMeans classifier

# + tags=[]
from sklearn.cluster import KMeans
from sklearn.preprocessing import LabelEncoder
import numpy as np

le = LabelEncoder()
le.fit(y_test)

# train
codes = encoder.predict(X_train)

kmeans = KMeans(n_clusters=num_classes, random_state=0)
kmeans.fit(codes)

# evaluate 
codes = encoder.predict(X_test)

y_true = y_test
y_pred = le.inverse_transform(kmeans.predict(codes))

print(classification_report(y_true=y_true, y_pred=y_pred))


# -

# ## NN classifier

# +
def build_nn(code_size):
    inputs = tf.keras.Input((code_size,))
    X = inputs
    X = layers.Dense(100, activation="relu")(X)
    X = layers.Dropout(0.1)(X)
    X = layers.Dense(100, activation="relu")(X)
    X = layers.Dropout(0.1)(X)
    X = layers.Dense(num_classes, activation="softmax")(X)
    outputs = X

    return tf.keras.Model(inputs=inputs, outputs=outputs)


codes_train = encoder.predict(X_train)
codes_test = encoder.predict(X_test)

nn_model = build_nn(CODE_SIZE)

adam_optimizer = tf.keras.optimizers.Adam()
loss_funct = tf.keras.losses.CategoricalCrossentropy()

early_stopping_callback = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=5)
callbacks = [early_stopping_callback]

nn_model.compile(optimizer=adam_optimizer, loss=loss_funct, metrics=["accuracy"])
nn_model.summary()

history = nn_model.fit(x=codes_train, y=y_train_hot,
    epochs=50,
    validation_data=(codes_test, y_test_hot), 
    batch_size=128,
    callbacks=callbacks,
    verbose=1)

# +
from sklearn.metrics import accuracy_score
from sklearn.metrics import classification_report
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
from sklearn.preprocessing import LabelEncoder

codes = encoder.predict(X_test)

y_true = LabelEncoder().fit_transform(y_test)
y_pred = np.argmax(nn_model.predict(codes), axis=1)
print(y_pred)

a = accuracy_score(y_true=y_true, y_pred=y_pred)
f1 = f1_score(y_true=y_true, y_pred=y_pred, average="weighted")
p = precision_score(y_true=y_true, y_pred=y_pred, average="weighted")
r = recall_score(y_true=y_true, y_pred=y_pred, average="weighted")

loss, accuracy = nn_model.evaluate(codes, y_test_hot)

print("LOSS =", loss)
print(classification_report(y_true=y_true, y_pred=y_pred))
print(f"accuracy = {accuracy} = {a}")
print(f"precision = {p}")
print(f"recall = {r}")
print(f"f1_score = {f1}")

show_loss(history)


# -

# # Grid Search
#

# We are searching for best hyperparams. More precisely we want to find the best autoencoder model given
#
# * code size
# * optimizer
# * loss func
# * epochs
# * batch size

# + tags=[]
def get_filename(hparams):
    return f'models/encoder{get_settings()}_cs-{hparams["code_size"]}_loss-{hparams["loss_func"]}_bs-{hparams["batch_size"]}.h5'

def compile_autoencoder(autoencoder, hparams):
    """
    Compile autoencoder with given hyperparams.
    """

    optimizer = hparams["optimizer"]
    loss = hparams["loss_func"]

    autoencoder.compile(optimizer=optimizer, loss=loss)

    return autoencoder


def fit_autoencoder(autoencoder, hparams):
    """
    Fit the autoencoder with given hyperparams on `X_train` with validation `X_test`.
    Please note that the best and last models are saved in `models` folder.
    """

    early_stopping_callback = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=3)
    callbacks = [early_stopping_callback]

    history = autoencoder.fit(
        x=X_train, y=X_train, 
        epochs=hparams["epochs"], 
        validation_data=(X_test, X_test), 
        batch_size=hparams["batch_size"],
        callbacks=callbacks,
        verbose=0)

    return history


def evaluate_autoencoder(autoencoder):
    """
    Evaluate the autoencoder with `X_test`.
    """
    return autoencoder.evaluate(X_test, X_test, verbose=0)


def run_model(hparams):
    """
    Setup, train and evaluate the autoencoder model with given hyperparams.

    :return: A tuple `(loss, hystory)`
    """
    encoder = build_encoder(DATA_SHAPE, hparams["code_size"])
    decoder = build_decoder(DATA_SHAPE, hparams["code_size"])

    autoencoder = build_autoencoder(encoder, decoder)
    autoencoder = compile_autoencoder(autoencoder, hparams)

    history = fit_autoencoder(autoencoder, hparams)
    loss = evaluate_autoencoder(autoencoder)

    # Note: we may also add the classifier training and evaluation 
    #       to select the best model

    filepath = get_filename(hparams)
    encoder.save(filepath)

    return loss, history


def select_model():
    """
    Run the model with hyperparams and save results.

    :return: An array of results with mse, hparams and training history.
    """
    results = []

    hparams_code_sizes = [2, 3, 4, 5, 6, 12, 18, 24, 30, 36, 42, 48, 54, 60, 72]
    hparams_opts = ["adam"]
    hparams_losses = ["mse"]
    hparams_epochs = [150]
    hparams_batch_sizes = [32, 128]

    n_iterations = np.prod([len(hparams_code_sizes), len(hparams_opts), len(hparams_losses), len(hparams_epochs), len(hparams_batch_sizes)])

    print(f"Starting model selection... {n_iterations} iterations needed.") 

    for code_size in hparams_code_sizes:
        for opt in hparams_opts:
            for loss in hparams_losses:
                for epochs in hparams_epochs:
                    for batch_size in hparams_batch_sizes:

                        hparams = {
                            "code_size": code_size,
                            "optimizer": opt,
                            "loss_func": loss,
                            "epochs": epochs,
                            "batch_size": batch_size
                        }

                        print(f"Starting run {len(results)}")
                    
                        loss_val, history = run_model(hparams)

                        print(f"hparams = {hparams}")
                        print(f"{loss} = {loss_val}")
                        print()

                        results += [{"hparams": hparams, "loss_val": loss_val, "history": history}]

    return results


# -

# Select the best model based on lowest loss.

# +
import time

results = []

if USE_GRID_SEARCH:
    start = time.time()
    results = select_model()
    end = time.time()

    print(f"Done in {end - start} s")


# -

# Print a CSV table for models and save them to file

# +
def get_csv(results):
    if ():
        return ""
    
    s = ",".join(["i", "type", "loss_val"] + [k for k,_ in results[0]["hparams"].items()] + ["filepath"]) + "\n"
    
    for i, result in zip(range(len(results)), results):
        hparams = result["hparams"]
        code_size = hparams["code_size"]
        batch_size = hparams["batch_size"]
        loss_val = result["loss_val"]

        filepath = get_filename(hparams)
        type = get_settings().lstrip("_") if get_settings() != "" else "default"

        s += ",".join([str(i), type, str(loss_val)] + [str(v) for _,v in hparams.items()] + [filepath]) + "\n"
    
    return s


def save_results(results, filename=f'models/res{get_settings()}.csv'):
    with open(filename, mode="w+") as f:
        content = get_csv(results)
        f.write(content)

if len(results) > 0:
    print(get_csv(results))

    if USE_GRID_SEARCH:
        save_results(results)
    else:
        print("WARNING: Results were not saved to file.")


# -

# Check best model

# + tags=[]
def test_best_model(best_filename, code_size):
    
    def run_nn(codes_train, codes_test):
        nn_model = build_nn(code_size)

        adam_optimizer = tf.keras.optimizers.Adam()
        loss_funct = tf.keras.losses.CategoricalCrossentropy()

        early_stopping_callback = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=3)
        callbacks = [early_stopping_callback]

        nn_model.compile(optimizer=adam_optimizer, loss=loss_funct, metrics=["accuracy"])
        #nn_model.summary()

        history = nn_model.fit(x=codes_train, y=y_train_hot,
            epochs=100,
            validation_data=(codes_test, y_test_hot), 
            batch_size=128,
            callbacks=callbacks,
            verbose=0)

        loss, accuracy = nn_model.evaluate(codes_test, y_test_hot)

        return accuracy, loss

        #show_loss(history)
    

    def run_knn(codes_train, codes_test):
        from sklearn.neighbors import KNeighborsClassifier

        # create the k-neighbors calssifier
        n_neighbors = num_classes
        metric = "euclidean"

        nbrs = KNeighborsClassifier(n_neighbors=n_neighbors, metric=metric)

        # fit the model using the codes
        nbrs.fit(codes_train, y_train)

        from sklearn.metrics import mean_squared_error
        from sklearn.metrics import accuracy_score

        y_true = y_test
        y_pred = nbrs.predict(codes_test)

        #loss = mean_squared_error(y_true=y_true, y_pred=y_pred)
        accuracy = accuracy_score(y_true=y_true, y_pred=y_pred)

        return accuracy


    encoder = tf.keras.models.load_model(best_filename, compile=False)
    encoder.compile(optimizer="adam", loss="mse")

    codes_train = encoder.predict(X_train, verbose=0)
    codes_test = encoder.predict(X_test, verbose=0)

    accuracy_knn = run_knn(codes_train, codes_test)
    accuracy_nn, loss_nn = run_nn(codes_train, codes_test)

    return {"accuracy_knn": accuracy_knn, "accuracy_nn": accuracy_nn, "loss_nn": loss_nn }
    

if USE_GRID_SEARCH_VALIDATION:
    validations_df = pd.DataFrame(columns=["accuracy_knn", "accuracy_nn", "loss_nn"])

    res_df = pd.read_csv(f"models/res{get_settings()}.csv")
    #res_oit_df = pd.read_csv("./models/res_oit.csv.txt")

    #data_df = res_df.append(res_oit_df, ignore_index=True)
    data_df = res_df
    
    for index, row in data_df.iterrows():
        print(f"Validating model {index} ({row['filepath']})")

        test_res = test_best_model(row["filepath"], row["code_size"])

        validations_df = validations_df.append(pd.DataFrame([{ **{"filepath": row['filepath']}, **test_res }]), ignore_index=True)
    
    print(validations_df)

    validations_df.to_csv(f"./models/validation{get_settings()}.csv")
    print("Done")
# -

fn = "models/manual/encoder_oit_bryan_allpos_cs-30_bs-128.h5"
print(test_best_model(fn, 30))

StarcoderdataJupyterScriptsDedupFiltered