import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt

# 1. LOAD & PREPROCESS MNIST DATASET
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0
x_train = x_train.reshape(-1, 784)
x_test = x_test.reshape(-1, 784)

# Model C Architecture from Task 1
BEST_ARCHITECTURE = [256, 128, 64]
EPOCHS = 20  # 20 epochs gives enough room to observe overfitting
BATCH_SIZE = 128


# ============================================================
# MODEL 1: WITHOUT REGULARIZATION (Baseline Model C from Task 1)
# ============================================================
model_base = keras.Sequential([
    layers.Input(shape=(784,)),
    layers.Dense(256, activation="relu"),
    layers.Dense(128, activation="relu"),
    layers.Dense(64, activation="relu"),
    layers.Dense(10, activation="softmax")
])

model_base.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)

print("=" * 60)
print("TRAINING BASELINE MODEL (WITHOUT REGULARIZATION)")
print("=" * 60)

history_base = model_base.fit(
    x_train, y_train,
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    validation_split=0.1,
    verbose=1
)


# ============================================================
# MODEL 2: WITH REGULARIZATION (Batch Normalization + Dropout 0.3)
# ============================================================
model_reg = keras.Sequential([
    layers.Input(shape=(784,)),
    
    layers.Dense(256),
    layers.BatchNormalization(),
    layers.Activation("relu"),
    layers.Dropout(0.3),
    
    layers.Dense(128),
    layers.BatchNormalization(),
    layers.Activation("relu"),
    layers.Dropout(0.3),
    
    layers.Dense(64),
    layers.BatchNormalization(),
    layers.Activation("relu"),
    layers.Dropout(0.3),
    
    layers.Dense(10, activation="softmax")
])

model_reg.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)

print("\n" + "=" * 60)
print("TRAINING REGULARIZED MODEL (WITH BATCHNORM + DROPOUT)")
print("=" * 60)

history_reg = model_reg.fit(
    x_train, y_train,
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    validation_split=0.1,
    verbose=1
)

# Optional: Save the trained regularized model to disk
model_reg.save("model_c_regularized.keras")
print("\n[INFO] Regularized model saved as 'model_c_regularized.keras' in your project directory.")


# ============================================================
# PLOTTING LOSS CURVES FOR COMPARISON
# ============================================================
plt.figure(figsize=(14, 5))

# Plot 1: Baseline (Without Regularization)
plt.subplot(1, 2, 1)
plt.plot(history_base.history['loss'], label='Training Loss', color='blue')
plt.plot(history_base.history['val_loss'], label='Validation Loss', color='orange', linestyle='--')
plt.title('Model C Without Regularization')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)

# Plot 2: Regularized (With BatchNorm + Dropout)
plt.subplot(1, 2, 2)
plt.plot(history_reg.history['loss'], label='Training Loss', color='blue')
plt.plot(history_reg.history['val_loss'], label='Validation Loss', color='green', linestyle='--')
plt.title('Model C With BatchNorm + Dropout(0.3)')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()