import os

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

import keras
import numpy as np
from keras import layers


keras.utils.set_random_seed(23)

num_classes = 2
image_shape = (96, 96, 3)

rng = np.random.default_rng(23)
x_train = rng.uniform(size=(8, *image_shape)).astype("float32")
y_train = np.array([0, 1, 0, 1, 0, 1, 0, 1], dtype="int32")
x_val = rng.uniform(size=(4, *image_shape)).astype("float32")
y_val = np.array([0, 1, 0, 1], dtype="int32")

base_model = keras.applications.MobileNetV2(
    input_shape=image_shape,
    include_top=False,
    weights="imagenet",
    pooling="avg",
)
base_model.trainable = False

inputs = keras.Input(shape=image_shape, name="image")
x = keras.applications.mobilenet_v2.preprocess_input(inputs)
x = base_model(x, training=False)
x = layers.Dropout(0.2)(x)
outputs = layers.Dense(num_classes, activation="softmax", name="class_probs")(x)
model = keras.Model(inputs, outputs, name="mobilenet_transfer_demo")

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.001),
    loss=keras.losses.SparseCategoricalCrossentropy(),
    metrics=[keras.metrics.SparseCategoricalAccuracy(name="accuracy")],
)
history = model.fit(x_train, y_train, validation_data=(x_val, y_val), epochs=1, batch_size=4, verbose=0)
frozen_trainable_weights = len(model.trainable_weights)

base_model.trainable = True
for layer in base_model.layers[:-20]:
    layer.trainable = False
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.0001),
    loss=keras.losses.SparseCategoricalCrossentropy(),
    metrics=[keras.metrics.SparseCategoricalAccuracy(name="accuracy")],
)
fine_tune_history = model.fit(
    x_train,
    y_train,
    validation_data=(x_val, y_val),
    epochs=2,
    initial_epoch=1,
    batch_size=4,
    verbose=0,
)
model.save("transfer-learning-demo.keras")

print(f"backend: {keras.backend.backend()}")
print(f"base model: {base_model.name}")
print("base weights: imagenet")
print(f"frozen trainable weights: {frozen_trainable_weights}")
print(f"fine-tune trainable weights: {len(model.trainable_weights)}")
print(f"final val accuracy: {fine_tune_history.history['val_accuracy'][-1]:.4f}")
print("saved model: transfer-learning-demo.keras")