# Get helper functions file
import os 
import matplotlib.pyplot as plt

if not os.path.exists("helper_functions.py"):
    !wget https://raw.githubusercontent.com/mrdbourke/tensorflow-deep-learning/main/extras/helper_functions.py
else:
    print("[INFO] 'helper_functions.py' already exists, skipping download.")

dir_path = os.path.dirname(os.path.realpath(os.getcwd()))
print(f'current dirPath: {dir_path}')

import tensorflow as tf
print(f"TensorFlow version: {tf.__version__}")

# Add timestamp
import datetime
print(f"Notebook last run (end-to-end): {datetime.datetime.now()}")

# 
# data-scient + ML
# 
from helper_functions import create_tensorboard_callback, plot_loss_curves, compare_historys
import tensorflow_datasets as tfds

# 
# detect if the dataset is, indeed, in the tensorflow_datasets world
# 

# Get all available datasets in TFDS
datasets_list = tfds.list_builders()

target_dataset = "food101"
print(f"'{target_dataset}' in TensorFlow Datasets: {target_dataset in datasets_list}")

# 
# get the data
# *** *** *** WARNING
#             downloads 5g of images
# *** *** *** WARNING

# Load in the data (takes about 5-6 minutes in Google Colab)
(train_data, test_data), ds_info = tfds.load(name=target_dataset, # target dataset to get from TFDS
                                             split=["train", "validation"], # what splits of data should we get? note: not all datasets have train, valid, test
                                             shuffle_files=True, # shuffle files on download?
                                             as_supervised=True, # download data in tuple format (sample, label), e.g. (image, label)
                                             with_info=True) # include dataset metadata? if so, tfds.load() returns tuple (data, ds_info)

ds_info.features

# Get class names
class_names = ds_info.features["label"].names
print(f'{len(class_names)} classes')
print(class_names[:10])

# 
# get a single sample of the data
# 
train_one_sample = train_data.take(1) # samples are in format (image_tensor, label)
train_one_sample

# view the data in the sample
# 
for image, label in train_one_sample:
  print(f"""
  Image shape: {image.shape}
  Image dtype: {image.dtype}
  Label (tensor form): {label}
  label-to-Class(str form): {class_names[label.numpy()]}
        """)

image

plt.imshow(image)
plt.title(class_names[label.numpy()]) # add title to image by indexing on class_names list
plt.axis(False);

def resizeAndType(image, label, img_shape=224):
    """
    Converts image datatype from 'uint8' -> 'float32' and reshapes image to
    [img_shape, img_shape, color_channels]
    """
    image = tf.image.resize(image, [img_shape, img_shape]) # reshape to img_shape
    return tf.cast(image, tf.float32), label # return (float32_image, label) tuple

# get a singlee sample image and check the outputs
preprocessed_img = resizeAndType(image, label)[0]
print(f"Image before preprocessing:\n {image[:2]}...,\nShape: {image.shape},\nDatatype: {image.dtype}\n")
print(f"Image after preprocessing:\n {preprocessed_img[:2]}...,\nShape: {preprocessed_img.shape},\nDatatype: {preprocessed_img.dtype}")

# We can still plot our preprocessed image as long as we 
# divide by 255 (for matplotlib capatibility)
plt.imshow(preprocessed_img/255.)
plt.title(class_names[label])
plt.axis(False);

# 
# TRAINING DATA
#   resizeAndType
#   shuffle, batch, prefetch
# 
resizedTrainedData = train_data.map(map_func=resizeAndType, num_parallel_calls=tf.data.AUTOTUNE)
preparedTrainingData = resizedTrainedData.shuffle(buffer_size=1000).batch(batch_size=32).prefetch(buffer_size=tf.data.AUTOTUNE)

# 
# TESTING DATA
#   resizeAndType
#   batch, prefetch (NO SHUFFLE)
# 
resizedTestData = test_data.map(resizeAndType, num_parallel_calls=tf.data.AUTOTUNE)
preparedTestingData = resizedTestData.batch(32).prefetch(tf.data.AUTOTUNE)

from helper_functions import create_tensorboard_callback

# Create ModelCheckpoint callback to save model's progress
checkpoint_path = "model_checkpoints/cp.ckpt" # saving weights requires ".ckpt" extension
model_checkpoint = tf.keras.callbacks.ModelCheckpoint(checkpoint_path,
                                                      monitor="val_accuracy", # save the model weights with best validation accuracy
                                                      save_best_only=True, # only save the best weights
                                                      save_weights_only=True, # only save model weights (not whole model)
                                                      verbose=0) # don't print out whether or not model is being saved 

# Turn on mixed precision training
# from tensorflow.keras import mixed_precision
# mixed_precision.set_global_policy(policy="mixed_float16") # set global policy to mixed precision