from keras.src.legacy.preprocessing.image import ImageDataGenerator
from keras.src.saving import load_model
from keras.src.callbacks import ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
import tensorflow as tf

# Run TensorFlow functions eagerly (useful for debugging)
tf.config.run_functions_eagerly(True)

# Load the model
model = load_model("models/skin_lesion_classifier.h5")

# Compile the model
model.compile(optimizer="adam", loss="categorical_crossentropy", metrics=["accuracy"])

# Define data generators for training with augmentation
train_datagen = ImageDataGenerator(
    validation_split=0.2,
    horizontal_flip=True,  
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    zoom_range=0.2,
    shear_range=0.2
)

train_generator = train_datagen.flow_from_directory(
    "dataset/processed_images/output/",
    target_size=(224, 224),
    batch_size=32,
    class_mode="categorical",
    subset="training"
)

# Define callbacks
checkpoint = ModelCheckpoint(
    "models/skin_lesion_classifier_trained.h5",  # Save the best model
    save_best_only=True,
    monitor="val_loss",
    mode="min",
    verbose=1
)

reduce_lr = ReduceLROnPlateau(
    monitor="val_loss",
    factor=0.5,
    patience=3,
    min_lr=1e-6
)

early_stopping = EarlyStopping(
    monitor="val_loss",
    patience=5,
    restore_best_weights=True
)

# Define validation data generator
val_datagen = ImageDataGenerator(validation_split=0.2)

val_generator = val_datagen.flow_from_directory(
    "dataset/processed_images/output/",
    target_size=(224, 224),
    batch_size=32,
    class_mode="categorical",
    subset="validation"
)

# Train the model with callbacks
weights_before = model.get_weights()  # Weights before training
history = model.fit(
    train_generator,
    validation_data=val_generator,
    epochs=5,  # Total number of epochs
    verbose=1,
    callbacks=[checkpoint, reduce_lr, early_stopping]
)
weights_after = model.get_weights()  # Weights after training

if weights_before == weights_after:
    print("Weights were successfully saved")
else:
    print("Weights have changed")

# Save the trained model
model.save("models/skin_lesion_classifier_trained.keras")