import pickle
from time import time
from datetime import datetime
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_datasets as tfds
from tensorflow.keras.layers import Dense, Flatten, Conv2D, MaxPool2D, Dropout
from tensorflow.keras.layers import BatchNormalization, Input, GlobalAveragePooling2D
from tensorflow.keras.callbacks import ReduceLROnPlateau
from tensorflow.keras import Model
import IPython
import optuna
from optuna.integration.tensorboard import TensorBoardCallback

def version_info(cls):
    print(f"{cls.__name__}: {cls.__version__}")

print("Version Used in this Notebook:")
version_info(tf)
version_info(tfds)
version_info(optuna)

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))

EPOCHS = 25
# EPOCHS = 2
BATCH_SIZE=128
IMG_SIZE=32
NUM_CLASSES=10

def get_dataset():
    (ds_train, ds_test), metadata = tfds.load(
        'cifar10', split=['train', 'test'], shuffle_files=True, 
        with_info=True, as_supervised=True
    )
    
    train_ds = ds_train \
        .cache() \
        .batch(BATCH_SIZE, drop_remainder=True) \
        .prefetch(tf.data.experimental.AUTOTUNE)
    
    test_ds = ds_test \
        .cache() \
        .batch(BATCH_SIZE, drop_remainder=True) \
        .prefetch(tf.data.experimental.AUTOTUNE)
    return (train_ds, test_ds)

transforms = tf.keras.Sequential([
    tf.keras.layers.experimental.preprocessing.RandomFlip('horizontal'),
    tf.keras.layers.experimental.preprocessing.RandomRotation(0.2),
])

## Create Model
def create_model(trial):
    """
    Create a simple CIFAR-10 model that uses ResNet50 as its backbone.
    
    Params:
    -------
    trial: optuna Trial object
    """
    inputs = Input(shape=(IMG_SIZE, IMG_SIZE, 3))
    x = transforms(inputs)
    x = tf.keras.applications.resnet.preprocess_input(x)
    x = tf.keras.applications.ResNet50(input_shape=(IMG_SIZE, IMG_SIZE, 3), include_top=False)(x, training=False)
    # Flatten or GAP
    use_gap = trial.suggest_categorical('use_gap', [True, False])
    if use_gap:
        x = GlobalAveragePooling2D()(x)
    else:
        x = Flatten()(x)
    x = Flatten()(x)
    # Dense Layer Units
    num_hidden = trial.suggest_int('dense_1', 32, 128)
    # Activation
    activation = trial.suggest_categorical('activation', ['relu', 'selu', 'elu'])
    x = Dense(128, activation=activation)(x)
    # Dropout rate    
    drop_rate = trial.suggest_float('drop_rate', 0.0, 0.8)
    x = Dropout(drop_rate)(x)
    outputs = Dense(NUM_CLASSES)(x)
    model = tf.keras.Model(inputs, outputs)
    return model

## Create Optimizer
def create_optimizer(trial):
    # LR
    lr = trial.suggest_float('learning_rate', 1e-5, 1e-1, log=True)
    optimizer = tf.keras.optimizers.Adam(learning_rate=lr)
    return optimizer

## Define Objective
def train(model, optimizer, dataset, mode="eval"):
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)    
    accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name="accuracy")
    mean_loss = tf.keras.metrics.Mean(name="loss")
    for images, labels in dataset:
        with tf.GradientTape() as tape:
            predictions = model(images, training=(mode=='train'))
            loss = loss_object(labels, predictions)            
            if mode == "train":
                gradients = tape.gradient(loss, model.trainable_variables)
                optimizer.apply_gradients(zip(gradients, model.trainable_variables))
            accuracy(labels, predictions)
            mean_loss(loss)
    return accuracy.result(), mean_loss.result()

def objective(trial):
    # dataset
    train_ds, test_ds = get_dataset()
    # model
    model = create_model(trial)
    # optimizer
    optimizer = create_optimizer(trial)
    # train
    for epoch in range(EPOCHS):
        train_acc, train_loss = train(model, optimizer, train_ds, "train")
        test_acc, test_loss = train(model, optimizer, test_ds, "eval")
        trial.report(test_acc, epoch)
        print(f"train_accuracy:{train_acc:.4f}, train_loss: {train_loss:.4f}, test_acc: {test_acc:.4f}")
        if trial.should_prune():
            raise optuna.exceptions.TrialPruned()
        
    return test_acc

# Track using Tensorboard
tensorboard_cb = TensorBoardCallback("./logs/", metric_name="accuracy")

start = time()
# Run
study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=1, timeout=600, callbacks=[tensorboard_cb])
stop = time()

took = stop - start
print(f"Total time: {took//60 : .0f}m {took%60:.0f}s")

pruned_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.PRUNED]
complete_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]

print(f"Finished Trials: {len(study.trials)}")
print(f"Pruned Trials: {len(pruned_trials)}")
print(f"Completed Trials: {len(complete_trials)}")

trial = study.best_trial
print(trial.value)

for k,v in trial.params.items():
    print(f"{k}: {v}")