import os
import shutil
from pathlib import Path

os.environ["KERAS_BACKEND"] = "jax"

import keras
import numpy as np


class InterruptAtEpoch(keras.callbacks.Callback):
    def __init__(self, epoch_to_stop):
        super().__init__()
        self.epoch_to_stop = epoch_to_stop

    def on_epoch_begin(self, epoch, logs=None):
        if epoch == self.epoch_to_stop:
            raise RuntimeError(f"simulated interruption at epoch {epoch}")


def build_model():
    model = keras.Sequential(
        [
            keras.Input(shape=(4,), name="features"),
            keras.layers.Dense(8, activation="relu"),
            keras.layers.Dense(1, activation="sigmoid"),
        ]
    )
    model.compile(
        optimizer=keras.optimizers.Adam(learning_rate=0.03),
        loss=keras.losses.BinaryCrossentropy(),
        metrics=[keras.metrics.BinaryAccuracy(name="accuracy")],
    )
    return model


keras.utils.set_random_seed(17)
backup_dir = Path("training_backup")
if backup_dir.exists():
    shutil.rmtree(backup_dir)

x_train = np.linspace(0.0, 1.0, 96, dtype="float32").reshape(24, 4)
y_train = ((x_train[:, 0] + x_train[:, 2]) > 0.75).astype("float32")

model = build_model()
callback = keras.callbacks.BackupAndRestore(backup_dir=backup_dir)

try:
    model.fit(
        x_train,
        y_train,
        epochs=6,
        batch_size=6,
        callbacks=[callback, InterruptAtEpoch(3)],
        verbose=0,
    )
except RuntimeError as exc:
    print(f"first run stopped: {exc}")

resumed_model = build_model()
resume_callback = keras.callbacks.BackupAndRestore(backup_dir=backup_dir)
history = resumed_model.fit(
    x_train,
    y_train,
    epochs=6,
    batch_size=6,
    callbacks=[resume_callback],
    verbose=0,
)

print(f"backend: {keras.backend.backend()}")
print(f"backup directory: {backup_dir}")
print("interrupted before epoch: 3")
print(f"resumed epochs completed: {len(history.history['loss'])}")
print(f"final accuracy: {history.history['accuracy'][-1]:.4f}")
print("training target reached: epoch 6")