import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import time


# ============================================================
# 1. LOAD DATASET MNIST
# ============================================================

print("=" * 60)
print("MEMUAT DATASET MNIST")
print("=" * 60)

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

print(f"Data training : {x_train.shape}")
print(f"Data testing  : {x_test.shape}")


# ============================================================
# 2. PREPROCESSING
# ============================================================

# Normalisasi nilai pixel dari 0-255 menjadi 0-1
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

# Mengubah gambar 28x28 menjadi 784 fitur
x_train = x_train.reshape(-1, 784)
x_test = x_test.reshape(-1, 784)

print(f"Training setelah preprocessing: {x_train.shape}")
print(f"Testing setelah preprocessing : {x_test.shape}")


# ============================================================
# 3. FUNGSI MEMBUAT MODEL
# ============================================================

def create_model(hidden_layers):

    model = keras.Sequential()

    # Input layer
    model.add(layers.Input(shape=(784,)))

    # Hidden layers
    for neurons in hidden_layers:
        model.add(
            layers.Dense(
                neurons,
                activation="relu"
            )
        )

    # Output layer
    model.add(
        layers.Dense(
            10,
            activation="softmax"
        )
    )

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

    return model


# ============================================================
# 4. KONFIGURASI 3 MODEL
# ============================================================

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


# ============================================================
# 5. TRAINING DAN EVALUASI
# ============================================================

results = []

EPOCHS = 10
BATCH_SIZE = 128

for model_name, architecture in models_config.items():

    print("\n")
    print("=" * 60)
    print(model_name)
    print("=" * 60)

    # Membuat model
    model = create_model(architecture)

    # Menampilkan arsitektur
    model.summary()

    # Training
    print("\nMemulai training...")

    start_time = time.time()

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

    end_time = time.time()

    # Total waktu training
    total_time = end_time - start_time

    # Rata-rata waktu per epoch
    average_time = total_time / EPOCHS

    # Evaluasi menggunakan test data
    test_loss, test_accuracy = model.evaluate(
        x_test,
        y_test,
        verbose=0
    )

    # Jumlah parameter
    total_parameters = model.count_params()

    # Menyimpan hasil
    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")


# ============================================================
# 6. HASIL PERBANDINGAN
# ============================================================

print("\n\n")
print("=" * 80)
print("HASIL PERBANDINGAN 3 MODEL MLP")
print("=" * 80)

print(
    f"{'Model':<30}"
    f"{'Parameter':>15}"
    f"{'Accuracy':>15}"
    f"{'Time/Epoch':>15}"
)

print("-" * 80)

for result in results:

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


# ============================================================
# 7. MENENTUKAN MODEL DENGAN ACCURACY TERTINGGI
# ============================================================

best_model = max(
    results,
    key=lambda x: x["accuracy"]
)

print("\n")
print("=" * 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")


# ============================================================
# 8. KESIMPULAN SINGKAT
# ============================================================

print("\n")
print("=" * 60)
print("KESIMPULAN")
print("=" * 60)

print(
    "Model yang lebih dalam memiliki jumlah parameter "
    "yang lebih banyak dan umumnya membutuhkan waktu "
    "training yang lebih besar."
)

print(
    "Namun, model yang lebih dalam tidak selalu menghasilkan "
    "peningkatan test accuracy yang signifikan."
)

print(
    "Pemilihan arsitektur sebaiknya mempertimbangkan "
    "accuracy, jumlah parameter, dan waktu training."
)