# Python ≥3.5 is required
import sys
assert sys.version_info >= (3, 5)

# Scikit-Learn ≥0.20 is required
import sklearn
assert sklearn.__version__ >= "0.20"

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

# Common imports
import numpy as np
import os

# Chapter imports
from functools import partial

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

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

# Where to save the files
PROJECT_ROOT_DIR = "D:\\AI\\handson-ml2-master\\"
CHAPTER_ID = "deep"
IMAGES_PATH = os.path.join(PROJECT_ROOT_DIR, "images", CHAPTER_ID)
os.makedirs(IMAGES_PATH, exist_ok=True)
LOGS_PATH = os.path.join(PROJECT_ROOT_DIR, "tensorboard_logs", CHAPTER_ID)
os.makedirs(LOGS_PATH, exist_ok=True)
H5_PATH = os.path.join(PROJECT_ROOT_DIR, 'h5', CHAPTER_ID)
os.makedirs(H5_PATH, exist_ok=True)


# backend
K = keras.backend


def get_run_logdir():
    import time
    run_id = time.strftime("run_%Y_%m_%d-%H_%M_%S")
    return os.path.join(LOGS_PATH, run_id)


def save_fig(fig_id, tight_layout=True, fig_extension="png", resolution=300):
    path = os.path.join(IMAGES_PATH, fig_id + "." + fig_extension)
    print("Saving figure", fig_id)
    if tight_layout:
        plt.tight_layout()
    plt.savefig(path, format=fig_extension, dpi=resolution)


def save_model(model, model_id, h5_extension="h5"):
    path = os.path.join(H5_PATH, model_id + "." + h5_extension)
    print("Saving model", model_id)
    model.save(path)


def split_dataset(X, y):
    y_5_or_6 = (y == 5) | (y == 6) # sandals or shirts
    y_A = y[~y_5_or_6]  # NOT 5 or 6
    y_A[y_A > 6] -= 2 # class indices 7, 8, 9 should be moved to 5, 6, 7
    y_B = (y[y_5_or_6] == 6).astype(np.float32) # If shirt (==6), y_B[i]=1, or =0
    return ((X[~y_5_or_6], y_A),
            (X[y_5_or_6], y_B))


# power scheduling
def power_decay(_lr0, _decay, _n_steps_per_epoch):
    def _power_decay_fn(epoch):
        return _lr0 / (1 + epoch * _n_steps_per_epoch * _decay)
    return _power_decay_fn


# exponential scheduling
def exponential_decay(lr0, s):  # s: the number of batches in one epoch, maybe 20
    def exponential_decay_fn(epoch):
        return lr0 * 0.1**(epoch / s)
    return exponential_decay_fn


# update lr at each iteration (batch) NOT epoch, so need to re-write the callback class
class ExponentialDecay(keras.callbacks.Callback):
    def __init__(self, s=40000):
        super().__init__()
        self.s = s  # s here is not the number of batches in one epoch but the total number of batches in training

    def on_batch_begin(self, batch, logs=None):
        # Note: the `batch` argument is reset at each epoch
        lr = K.get_value(self.model.optimizer.lr)
        K.set_value(self.model.optimizer.lr, lr * 0.1**(1 / self.s))  # change s into self.s

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        logs['lr'] = K.get_value(self.model.optimizer.lr)


def piecewise_constant_fn(epoch):
    if epoch < 5:
        return 0.01
    elif epoch < 15:
        return 0.005
    else:
        return 0.001


def piecewise_constant(boundaries, values):
    boundaries = np.array([0] + boundaries)
    values = np.array(values)

    def piecewise_constant_fn(epoch):
        return values[np.argmax(boundaries > epoch) - 1]
    return piecewise_constant_fn


class ExponentialLearningRate(keras.callbacks.Callback):
    def __init__(self, factor):
        self.factor = factor
        self.rates = []
        self.losses = []

    def on_batch_end(self, batch, logs):
        self.rates.append(K.get_value(self.model.optimizer.lr))
        self.losses.append(logs["loss"])
        K.set_value(self.model.optimizer.lr, self.model.optimizer.lr * self.factor)


def find_learning_rate(model, X, y, epochs=1, batch_size=32, min_rate=10**-5, max_rate=10):
    init_weights = model.get_weights()
    iterations = len(X) // batch_size * epochs
    factor = np.exp(np.log(max_rate / min_rate) / iterations)
    init_lr = K.get_value(model.optimizer.lr)
    K.set_value(model.optimizer.lr, min_rate)
    exp_lr = ExponentialLearningRate(factor)
    history = model.fit(X, y, epochs=epochs, batch_size=batch_size,
                        callbacks=[exp_lr])
    K.set_value(model.optimizer.lr, init_lr)
    model.set_weights(init_weights)
    return exp_lr.rates, exp_lr.losses


def plot_lr_vs_loss(rates, losses):
    plt.plot(rates, losses)
    plt.gca().set_xscale('log')
    plt.hlines(min(losses), min(rates), max(rates))
    plt.axis([min(rates), max(rates), min(losses), (losses[0] + min(losses)) / 2])
    plt.xlabel("Learning rate")
    plt.ylabel("Loss")


class OneCycleScheduler(keras.callbacks.Callback):
    def __init__(self, iterations, max_rate, start_rate=None,
                 last_iterations=None, last_rate=None):
        self.iterations = iterations
        self.max_rate = max_rate
        self.start_rate = start_rate or max_rate / 10
        self.last_iterations = last_iterations or iterations // 10 + 1
        self.half_iteration = (iterations - self.last_iterations) // 2
        self.last_rate = last_rate or self.start_rate / 1000
        self.iteration = 0

    def _interpolate(self, iter1, iter2, rate1, rate2):
        return ((rate2 - rate1) * (self.iteration - iter1)
                / (iter2 - iter1) + rate1)

    def on_batch_begin(self, batch, logs):
        if self.iteration < self.half_iteration:
            rate = self._interpolate(0, self.half_iteration, self.start_rate, self.max_rate)
        elif self.iteration < 2 * self.half_iteration:
            rate = self._interpolate(self.half_iteration, 2 * self.half_iteration,
                                     self.max_rate, self.start_rate)
        else:
            rate = self._interpolate(2 * self.half_iteration, self.iterations,
                                     self.start_rate, self.last_rate)
            rate = max(rate, self.last_rate)
        self.iteration += 1
        K.set_value(self.model.optimizer.lr, rate)


class MCDropout(keras.layers.Dropout):
    def call(self, inputs):
        return super().call(inputs, training=True)


class MCAlphaDropout(keras.layers.AlphaDropout):
    def call(self, inputs):
        return super().call(inputs, training=True)


if __name__ == '__main__':

    # data set
    (X_train_full, y_train_full), (X_test, y_test) = keras.datasets.fashion_mnist.load_data()
    X_train_full = X_train_full / 255.0
    X_test = X_test / 255.0
    X_valid, X_train = X_train_full[:5000], X_train_full[5000:]
    y_valid, y_train = y_train_full[:5000], y_train_full[5000:]

    pixel_means = X_train.mean(axis=0, keepdims=True)
    pixel_stds = X_train.std(axis=0, keepdims=True)
    X_train_scaled = (X_train - pixel_means) / pixel_stds
    X_valid_scaled = (X_valid - pixel_means) / pixel_stds
    X_test_scaled = (X_test - pixel_means) / pixel_stds

    # L1 and L2 regularization

    # build a model with L1, L2 regularization
    model = keras.models.Sequential([
        keras.layers.Flatten(input_shape=[28, 28]),
        keras.layers.Dense(300, activation="elu",
                           kernel_initializer="he_normal",
                           kernel_regularizer=keras.regularizers.l1(0.01)),
        keras.layers.Dense(100, activation="elu",
                           kernel_initializer="he_normal",
                           kernel_regularizer=keras.regularizers.l2(0.01)),
        keras.layers.Dense(10, activation="softmax",
                           kernel_regularizer=keras.regularizers.l1_l2(0.1, 0.01))
    ])  # L1, L2 regularization added into this model is just for showing how to use, so maybe not suitable in the model

    model.compile(loss="sparse_categorical_crossentropy", optimizer="nadam", metrics=["accuracy"])
    n_epochs = 2
    history = model.fit(X_train_scaled, y_train, epochs=n_epochs,
                        validation_data=(X_valid_scaled, y_valid))

    # functools.partial() to make a general layer for using later.
    RegularizedDense = partial(keras.layers.Dense,
                               activation="elu",
                               kernel_initializer="he_normal",
                               kernel_regularizer=keras.regularizers.l2(0.01))

    model = keras.models.Sequential([
        keras.layers.Flatten(input_shape=[28, 28]),
        RegularizedDense(300),  # other params are defined in partial()
        RegularizedDense(100),
        RegularizedDense(10, activation="softmax")  # also can rewrite params which are already defined in partial()
    ])

    model.compile(loss="sparse_categorical_crossentropy", optimizer="nadam", metrics=["accuracy"])

    n_epochs = 2
    l1_l2_partial_history = model.fit(X_train_scaled, y_train, epochs=n_epochs,
                                      validation_data=(X_valid_scaled, y_valid))

    # Dropout

    model = keras.models.Sequential([
        keras.layers.Flatten(input_shape=[28, 28]),
        keras.layers.Dropout(rate=0.2),  # dropout rate = 20%
        keras.layers.Dense(300, activation="elu", kernel_initializer="he_normal"),
        keras.layers.Dropout(rate=0.2),
        keras.layers.Dense(100, activation="elu", kernel_initializer="he_normal"),
        keras.layers.Dropout(rate=0.2),
        keras.layers.Dense(10, activation="softmax")
    ])
    model.compile(loss="sparse_categorical_crossentropy", optimizer="nadam", metrics=["accuracy"])
    n_epochs = 2
    dropout_history = model.fit(X_train_scaled, y_train, epochs=n_epochs,
                                validation_data=(X_valid_scaled, y_valid))

    # Alpha Dropout

    tf.random.set_random_seed(42)
    np.random.seed(42)

    # Alpha Dropout needs to work with SELU activation
    model = keras.models.Sequential([
        keras.layers.Flatten(input_shape=[28, 28]),
        keras.layers.AlphaDropout(rate=0.2),
        keras.layers.Dense(300, activation="selu", kernel_initializer="lecun_normal"),  # SELU + Alpha Dropout
        keras.layers.AlphaDropout(rate=0.2),                                            # Alpha
        keras.layers.Dense(100, activation="selu", kernel_initializer="lecun_normal"),
        keras.layers.AlphaDropout(rate=0.2),
        keras.layers.Dense(10, activation="softmax")
    ])
    optimizer = keras.optimizers.SGD(lr=0.01, momentum=0.9, nesterov=True)
    model.compile(loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])

    n_epochs = 20
    alpha_dropout_history = model.fit(X_train_scaled, y_train, epochs=n_epochs,
                                      validation_data=(X_valid_scaled, y_valid))

    model.evaluate(X_test_scaled, y_test)  # [0.45350628316402436, 0.868] more loss and fewer accuracy
    model.evaluate(X_train_scaled, y_train)  # [0.335701530437036, 0.88872725]

    # MC Dropout

    tf.random.set_random_seed(42)
    np.random.seed(42)

    # code from book
    mc_model = keras.models.Sequential(
        [MCAlphaDropout(layer.rate) if isinstance(layer, keras.layers.AlphaDropout)
         else layer for layer in model.layers])  # replace regular Dropout with MC Dropout layer
    print(mc_model.summary())

    optimizer = keras.optimizers.SGD(lr=0.01, momentum=0.9, nesterov=True)
    mc_model.compile(loss="sparse_categorical_crossentropy", optimizer=optimizer, metrics=["accuracy"])

    mc_model.set_weights(model.get_weights())  # no need to fit again
    # MC Dropout model needs the average prediction.
    np.round(np.mean([mc_model.predict(X_test_scaled[:1]) for sample in range(100)], axis=0), 2)

    # how to use MC Dropout
    mnist = tf.keras.datasets.mnist

    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train, x_test = x_train / 255.0, x_test / 255.0

    inp = keras.layers.Input(shape=(28, 28))
    x = keras.layers.Flatten()(inp)
    x = keras.layers.Dense(512, activation=tf.nn.relu)(x)
    x = keras.layers.Dropout(0.5)(x, training=True)  # dropout is active in training and prediction
    out = keras.layers.Dense(10, activation="softmax")(x)
    model = keras.Model(inp, out)

    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    model.fit(x_train, y_train, epochs=3)
    # dropout is active while the whole prediction. so results are different.
    for _ in range(10):
        print(model.predict(x_test[:1]))  # need the average prediction probabilites

    # Max-Norm Regularization    Max norm

    # how to use max norm in Keras
    layer = keras.layers.Dense(100, activation="selu", kernel_initializer="lecun_normal",
                               kernel_constraint=keras.constraints.max_norm(1.))
    # max_norm(1., axis=[0,1,2...]) default axis=0 in Dense. And axis should be changed in convolutional layers

    # partial() with max norm
    MaxNormDense = partial(keras.layers.Dense,
                           activation="selu", kernel_initializer="lecun_normal",
                           kernel_constraint=keras.constraints.max_norm(1.))  # partial(layer type, params args...)
    # also the para "bias_constraint=" can be added into this partial to constrain the bias in layers

    model = keras.models.Sequential([
        keras.layers.Flatten(input_shape=[28, 28]),
        MaxNormDense(300),  # max norm does can also reduce gradient explode/vanish problem as BN can.
        MaxNormDense(100),
        keras.layers.Dense(10, activation="softmax")
    ])
    model.compile(loss="sparse_categorical_crossentropy", optimizer="nadam", metrics=["accuracy"])

    n_epochs = 2
    history = model.fit(X_train_scaled, y_train, epochs=n_epochs,
                        validation_data=(X_valid_scaled, y_valid))
