"""
### This Notebook has been taken from Keras.io. It shows how to detect pneumonia from chest x-ray using a TPU

#### see https://keras.io/examples/vision/xray_classification_with_tpus/
#### Done some changes: changed using Efficient Net for the convolutional base
"""
!pip install -q efficientnet >> /dev/null
import re
import os
import random
import numpy as np
import pandas as pd
import tensorflow as tf
import matplotlib.pyplot as plt

from kaggle_datasets import KaggleDatasets

try:
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
    print("Device:", tpu.master())
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    strategy = tf.distribute.experimental.TPUStrategy(tpu)
except:
    strategy = tf.distribute.get_strategy()

print("Number of replicas:", strategy.num_replicas_in_sync)
# for now using TF 2.2 (compatibility with TPU)
print(tf.__version__)
AUTOTUNE = tf.data.experimental.AUTOTUNE

GCS_PATH = KaggleDatasets().get_gcs_path('chest-xray-pneumonia')

BATCH_SIZE = 25 * strategy.num_replicas_in_sync

IMAGE_SIZE = [256, 256]
# ho aumentatop il train set
PERC_TRAIN = 0.9

filenames = tf.io.gfile.glob(str(GCS_PATH + "/chest_xray/train/*/*"))
filenames.extend(tf.io.gfile.glob(str(GCS_PATH + "/chest_xray/val/*/*")))

random.shuffle(filenames)

split_ind = int(PERC_TRAIN* len(filenames))

train_filenames, val_filenames = filenames[:split_ind], filenames[split_ind:]
COUNT_NORMAL = len([filename for filename in train_filenames if "NORMAL" in filename])
print("Normal images count in training set: " + str(COUNT_NORMAL))

COUNT_PNEUMONIA = len(
    [filename for filename in train_filenames if "PNEUMONIA" in filename]
)
print("Pneumonia images count in training set: " + str(COUNT_PNEUMONIA))
train_list_ds = tf.data.Dataset.from_tensor_slices(train_filenames)

val_list_ds = tf.data.Dataset.from_tensor_slices(val_filenames)

for f in train_list_ds.take(5):
    print(f.numpy())
TRAIN_IMG_COUNT = tf.data.experimental.cardinality(train_list_ds).numpy()
print("Training images count: " + str(TRAIN_IMG_COUNT))

VAL_IMG_COUNT = tf.data.experimental.cardinality(val_list_ds).numpy()
print("Validating images count: " + str(VAL_IMG_COUNT))
CLASS_NAMES = [
    str(tf.strings.split(item, os.path.sep)[-1].numpy())[2:-1]
    for item in tf.io.gfile.glob(str(GCS_PATH + "/chest_xray/train/*"))
]
print("Class names: %s" % (CLASS_NAMES,))
def get_label(file_path):
    # convert the path to a list of path components
    parts = tf.strings.split(file_path, os.path.sep)
    # The second to last is the class-directory
    return parts[-2] == "PNEUMONIA"


def decode_img(img):
    # convert the compressed string to a 3D uint8 tensor
    img = tf.image.decode_jpeg(img, channels=3)
    # resize the image to the desired size.
    return tf.image.resize(img, IMAGE_SIZE)


def process_path(file_path):
    label = get_label(file_path)
    # load the raw data from the file as a string
    img = tf.io.read_file(file_path)
    img = decode_img(img)
    return img, label


train_ds = train_list_ds.map(process_path, num_parallel_calls=AUTOTUNE)

val_ds = val_list_ds.map(process_path, num_parallel_calls=AUTOTUNE)
for image, label in train_ds.take(1):
    print("Image shape: ", image.numpy().shape)
    print("Label: ", label.numpy())
# ho fatto alcune modifiche rispetto all'esempio su Keras.io per correggere un bug.

test_filenames = tf.io.gfile.glob(str(GCS_PATH + "/chest_xray/test/*/*"))

test_list_ds = tf.data.Dataset.from_tensor_slices(test_filenames)

TEST_IMG_COUNT = tf.data.experimental.cardinality(test_list_ds).numpy()

print("Test images count: " + str(TEST_IMG_COUNT))
test_ds = test_list_ds.map(process_path, num_parallel_calls=AUTOTUNE)
test_ds = test_ds.batch(BATCH_SIZE)
def prepare_for_training(ds, cache=True, shuffle_buffer_size=1000):
    # This is a small dataset, only load it once, and keep it in memory.
    # use `.cache(filename)` to cache preprocessing work for datasets that don't
    # fit in memory.
    if cache:
        if isinstance(cache, str):
            ds = ds.cache(cache)
        else:
            ds = ds.cache()

    ds = ds.shuffle(buffer_size=shuffle_buffer_size)
    ds = ds.batch(BATCH_SIZE)

    # `prefetch` lets the dataset fetch batches in the background while the model
    # is training.
    ds = ds.prefetch(buffer_size=AUTOTUNE)

    return ds
train_ds = prepare_for_training(train_ds)
val_ds = prepare_for_training(val_ds)

image_batch, label_batch = next(iter(train_ds))
def show_batch(image_batch, label_batch):
    plt.figure(figsize=(10, 10))
    for n in range(25):
        ax = plt.subplot(5, 5, n + 1)
        plt.imshow(image_batch[n] / 255)
        if label_batch[n]:
            plt.title("PNEUMONIA")
        else:
            plt.title("NORMAL")
        plt.axis("off")
show_batch(image_batch.numpy(), label_batch.numpy())
"""
### build the CNN
"""
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.layers.experimental import preprocessing

import efficientnet.tfkeras as efn
EFNS = [efn.EfficientNetB0, efn.EfficientNetB1, efn.EfficientNetB2, efn.EfficientNetB3, 
        efn.EfficientNetB4, efn.EfficientNetB5, efn.EfficientNetB6, efn.EfficientNetB7]

# as default it used B0

def build_model(dim = 256, ef = 0):
    inp = tf.keras.layers.Input(shape=(*IMAGE_SIZE, 3))

    x = preprocessing.Rescaling(1.0 / 255)(inp)
    
    base = EFNS[ef](input_shape=(*IMAGE_SIZE, 3), weights='imagenet', include_top = False)
    
    x = base(x)
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    x = tf.keras.layers.Dense(128, activation='relu')(x)
    x = tf.keras.layers.Dropout(0.3)(x)
    x = tf.keras.layers.Dense(1, activation='sigmoid')(x)
    
    model = tf.keras.Model(inputs = inp,outputs = x)
    
    return model
# see how many params the model has

model_s = build_model(ef = 4)

model_s.summary()
checkpoint_cb = tf.keras.callbacks.ModelCheckpoint("xray_model.h5", save_best_only=True)

early_stopping_cb = tf.keras.callbacks.EarlyStopping(monitor="val_loss",
    patience=10, restore_best_weights=True, verbose=1)
def get_lr_callback(batch_size=8):
    lr_start   = 0.000001
    lr_max     = 0.000005 * 2 * batch_size/16 # 0.000020 * 8 * batch_size/16
    lr_min     = 0.000001
    lr_ramp_ep = 8
    lr_sus_ep  = 4
    lr_decay   = 0.8
   
    def lrfn(epoch):
        if epoch < lr_ramp_ep:
            lr = (lr_max - lr_start) / lr_ramp_ep * epoch + lr_start
            
        elif epoch < lr_ramp_ep + lr_sus_ep:
            lr = lr_max
            
        else:
            lr = (lr_max - lr_min) * lr_decay**(epoch - lr_ramp_ep - lr_sus_ep) + lr_min
            
        return lr

    lr_callback = tf.keras.callbacks.LearningRateScheduler(lrfn, verbose=True)
    
    return lr_callback
with strategy.scope():
    # proviamo con EF4
    model = build_model(ef = 4)

    METRICS = [
        tf.keras.metrics.BinaryAccuracy(),
        tf.keras.metrics.Precision(name="precision"),
        tf.keras.metrics.Recall(name="recall"),
        tf.keras.metrics.AUC(name='auc')
    ]
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate = 0.001),
        loss="binary_crossentropy",
        metrics=METRICS,
    )

history = model.fit(
    train_ds,
    epochs=60,
    validation_data=val_ds,
    # class_weight=class_weight,
    callbacks=[checkpoint_cb, early_stopping_cb, get_lr_callback(BATCH_SIZE)],
)
fig, ax = plt.subplots(1, 4, figsize=(20, 5))
ax = ax.ravel()

for i, met in enumerate(["auc", "recall", "binary_accuracy", "loss"]):
    ax[i].plot(history.history[met])
    ax[i].plot(history.history["val_" + met])
    ax[i].set_title("Model {}".format(met))
    ax[i].set_xlabel("epochs")
    ax[i].set_ylabel(met)
    ax[i].legend(["train", "val"])
model.evaluate(test_ds, return_dict=True)
img = tf.io.read_file(
    str(GCS_PATH + "/chest_xray/test/PNEUMONIA/person100_bacteria_475.jpeg")
)

img = decode_img(img)/255.
plt.imshow(img)

img_array = tf.keras.preprocessing.image.img_to_array(img)
img_array = tf.expand_dims(img_array, 0)  # Create batch axis

prediction = model.predict(img_array)[0]
scores = [1 - prediction, prediction]

for score, name in zip(scores, CLASS_NAMES):
    print("This image is %.2f percent %s" % ((100 * score), name))