"""
#### this notebook is part of the documentation on my HPA approach  
    -> main notebook: https://www.kaggle.com/philipjamessullivan/0-hpa-approach-summary


## 7: network training
    -> https://www.kaggle.com/philipjamessullivan/7-train-effnetb0-version-a-part-1  
    -> https://www.kaggle.com/philipjamessullivan/7-train-effnetb1-version-a-part-1  
    -> https://www.kaggle.com/philipjamessullivan/7-train-effnetb2-version-a-part-1  
    -> https://www.kaggle.com/philipjamessullivan/7-train-effnetb3-version-a-part-1  
    
##### Network Architecture: Efficientnet (B0 to B3) with a Dropout layer of 0.3  

#### Goal: further refine model  

##### INPUT: 
**resized cropped RGB cell images**:  

EfficientnetB0 -> 224x224  
EfficientnetB1 -> 240x240  
EfficientnetB2 -> 260x260  
EfficientnetB3 -> 300x300  

**specific subset of ground truth dataframe**:  

Colums: img_id, Label, Cell#, bbox coords, path  

##### OUTPUT:
**model checkpoints**:  

one after each epoch

**training history**:  

shown as printed array on screen with f1 loss and f1 score per step
"""
import pickle
import pandas as pd
hpa_df = pd.read_pickle("../input/gtdataframes/dataset_partial_300.pkl")
hpa_df
LABELS= {
0: "Nucleoplasm",
1: "Nuclear membrane",
2: "Nucleoli",
3: "Nucleoli fibrillar center",
4: "Nuclear speckles",
5: "Nuclear bodies",
6: "Endoplasmic reticulum",
7: "Golgi apparatus",
8: "Intermediate filaments",
9: "Actin filaments",
10: "Microtubules",
11: "Mitotic spindle",
12: "Centrosome",
13: "Plasma membrane",
14: "Mitochondria",
15: "Aggresome",
16: "Cytosol",
17: "Vesicles and punctate cytosolic patterns",
18: "Negative"
}
import tensorflow as tf
print(tf.__version__)

from tensorflow.keras.layers import *
from tensorflow.keras.models import *
import tensorflow_addons as tfa

import os
import re
import cv2
import glob
import numpy as np
import pandas as pd
import seaborn as sns
from functools import partial
import matplotlib.pyplot as plt
IMG_WIDTH = 300
IMG_HEIGHT = 300
BATCH_SIZE = 16

AUTOTUNE = tf.data.experimental.AUTOTUNE
#%%script echo skipping
#set an amount of folds to split dataframe into --> k-fold cross validation
# explanation: https://towardsdatascience.com/cross-validation-explained-evaluating-estimator-performance-e51e5430ff85
N_FOLDS=5
#choose which one of the 5 folds will be used as validation set this time
i_VAL_FOLD=1
hpa_df=np.array_split(hpa_df, N_FOLDS+1) #add one extra part for testing set
df_test_split=hpa_df[-1]
hpa_df=hpa_df[:-1]
i_training = [i for i in range(N_FOLDS)]
i_training.pop(i_VAL_FOLD-1)
i_validation=i_VAL_FOLD-1
df_train_split=list()
for i in i_training:
    df_train_split.append(hpa_df[i])
df_train_split=pd.concat(df_train_split)
df_val_split=hpa_df[i_validation]
print(len(df_train_split))
print(len(df_val_split))
print(len(df_test_split))
#%%script echo skipping
#analyze class imbalance and set up class weights here
#https://www.analyticsvidhya.com/blog/2020/10/improve-class-imbalance-class-weights/
y_train=df_train_split["Label"].apply(lambda x:list(map(int, x.split("|"))))
y_train=y_train.values
y_train=np.concatenate(y_train)
from sklearn.utils import class_weight
class_weights = class_weight.compute_class_weight('balanced',
                                                 np.unique(y_train),
                                                 y_train)
#%%script echo skipping
tmp_dict={}
for i in range(len(LABELS)):
    tmp_dict[i]=class_weights[i]
class_weights=tmp_dict
class_weights
#%%script echo skipping
# adapted from https://www.kaggle.com/ayuraj/hpa-multi-label-classification-with-tf-and-w-b
@tf.function
def multiple_one_hot(cat_tensor, depth_list):
    """Creates one-hot-encodings for multiple categorical attributes and
    concatenates the resulting encodings

    Args:
        cat_tensor (tf.Tensor): tensor with mutiple columns containing categorical features
        depth_list (list): list of the no. of values (depth) for each categorical

    Returns:
        one_hot_enc_tensor (tf.Tensor): concatenated one-hot-encodings of cat_tensor
    """
    one_hot_enc_tensor = tf.one_hot(cat_int_tensor[:,0], depth_list[0], axis=1)
    for col in range(1, len(depth_list)):
        add = tf.one_hot(cat_int_tensor[:,col], depth_list[col], axis=1)
        one_hot_enc_tensor = tf.concat([one_hot_enc_tensor, add], axis=1)
    return one_hot_enc_tensor

@tf.function
def load_image(df_dict):
    # Load image
    rgb = tf.io.read_file(df_dict['path'])
    image = tf.image.decode_png(rgb, channels=3)
    #https://medium.com/@kyawsawhtoon/a-tutorial-to-histogram-equalization-497600f270e2
    image=tf.image.per_image_standardization(image)
    
    # Parse label
    label = tf.strings.split(df_dict['Label'], sep='|')
    label = tf.strings.to_number(label, out_type=tf.int32)
    label = tf.reduce_sum(tf.one_hot(indices=label, depth=19), axis=0)
    
    return image, label
#%%script echo skipping
train_ds = tf.data.Dataset.from_tensor_slices(dict(df_train_split))
val_ds = tf.data.Dataset.from_tensor_slices(dict(df_val_split))

# Training Dataset
train_ds = (
    train_ds
    .shuffle(1024)
    .map(load_image, num_parallel_calls=AUTOTUNE)
    .batch(BATCH_SIZE)
    .prefetch(tf.data.experimental.AUTOTUNE)
)

# Validation Dataset
val_ds = (
    val_ds
    .shuffle(1024)
    .map(load_image, num_parallel_calls=AUTOTUNE)
    .batch(BATCH_SIZE)
    .prefetch(tf.data.experimental.AUTOTUNE)
)
#%%script echo skipping
def get_label_name(labels):
    l = np.where(labels == 1.)[0]
    label_names = []
    for label in l:
        label_names.append(LABELS[label])
        
    return '-'.join(str(label_name) for label_name in label_names)

def show_batch(image_batch, label_batch):
  plt.figure(figsize=(20,20))
  for n in range(10):
      ax = plt.subplot(5,5,n+1)
      plt.imshow(image_batch[n])
      plt.title(get_label_name(label_batch[n].numpy()))
      plt.axis('off')
#%%script echo skipping
# Training batch
image_batch, label_batch = next(iter(train_ds))
show_batch(image_batch, label_batch)
#print(label_batch)
#%%script echo skipping
def get_model():
    base_model = tf.keras.applications.EfficientNetB3(include_top=False, weights='imagenet')
    base_model.trainable = True

    inputs = Input((IMG_HEIGHT, IMG_WIDTH, 3))
    x = base_model(inputs, training=True)
    x = GlobalAveragePooling2D()(x)
    x = Dropout(0.3)(x)
    outputs = Dense(len(LABELS), activation='sigmoid')(x)
    
    return Model(inputs, outputs)

tf.keras.backend.clear_session()
model = get_model()
model.summary()
#%%script echo skipping
time_stopping_callback = tfa.callbacks.TimeStopping(seconds=int(round(60*60*8)), verbose=1) #8h to not exceed allowance

earlystopper = tf.keras.callbacks.EarlyStopping(
    monitor='val_loss', patience=10, verbose=0, mode='min',
    restore_best_weights=True
)

lronplateau = tf.keras.callbacks.ReduceLROnPlateau(
    monitor='val_loss', factor=0.5, patience=5, verbose=0,
    mode='auto', min_delta=0.0001, cooldown=0, min_lr=0
)
#%%script echo skipping
#set up checkpoint save
#source:https://www.tensorflow.org/tutorials/keras/save_and_load
!pip install -q pyyaml h5py
import os
checkpoint_path_input = "../input/traineffnetb3versionapart1/cp.ckpt"
checkpoint_path = "./cp.ckpt"
checkpoint_dir = os.path.dirname(checkpoint_path)

# Create a callback that saves the model's weights
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1)
#%%script echo skipping
import keras.backend as K
K_epsilon = K.epsilon()
def f1(y_true, y_pred):
    #y_pred = K.round(y_pred)
    y_pred = K.cast(K.greater(K.clip(y_pred, 0, 1), 0.5), K.floatx())
    tp = K.sum(K.cast(y_true*y_pred, 'float'), axis=0)
    tn = K.sum(K.cast((1-y_true)*(1-y_pred), 'float'), axis=0)
    fp = K.sum(K.cast((1-y_true)*y_pred, 'float'), axis=0)
    fn = K.sum(K.cast(y_true*(1-y_pred), 'float'), axis=0)

    p = tp / (tp + fp + K_epsilon)
    r = tp / (tp + fn + K_epsilon)

    f1 = 2*p*r / (p+r+K_epsilon)
    f1 = tf.where(tf.math.is_nan(f1), tf.zeros_like(f1), f1)
    return K.mean(f1)
def f1_loss(y_true, y_pred):
    
    #y_pred = K.cast(K.greater(K.clip(y_pred, 0, 1), THRESHOLD), K.floatx())
    tp = K.sum(K.cast(y_true*y_pred, 'float'), axis=0)
    tn = K.sum(K.cast((1-y_true)*(1-y_pred), 'float'), axis=0)
    fp = K.sum(K.cast((1-y_true)*y_pred, 'float'), axis=0)
    fn = K.sum(K.cast(y_true*(1-y_pred), 'float'), axis=0)

    p = tp / (tp + fp + K_epsilon)
    r = tp / (tp + fn + K_epsilon)

    f1 = 2*p*r / (p+r+K_epsilon)
    f1 = tf.where(tf.math.is_nan(f1), tf.zeros_like(f1), f1)
    return 1-K.mean(f1)
#%%script echo skipping
import tensorflow as tf
import timeit

device_name = tf.test.gpu_device_name()
if "GPU" not in device_name:
    print("GPU device not found")
print('Found GPU at: {}'.format(device_name))
#https://stackoverflow.com/questions/47490834/how-can-i-print-the-learning-rate-at-each-epoch-with-adam-optimizer-in-keras
import keras
def get_lr_metric(optimizer):
    def lr(y_true, y_pred):
        return optimizer._decayed_lr(tf.float32)
    return lr
optimizer = keras.optimizers.Adam()
lr_metric = get_lr_metric(optimizer)
#%%script echo skipping
# Initialize model
tf.keras.backend.clear_session()

#model = get_model()
model.load_weights(checkpoint_path_input)

# Compile model
model.compile(
    optimizer=optimizer, 
    loss=f1_loss, 
    metrics=[f1, lr_metric])

# Train
history=model.fit(
    train_ds,
    epochs=1000,
    validation_data=val_ds,
    class_weight=class_weights,
    callbacks=[cp_callback,earlystopper,time_stopping_callback])
#%%script echo skipping
history.history
#%%script echo skipping
#source: https://machinelearningmastery.com/display-deep-learning-model-training-history-in-keras/
# list all data in history
print(history.history.keys())
# summarize history for f1
plt.plot(history.history['f1'])
plt.plot(history.history['val_f1'])
plt.title('f1')
plt.ylabel('f1')
plt.xlabel('epoch')
plt.legend(['train', 'val'], loc='upper left')
plt.show()
# summarize history for loss
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'val'], loc='upper left')
plt.show()
# summarize history for lr
plt.plot(history.history['lr'])
plt.title('learning rate')
plt.ylabel('learning rate')
plt.xlabel('epoch')
plt.legend('train', loc='upper left')
plt.show()
with open('./historyhistory.pkl', 'wb') as handle:
    pickle.dump(history.history, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open('./history.pkl', 'wb') as handle:
    pickle.dump(history, handle, protocol=pickle.HIGHEST_PROTOCOL)