import keras
import numpy as np


keras.utils.set_random_seed(21)

x_train = np.linspace(-1.0, 1.0, 48, dtype="float32").reshape(-1, 1)
y_train = (3.0 * x_train) - 0.25


class LearningRatePrinter(keras.callbacks.Callback):
    def on_train_begin(self, logs=None):
        self.rates = []

    def on_epoch_begin(self, epoch, logs=None):
        lr = float(keras.ops.convert_to_numpy(self.model.optimizer.learning_rate))
        self.rates.append(lr)
        print(f"epoch={epoch + 1} lr={lr:.4f}")


def schedule(epoch, lr):
    if epoch < 2:
        return lr
    return lr * 0.5


model = keras.Sequential(
    [
        keras.layers.Input(shape=(1,)),
        keras.layers.Dense(8, activation="relu"),
        keras.layers.Dense(1),
    ]
)
model.compile(
    optimizer=keras.optimizers.SGD(learning_rate=0.1),
    loss="mse",
)

lr_printer = LearningRatePrinter()
lr_scheduler = keras.callbacks.LearningRateScheduler(schedule)

history = model.fit(
    x_train,
    y_train,
    epochs=5,
    batch_size=8,
    callbacks=[lr_scheduler, lr_printer],
    verbose=0,
)

rates = ", ".join(f"{rate:.4f}" for rate in lr_printer.rates)
print(f"rates seen: [{rates}]")
print(f"final loss: {history.history['loss'][-1]:.4f}")