import sys
import os
import time

# 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


def main():
    print("=" * 60)
    print("TUGAS 1: MEMUAT DATASET MNIST & EKSPERIMEN MLP")
    print("=" * 60)

    # 1. Load Data
    (x_train, y_train), (x_test, y_test) = load_mnist_for_mlp()
    print(f"Data training : {x_train.shape}")
    print(f"Data testing  : {x_test.shape}")

    # 2. Konfigurasi 3 Model MLP
    models_config = {
        "Model A - 1 Hidden Layer": [64],
        "Model B - 2 Hidden Layer": [128, 64],
        "Model C - 3 Hidden Layer": [256, 128, 64]
    }

    results = []
    EPOCHS = 10
    BATCH_SIZE = 128

    # 3. Training & Evaluasi
    for model_name, architecture in models_config.items():
        print("\n" + "=" * 60)
        print(model_name)
        print("=" * 60)

        model = create_mlp_model(architecture)
        model.summary()

        print("\nMemulai training...")
        start_time = time.time()

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

        total_time = time.time() - start_time
        average_time = total_time / EPOCHS

        test_loss, test_accuracy = model.evaluate(x_test, y_test, verbose=0)
        total_parameters = model.count_params()

        results.append({
            "model": model_name,
            "parameters": total_parameters,
            "accuracy": test_accuracy * 100,
            "time_per_epoch": average_time
        })

        print("\nHasil:")
        print(f"Test Accuracy     : {test_accuracy * 100:.2f}%")
        print(f"Jumlah Parameter  : {total_parameters:,}")
        print(f"Waktu/Epoch       : {average_time:.2f} detik")

    # 4. Ringkasan Perbandingan
    print("\n\n" + "=" * 80)
    print("HASIL PERBANDINGAN 3 MODEL MLP")
    print("=" * 80)
    print(f"{'Model':<30}{'Parameter':>15}{'Accuracy':>15}{'Time/Epoch':>15}")
    print("-" * 80)

    for res in results:
        print(
            f"{res['model']:<30}"
            f"{res['parameters']:>15,}"
            f"{res['accuracy']:>14.2f}%"
            f"{res['time_per_epoch']:>14.2f}s"
        )

    best_model = max(results, key=lambda x: x["accuracy"])
    print("\n" + "=" * 60)
    print("MODEL DENGAN TEST ACCURACY TERTINGGI")
    print("=" * 60)
    print(f"Model               : {best_model['model']}")
    print(f"Test Accuracy       : {best_model['accuracy']:.2f}%")
    print(f"Jumlah Parameter    : {best_model['parameters']:,}")
    print(f"Waktu Training/Epoch: {best_model['time_per_epoch']:.2f} detik")


if __name__ == "__main__":
    main()