import os

from datetime import datetime
from packaging import version

import tensorflow as tf
from tensorflow import keras

import numpy as np

print("TensorFlow version: ", tf.__version__)
assert (
    version.parse(tf.__version__).release[0] >= 2
), "This script requires TensorFlow 2.0 or above."

data_size = 1000
# 80% of the data is for training.
train_pct = 0.8
lr = 0.2
train_epochs = 100

train_size = int(data_size * train_pct)

# Create some input data between -1 and 1 and randomize it.
x = np.linspace(-1, 1, data_size)
np.random.shuffle(x)

# Generate the output data.
# y = 0.5x + 2 + noise
y = 0.5 * x + 2 + np.random.normal(0, 0.05, (data_size,))

# Split into test and train pairs.
x_train, y_train = x[:train_size], y[:train_size]
x_test, y_test = x[train_size:], y[train_size:]

logdir = (
    os.getenv("LOGDIR") or ("logs/scalars/" + datetime.now().strftime("%Y%m%d-%H%M%S"))
)
file_writer = tf.summary.create_file_writer(logdir + "/metrics")
file_writer.set_as_default()


def lr_schedule(epoch):
    """
    Returns a custom learning rate that decreases as epochs progress.
    """
    learning_rate = 0.2
    if epoch > 10:
        learning_rate = 0.02
    if epoch > 20:
        learning_rate = 0.01
    if epoch > 50:
        learning_rate = 0.005

    tf.summary.scalar('learning rate', data=learning_rate, step=epoch)
    return learning_rate


lr_callback = keras.callbacks.LearningRateScheduler(lr_schedule)
tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logdir)

model = keras.models.Sequential(
    [
        keras.layers.Dense(16, input_dim=1),
        keras.layers.Dense(1),
    ]
)

model.compile(
    loss='mse',
    optimizer=keras.optimizers.SGD(),  # keras.losses.mean_squared_error
)

training_history = model.fit(
    x_train,  # input
    y_train,  # output
    batch_size=train_size,
    verbose=0,  # Suppress chatty output; use Tensorboard instead
    epochs=train_epochs,
    validation_data=(x_test, y_test),
    callbacks=[tensorboard_callback, lr_callback],
)
