import tensorflow as tf
from tensorflow.keras import layers, models

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
y_train = tf.keras.utils.to_categorical(y_train, 10)
y_test = tf.keras.utils.to_categorical(y_test, 10)

def build_tl_model(base_cls):
    base = base_cls(include_top=False, weights='imagenet', input_shape=(32, 32, 3))
    base.trainable = False
    return models.Sequential([
        base,
        layers.GlobalAveragePooling2D(),
        layers.Dense(128, activation='relu'),
        layers.Dense(10, activation='softmax')
    ])

model_fe = build_tl_model(tf.keras.applications.DenseNet121)

model_ft = model_fe
for layer in model_ft.layers:
    layer.trainable = True

model_ft.compile(
    optimizer=tf.keras.optimizers.Adam(1e-5),  # note: 100x smaller LR
    loss='categorical_crossentropy',
    metrics=['accuracy']
)
history_ft = model_ft.fit(x_train, y_train, epochs=5, validation_data=(x_test, y_test))

model_nf = build_tl_model(tf.keras.applications.DenseNet121)
model_nf.trainable = True
model_nf.compile(optimizer=tf.keras.optimizers.Adam(1e-4),
                  loss='categorical_crossentropy', metrics=['accuracy'])
history_nf = model_nf.fit(x_train, y_train, epochs=5, validation_data=(x_test, y_test))

base = tf.keras.applications.DenseNet121(include_top=False, weights="imagenet",
                                          input_shape=(32, 32, 3))
for layer in base.layers[:100]:
    layer.trainable = False
for layer in base.layers[100:]:
    layer.trainable = True

model_cf = models.Sequential([base, layers.GlobalAveragePooling2D(),
                               layers.Dense(128, activation='relu'),
                               layers.Dense(10, activation='softmax')])
model_cf.compile(optimizer=tf.keras.optimizers.Adam(1e-4),
                  loss="categorical_crossentropy", metrics=["accuracy"])

scratch = models.Sequential([
    layers.Input((32, 32, 3)),
    layers.Conv2D(32, 3, activation='relu', padding='same'), layers.MaxPool2D(),
    layers.Conv2D(64, 3, activation='relu', padding='same'), layers.MaxPool2D(),
    layers.GlobalAveragePooling2D(),
    layers.Dropout(0.2),
    layers.Dense(10, activation='softmax')
])
scratch.compile('adam', loss="categorical_crossentropy", metrics=["accuracy"])
history_s = scratch.fit(x_train, y_train, epochs=10, validation_data=(x_test, y_test))

plt.plot(history_fe.history['val_accuracy'], 'o-', label='Freeze+TrainHead')
plt.plot(history_ft.history['val_accuracy'], 'o-', label='Unfreeze+FineTune')
plt.plot(history_nf.history['val_accuracy'], 'o-', label='No Freeze->Fine-tuneAll')
plt.plot(history_cf.history['val_accuracy'], 'o-', label='Custom-Freeze')
plt.plot(history_s.history['val_accuracy'], 'o-', label='Scratch CNN')
plt.title("Validation Accuracy Comparison")
plt.legend()