# Preliminary steps
# Common imports
import numpy as np
import pandas as pd

# Import sklearn
import sklearn
from sklearn.datasets import load_files
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
assert sklearn.__version__ >= "0.20"

# TensorFlow ≥2.0 is required
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import datasets, layers, models
assert tf.__version__ >= "2.0"

# To plot pretty figures
import seaborn as sns
import matplotlib as mpl
import matplotlib.pyplot as plt
mpl.rc('axes', labelsize=14)
mpl.rc('xtick', labelsize=12)
mpl.rc('ytick', labelsize=12)

# To make this notebook's output stable across runs
np.random.seed(42)
tf.random.set_seed(42)

# Loading data for the new/worn blade
new = [pd.read_csv("data/NewBlade00%d.csv"%i) for i in range(1,4)]
worn = [pd.read_csv("data/WornBlade00%d.csv"%i) for i in range(1,4)]

new[0].head()

new[0].info()

new[0].describe()

# Preprocessor for the data
scaler = StandardScaler()
scaler.fit(np.concatenate([new_.values for new_ in new]))

# Give each column a name
columns = ["timestamp", "cut torque", "cut lag error", "cut position",
           "cut speed", "film position", "film speed",
           "film lag error"]

for i in range(3):
    new[i].iloc[:,1:] = scaler.transform(new[i].astype(float))[:,1:]
    worn[i].iloc[:,1:] = scaler.transform(worn[i].astype(float))[:,1:]
    new[i].columns = columns
    worn[i].columns = columns
    print("New  {0}, shape={1}".format(i+1, new[i].shape))
    print("Worn {0}, shape={1}".format(i+1, worn[i].shape))

data_normal = np.concatenate([new_.values for new_ in new])
n = len(data_normal)

valid_size = 20.
valid_size = int(valid_size*n/100.)

x_train = data_normal[:-valid_size, 1:]
x_valid = data_normal[-valid_size:, 1:]

print("Train data: {0}".format(x_train.shape))
print("Valid data: {0}".format(x_valid.shape))

data_worn = np.concatenate([worn_.values for worn_ in worn])
n = len(data_worn)

x_worn = data_worn[:, 1:]
print("Data worn: {0}".format(x_worn.shape))

# Define encoder
encoder = keras.models.Sequential([
    keras.layers.Dense(256, activation="elu"),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(32, activation="elu"),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(5, activation="elu")
    ])

# Define decoder
decoder = keras.models.Sequential([
    keras.layers.Dense(32, activation="elu", input_shape = [5] ),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(256, activation="elu"),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(x_train.shape[1])
    ])

# Put together
autoencoder = keras.models.Sequential([
    encoder, decoder
    ])

# Define optimizer and compile model
autoencoder.compile(optimizer="adam", loss="mse")

history = autoencoder.fit(np.array(x_train),np.array(x_train),
                epochs=40,
                batch_size=64,
                shuffle=True,
                validation_split=0.1,
               )

plt.plot(history.history["loss"], label="Training Loss")
plt.plot(history.history["val_loss"], label="Validation Loss")
plt.xlabel('Epoch')
plt.ylabel('mse')
plt.legend()

# Transform input into latent space
z_train = encoder.predict(x_train)
z_valid = encoder.predict(x_valid)
z_worn = encoder.predict(x_worn)

fig = plt.figure(figsize=(14, 8))
ax = fig.add_subplot(projection='3d')

for m, color, label, scale, z in [('o', 'tab:green', 'Train', 0.5, z_train), ('x', 'tab:blue', 'Valid', 0.5, z_valid), ('^', 'tab:red', 'Worn', 0.5, z_worn)]:
    xs = z[:,0]
    ys = z[:,3]
    zs = z[:,4]
    ax.scatter(xs, zs, ys, marker=m, s=scale, c=color, label=label)

ax.set_xlabel('Feature 1')
ax.set_ylabel('Feature 2')
ax.set_zlabel('Feature 3')
ax.legend()

plt.show()

columns_ = ["cut torque", "cut lag error", "cut position",
           "cut speed", "film position", "film speed",
           "film lag error"]

# Reconstruction train data
x_pred = autoencoder.predict(np.array(x_train))

# Creating dataframe for reconstructed data
x_pred = pd.DataFrame(x_pred,columns = columns_)
x_pred.index = pd.DataFrame(x_train).index

# Dataframe to get the difference of predicted data and real data. 
scored_normal = pd.DataFrame(index = pd.DataFrame(x_train).index)

# Returning the mean of the loss for each column
scored_normal['Loss_mae'] = np.mean(np.abs(x_pred - x_train), axis = 1)

# Selecting a Threshold
train_loss = np.mean(np.abs(x_pred - x_train), axis = 1)
threshold = np.mean(train_loss) + np.std(train_loss)
print("Threshold: ", threshold)

# Reconstruction worn data
x_pred = autoencoder.predict(np.array(x_worn))

# Creating dataframe for reconstructed data
x_pred = pd.DataFrame(x_pred,columns = columns_)
x_pred.index = pd.DataFrame(x_worn).index

# Dataframe to get the difference of predicted data and real data. 
scored_worn = pd.DataFrame(index = pd.DataFrame(x_worn).index)

# Returning the mean of the loss for each column
scored_worn['Loss_mae'] = np.mean(np.abs(x_pred - x_worn), axis = 1)

# Plot reconstruction error for worn (red) and normal (green) data
plt.subplots(figsize=(15,7))
sns.distplot(scored_worn['Loss_mae'],
            bins = 15,
            kde= True,
            color = 'tab:red')
sns.distplot(scored_normal['Loss_mae'],
            bins = 15,
            kde= True,
            color = 'tab:green')

error_normal = np.mean(np.abs(autoencoder.predict(data_normal[:,1:]) - data_normal[:,1:]), axis=1)
error_worn = np.mean(np.abs(autoencoder.predict(data_worn[:,1:]) - data_worn[:,1:]), axis=1)
error_normal_worn = np.concatenate((error_normal, error_worn))

def moving_average(x, w):
    return np.convolve(x, np.ones(w), 'valid') / w

error_normal_worn_avg = moving_average(error_normal_worn, 200)

plt.grid()
plt.plot(error_normal_worn)
plt.plot(error_normal_worn_avg)
plt.title("Rekonstruktionsfehler")
plt.show()

n = len(data_normal)

max_normal = np.amax(error_normal)
max_full = max(max_normal, np.amax(error_worn))

sh_train = 0.95
sh_worn = 0.05
expit = lambda x: 1./(1 + np.exp(-x))
logit = lambda p: np.log(p/(1 - p))

a = -(logit(sh_train) - logit(sh_worn))/(max_full - max_normal)
b = -(logit(sh_worn)*max_normal - logit(sh_train)*max_full)/(max_full - max_normal)
system_health = lambda error: 100*expit(a*error + b)

health_normal_worn = np.concatenate((system_health(error_normal[:-valid_size]),
    system_health(error_normal[-valid_size:]),
    system_health(error_worn)
    ))

health_normal_worn_avg = moving_average(health_normal_worn, 500)

# Plot
plt.figure(figsize=(14, 4), dpi=80)
plt.plot(health_normal_worn)
plt.plot(health_normal_worn_avg, '--')
plt.axvline(x= n - valid_size, linewidth=2, color='k')
plt.axvline(x= n, linewidth=2, color='k')
plt.xlim(0, n + len(error_worn))

# Description
plt.title('Systemgesundheit')
plt.ylabel('Systemgesundheit in %')
plt.xlabel('Samples')
plt.figtext(0.5, -0.15, 'Systemgesundheit über die Zeit. Alle 6 Datensätze werden nacheinander ausgegeben.\n Die vertikalen Linien geben jeweils den Bereich für train, validation und worn Daten an.', wrap=True, horizontalalignment='center', fontsize=11)
plt.text( 150, 5, 'train')
plt.text( (n - valid_size) + 150, 5, 'valid')
plt.text( n + 150, 5, 'worn')
plt.grid()
plt.show()

_, p = x_train.shape

error_normal_all = np.abs(autoencoder.predict(data_normal[:,1:]) - data_normal[:,1:])
error_worn_all = np.abs(autoencoder.predict(data_worn[:,1:]) - data_worn[:,1:])

distance_all = np.concatenate((error_normal_all, error_worn_all))
fractions = (distance_all.T/np.sum(distance_all, axis=1)).T
t = np.arange(len(fractions))

# Plot
plt.figure(figsize=(16, 4), dpi=80)
for j in range(p):
    lower = np.sum(fractions[:,:j], axis=1)
    upper = np.sum(fractions[:,:j + 1], axis=1)
    plt.fill_between(t, lower, upper, label=new[0].columns[j + 1])

plt.axvline(x= n - valid_size, linewidth=2, color='k')
plt.axvline(x= n, linewidth=2, color='k')
plt.xlim(0, n + len(error_worn))
plt.ylim(0, 1)

# Description
plt.title("Anteile der einzelnen Komponenten an dem Rekonstruktionsfehler")
plt.xlabel('Samples')
plt.ylabel('Rekonstruktionsfehler (normalisiert)')
plt.text( 150, -0.05, 'train', color='k')
plt.text( (n - valid_size) + 150, -0.05, 'valid', color='k')
plt.text( n + 150, -0.05, 'worn', color='k')
plt.legend(loc='center left', bbox_to_anchor=(1, 0.5), ncol=1)
plt.figtext(0.5, -0.17, 'Normalisierter Rekonstruktionsfehler über die Zeit. Alle 6 Datensätze werden nacheinander ausgegeben.\n Die vertikalen Linien geben jeweils den Bereich für train, validation und worn Daten an. \nDer Beitrag zum Rekonstruktionsfehler ist eingefärbt.', wrap=True, horizontalalignment='center', fontsize=11)

means = np.mean(fractions[-1000:], axis = 0)
for i, amean in enumerate(means):
    print("{0: <15s}: {1: 0.2f}%".format(new[0].columns[i + 1], amean * 100))