from __future__ import absolute_import, division, print_function, unicode_literals

import tensorflow as tf
from tensorflow.keras import layers

print(tf.__version__)

# !pip install -q git+https://github.com/tensorflow/docs

# import tensorflow_docs as tfdocs
# import tensorflow_docs.modeling
# import tensorflow_docs.plots

from IPython import display
from matplotlib import pyplot as plt

import numpy as np

import pathlib
import shutil
import tempfile

logdir = pathlib.Path(tempfile.mkdtemp()) / "tensorboard_logs"
shutil.rmtree(logdir, ignore_errors=True)

gz = tf.keras.utils.get_file('HIGGS.csv.gz',
                             'https://archive.ics.uci.edu/ml/machine-learning-databases/00280/HIGGS.csv.gz')

FEATURES = 28

ds = tf.data.experimental.CsvDataset(gz, [float(), ] * (FEATURES + 1), compression_type="GZIP")


def pack_row(*row):
    label = row[0]
    features = tf.stack(row[1:], 1)
    return features, label


packed_ds = ds.batch(10000).map(pack_row).unbatch()

for features, label in packed_ds.batch(1000).take(1):
    print(features[0])
    plt.hist(features.numpy().flatten(), bins=101)

N_VALIDATION = int(1e3)
N_TRAIN = int(1e4)
BUFFER_SIZE = int(1e4)
BATCH_SIZE = 500
STEPS_PER_EPOCH = N_TRAIN // BATCH_SIZE

validate_ds = packed_ds.take(N_VALIDATION).cache()
train_ds = packed_ds.skip(N_VALIDATION).take(N_TRAIN).cache()

train_ds

validate_ds = validate_ds.batch(BATCH_SIZE)
train_ds = train_ds.shuffle(BUFFER_SIZE).repeat().batch(BATCH_SIZE)

lr_schedule = tf.keras.optimizers.schedules.InverseTimeDecay(
    0.001,
    decay_steps=STEPS_PER_EPOCH * 1000,
    decay_rate=1,
    staircase=False)


def get_optimizer():
    return tf.keras.optimizers.Adam(lr_schedule)


step = np.linspace(0, 100000)
lr = lr_schedule(step)
plt.figure(figsize=(8, 6))
plt.plot(step / STEPS_PER_EPOCH, lr)
plt.ylim([0, max(plt.ylim())])
plt.xlabel('Epoch')
_ = plt.ylabel('Learning Rate')


def get_callbacks(name):
    return [
        tfdocs.modeling.EpochDots(),
        tf.keras.callbacks.EarlyStopping(monitor='val_binary_crossentropy', patience=200),
        tf.keras.callbacks.TensorBoard(logdir / name),
    ]


def compile_and_fit(model, name, optimizer=None, max_epochs=10000):
    if optimizer is None:
        optimizer = get_optimizer()
    model.compile(optimizer=optimizer,
                  loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
                  metrics=[
                      tf.keras.losses.BinaryCrossentropy(
                          from_logits=True, name='binary_crossentropy'),
                      'accuracy'])

    model.summary()

    history = model.fit(
        train_ds,
        steps_per_epoch=STEPS_PER_EPOCH,
        epochs=max_epochs,
        validation_data=validate_ds,
        callbacks=get_callbacks(name),
        verbose=0)
    return history


tiny_model = tf.keras.Sequential([
    layers.Dense(16, activation='elu', input_shape=(FEATURES,)),
    layers.Dense(1)
])

size_histories = {}

size_histories['Tiny'] = compile_and_fit(tiny_model, 'sizes/Tiny')

plotter = tfdocs.plots.HistoryPlotter(metric='binary_crossentropy', smoothing_std=10)
plotter.plot(size_histories)
plt.ylim([0.5, 0.7])

small_model = tf.keras.Sequential([
    # `input_shape` is only required here so that `.summary` works.
    layers.Dense(16, activation='elu', input_shape=(FEATURES,)),
    layers.Dense(16, activation='elu'),
    layers.Dense(1)
])

size_histories['Small'] = compile_and_fit(small_model, 'sizes/Small')

medium_model = tf.keras.Sequential([
    layers.Dense(64, activation='elu', input_shape=(FEATURES,)),
    layers.Dense(64, activation='elu'),
    layers.Dense(64, activation='elu'),
    layers.Dense(1)
])

size_histories['Medium'] = compile_and_fit(medium_model, "sizes/Medium")

large_model = tf.keras.Sequential([
    layers.Dense(512, activation='elu', input_shape=(FEATURES,)),
    layers.Dense(512, activation='elu'),
    layers.Dense(512, activation='elu'),
    layers.Dense(512, activation='elu'),
    layers.Dense(1)
])

size_histories['large'] = compile_and_fit(large_model, "sizes/large")

plotter.plot(size_histories)
a = plt.xscale('log')
plt.xlim([5, max(plt.xlim())])
plt.ylim([0.5, 0.7])
plt.xlabel("Epochs [Log Scale]")

display.IFrame(
    src="https://tensorboard.dev/experiment/vW7jmmF9TmKmy3rbheMQpw/#scalars&_smoothingWeight=0.97",
    width="100%", height="800px")

shutil.rmtree(logdir / 'regularizers/Tiny', ignore_errors=True)
shutil.copytree(logdir / 'sizes/Tiny', logdir / 'regularizers/Tiny')

numpy.linspace matplotlib.pyplot.ylim matplotlib.pyplot.xscale tensorflow.stack tensorflow.keras.layers.Dense tensorflow.keras.losses.BinaryCrossentropy matplotlib.pyplot.plot tensorflow.keras.optimizers.schedules.InverseTimeDecay matplotlib.pyplot.ylabel tensorflow.keras.optimizers.Adam tensorflow.keras.utils.get_file matplotlib.pyplot.xlim tensorflow.keras.callbacks.TensorBoard matplotlib.pyplot.xlabel tensorflow.keras.callbacks.EarlyStopping matplotlib.pyplot.figure
[(24, 'shutil.rmtree', 'shutil.rmtree', (['logdir'], {'ignore_errors': '(True)'}), False, 'import shutil\n'), (26, 'tensorflow.keras.utils.get_file', 'tf.keras.utils.get_file', (['"""HIGGS.csv.gz"""', '"""https://archive.ics.uci.edu/ml/machine-learning-databases/00280/HIGGS.csv.gz"""'], {}), True, 'import tensorflow as tf\n'), (60, 'tensorflow.keras.optimizers.schedules.InverseTimeDecay', 'tf.keras.optimizers.schedules.InverseTimeDecay', (['(0.001)'], {'decay_steps': '(STEPS_PER_EPOCH * 1000)', 'decay_rate': '(1)', 'staircase': '(False)'}), True, 'import tensorflow as tf\n'), (71, 'numpy.linspace', 'np.linspace', (['(0)', '(100000)'], {}), True, 'import numpy as np\n'), (73, 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(8, 6)'}), True, 'from matplotlib import pyplot as plt\n'), (74, 'matplotlib.pyplot.plot', 'plt.plot', (['(step / STEPS_PER_EPOCH)', 'lr'], {}), True, 'from matplotlib import pyplot as plt\n'), (76, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Epoch"""'], {}), True, 'from matplotlib import pyplot as plt\n'), (77, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Learning Rate"""'], {}), True, 'from matplotlib import pyplot as plt\n'), (121, 'matplotlib.pyplot.ylim', 'plt.ylim', (['[0.5, 0.7]'], {}), True, 'from matplotlib import pyplot as plt\n'), (152, 'matplotlib.pyplot.xscale', 'plt.xscale', (['"""log"""'], {}), True, 'from matplotlib import pyplot as plt\n'), (154, 'matplotlib.pyplot.ylim', 'plt.ylim', (['[0.5, 0.7]'], {}), True, 'from matplotlib import pyplot as plt\n'), (155, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Epochs [Log Scale]"""'], {}), True, 'from matplotlib import pyplot as plt\n'), (157, 'IPython.display.IFrame', 'display.IFrame', ([], {'src': '"""https://tensorboard.dev/experiment/vW7jmmF9TmKmy3rbheMQpw/#scalars&_smoothingWeight=0.97"""', 'width': '"""100%"""', 'height': '"""800px"""'}), False, 'from IPython import display\n'), (161, 'shutil.rmtree', 'shutil.rmtree', (["(logdir / 'regularizers/Tiny')"], {'ignore_errors': '(True)'}), False, 'import shutil\n'), (162, 'shutil.copytree', 'shutil.copytree', (["(logdir / 'sizes/Tiny')", "(logdir / 'regularizers/Tiny')"], {}), False, 'import shutil\n'), (36, 'tensorflow.stack', 'tf.stack', (['row[1:]', '(1)'], {}), True, 'import tensorflow as tf\n'), (68, 'tensorflow.keras.optimizers.Adam', 'tf.keras.optimizers.Adam', (['lr_schedule'], {}), True, 'import tensorflow as tf\n'), (23, 'tempfile.mkdtemp', 'tempfile.mkdtemp', ([], {}), False, 'import tempfile\n'), (83, 'tensorflow.keras.callbacks.EarlyStopping', 'tf.keras.callbacks.EarlyStopping', ([], {'monitor': '"""val_binary_crossentropy"""', 'patience': '(200)'}), True, 'import tensorflow as tf\n'), (84, 'tensorflow.keras.callbacks.TensorBoard', 'tf.keras.callbacks.TensorBoard', (['(logdir / name)'], {}), True, 'import tensorflow as tf\n'), (111, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(16)'], {'activation': '"""elu"""', 'input_shape': '(FEATURES,)'}), False, 'from tensorflow.keras import layers\n'), (112, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), False, 'from tensorflow.keras import layers\n'), (125, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(16)'], {'activation': '"""elu"""', 'input_shape': '(FEATURES,)'}), False, 'from tensorflow.keras import layers\n'), (126, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(16)'], {'activation': '"""elu"""'}), False, 'from tensorflow.keras import layers\n'), (127, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), False, 'from tensorflow.keras import layers\n'), (133, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(64)'], {'activation': '"""elu"""', 'input_shape': '(FEATURES,)'}), False, 'from tensorflow.keras import layers\n'), (134, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(64)'], {'activation': '"""elu"""'}), False, 'from tensorflow.keras import layers\n'), (135, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(64)'], {'activation': '"""elu"""'}), False, 'from tensorflow.keras import layers\n'), (136, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), False, 'from tensorflow.keras import layers\n'), (142, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(512)'], {'activation': '"""elu"""', 'input_shape': '(FEATURES,)'}), False, 'from tensorflow.keras import layers\n'), (143, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(512)'], {'activation': '"""elu"""'}), False, 'from tensorflow.keras import layers\n'), (144, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(512)'], {'activation': '"""elu"""'}), False, 'from tensorflow.keras import layers\n'), (145, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(512)'], {'activation': '"""elu"""'}), False, 'from tensorflow.keras import layers\n'), (146, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(1)'], {}), False, 'from tensorflow.keras import layers\n'), (75, 'matplotlib.pyplot.ylim', 'plt.ylim', ([], {}), True, 'from matplotlib import pyplot as plt\n'), (92, 'tensorflow.keras.losses.BinaryCrossentropy', 'tf.keras.losses.BinaryCrossentropy', ([], {'from_logits': '(True)'}), True, 'import tensorflow as tf\n'), (153, 'matplotlib.pyplot.xlim', 'plt.xlim', ([], {}), True, 'from matplotlib import pyplot as plt\n'), (94, 'tensorflow.keras.losses.BinaryCrossentropy', 'tf.keras.losses.BinaryCrossentropy', ([], {'from_logits': '(True)', 'name': '"""binary_crossentropy"""'}), True, 'import tensorflow as tf\n')]