import sys
import os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Tambahkan root directory ke path agar dapat import dari src
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from src.data_loader import load_mnist_for_mlp
from src.models import create_mlp_model, create_regularized_mlp_model
from src.utils import ensure_dir


def main():
    print("=" * 60)
    print("TUGAS 2: EKSPERIMEN REGULARISASI (BASELINE vs BATCHNORM + DROPOUT)")
    print("=" * 60)

    # 1. Load Data
    (x_train, y_train), (x_test, y_test) = load_mnist_for_mlp()

    EPOCHS = 20
    BATCH_SIZE = 128

    # 2. Baseline Model (Without Regularization)
    print("\n" + "=" * 60)
    print("TRAINING BASELINE MODEL (WITHOUT REGULARIZATION)")
    print("=" * 60)
    model_base = create_mlp_model([256, 128, 64])
    history_base = model_base.fit(
        x_train, y_train,
        epochs=EPOCHS,
        batch_size=BATCH_SIZE,
        validation_split=0.1,
        verbose=1
    )

    # 3. Regularized Model (BatchNorm + Dropout 0.3)
    print("\n" + "=" * 60)
    print("TRAINING REGULARIZED MODEL (WITH BATCHNORM + DROPOUT)")
    print("=" * 60)
    model_reg = create_regularized_mlp_model(hidden_layers=[256, 128, 64], dropout_rate=0.3)
    history_reg = model_reg.fit(
        x_train, y_train,
        epochs=EPOCHS,
        batch_size=BATCH_SIZE,
        validation_split=0.1,
        verbose=1
    )

    # 4. Save Model Artifact
    results_model_dir = os.path.join("results", "models")
    ensure_dir(results_model_dir)
    model_reg_path = os.path.join(results_model_dir, "model_c_regularized.keras")
    model_reg.save(model_reg_path)
    print(f"\n[INFO] Model teregularisasi disimpan di: {model_reg_path}")

    # 5. Plot Comparison Curves
    results_fig_dir = os.path.join("results", "figures")
    ensure_dir(results_fig_dir)
    fig_path = os.path.join(results_fig_dir, "loss_regularization.png")

    plt.figure(figsize=(14, 5))

    # Subplot 1: Baseline
    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, alpha=0.3)

    # Subplot 2: Regularized
    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, alpha=0.3)

    plt.tight_layout()
    plt.savefig(fig_path, dpi=130)
    print(f"[INFO] Grafik perbandingan loss disimpan di: {fig_path}")


if __name__ == "__main__":
    main()