"""
Tugas No. 4 - Implementasi CNN sederhana untuk MNIST
dan perbandingannya dengan MLP dari Tugas No. 1.

Cara menjalankan di VS Code:
    1. Aktifkan virtual environment
    2. pip install tensorflow numpy matplotlib
    3. Klik tombol Run, atau jalankan: python task4_cnn.py

Dataset akan diunduh otomatis saat pertama kali dijalankan
dan disimpan di ~/.keras/datasets/
"""

import os

# Sembunyikan log informasi TensorFlow (harus sebelum import tensorflow)
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import time

import matplotlib
matplotlib.use("Agg")  # simpan plot sebagai file, tidak membuka jendela
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import keras
from keras import layers

#Konfigurasi
SEED = 42
EPOCHS = 10
BATCH_SIZE = 128
VAL_SPLIT = 0.1

tf.random.set_seed(SEED)
np.random.seed(SEED)

# Hasil MLP terbaik dari Tugas No. 1 (Model C: 256-128-64).
# Ganti angka di bawah ini dengan hasil eksperimen Anda sendiri.
MLP_ACCURACY = 0.9764
MLP_PARAMS = 242_762
MLP_EPOCH_TIME = 2.50

#1. Load dan preprocessing data
print("Memuat dataset MNIST...")
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

print(f"Data training : {x_train.shape[0]:,} gambar")
print(f"Data testing  : {x_test.shape[0]:,} gambar")
print(f"Ukuran gambar : {x_train.shape[1]} x {x_train.shape[2]} piksel\n")

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

#CNN membutuhkan dimensi channel: (N, 28, 28) -> (N, 28, 28, 1)
#Berbeda dengan MLP yang meratakan gambar menjadi (N, 784)
x_train = x_train.reshape(-1, 28, 28, 1)
x_test = x_test.reshape(-1, 28, 28, 1)

#2. Bangun model CNN

model = keras.Sequential(
    name="cnn_mnist",
    layers=[
        layers.Input(shape=(28, 28, 1)),

        # Blok konvolusi 1: deteksi fitur sederhana (tepi, garis)
        layers.Conv2D(32, (3, 3), activation="relu"),
        layers.MaxPooling2D((2, 2)),

        # Blok konvolusi 2: gabungkan menjadi fitur lebih kompleks
        layers.Conv2D(64, (3, 3), activation="relu"),
        layers.MaxPooling2D((2, 2)),

        # Klasifikasi
        layers.Flatten(),
        layers.Dense(64, activation="relu"),
        layers.Dropout(0.3),
        layers.Dense(10, activation="softmax"),
    ],
)

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

model.summary()

# 3. Training dengan pencatatan waktu per epoch
epoch_times = []

class TimeCallback(keras.callbacks.Callback):
    """Mencatat durasi setiap epoch."""

    def on_epoch_begin(self, epoch, logs=None):
        self._start = time.time()

    def on_epoch_end(self, epoch, logs=None):
        epoch_times.append(time.time() - self._start)


print("\nMemulai training CNN...")
print("Perkiraan waktu: 15-20 detik per epoch di CPU (total sekitar 3 menit).\n")

history = model.fit(
    x_train,
    y_train,
    validation_split=VAL_SPLIT,
    epochs=EPOCHS,
    batch_size=BATCH_SIZE,
    verbose=2,
    callbacks=[TimeCallback()],
)
# 4. Evaluasi
test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0)
n_params = model.count_params()
avg_epoch_time = float(np.mean(epoch_times))

predictions = np.argmax(model.predict(x_test, verbose=0), axis=1)
n_wrong = int(np.sum(predictions != y_test))

# 5. Perbandingan CNN vs MLP
error_mlp = (1 - MLP_ACCURACY) * 100
error_cnn = (1 - test_acc) * 100
gain_point = (test_acc - MLP_ACCURACY) * 100
error_reduction = (1 - error_cnn / error_mlp) * 100

print("\n" + "=" * 62)
print("HASIL EKSPERIMEN CNN")
print("=" * 62)
print(f"Test accuracy         : {test_acc * 100:.2f}%")
print(f"Test loss             : {test_loss:.4f}")
print(f"Jumlah parameter      : {n_params:,}")
print(f"Waktu training/epoch  : {avg_epoch_time:.2f} detik")
print(f"Salah klasifikasi     : {n_wrong} dari {len(y_test):,} gambar")

print("\n" + "=" * 62)
print("PERBANDINGAN CNN vs MLP")
print("=" * 62)
print(f"{'Metrik':<24}{'MLP':>14}{'CNN':>14}")
print("-" * 62)
print(f"{'Test accuracy':<24}{MLP_ACCURACY * 100:>13.2f}%{test_acc * 100:>13.2f}%")
print(f"{'Error rate':<24}{error_mlp:>13.2f}%{error_cnn:>13.2f}%")
print(f"{'Jumlah parameter':<24}{MLP_PARAMS:>14,}{n_params:>14,}")
print(f"{'Waktu/epoch (detik)':<24}{MLP_EPOCH_TIME:>14.2f}{avg_epoch_time:>14.2f}")
print("-" * 62)
print(f"Improvement akurasi      : +{gain_point:.2f} poin persentase")
print(f"Pengurangan error rate   : {error_reduction:.1f}% (relatif)")
print(f"Selisih parameter        : {n_params - MLP_PARAMS:+,}")
print("=" * 62)

# 6. Simpan grafik
fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))

epochs_range = range(1, EPOCHS + 1)
axes[0].plot(epochs_range, history.history["accuracy"], label="Train accuracy")
axes[0].plot(epochs_range, history.history["val_accuracy"], label="Validation accuracy")
axes[0].set_title("Learning Curve CNN")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("Accuracy")
axes[0].legend()
axes[0].grid(alpha=0.3)

axes[1].bar(
    ["MLP terbaik", "CNN"],
    [MLP_ACCURACY * 100, test_acc * 100],
    color=["#55A868", "#C44E52"],
)
axes[1].set_ylim(96, 100)
axes[1].set_ylabel("Test Accuracy (%)")
axes[1].set_title("Perbandingan MLP vs CNN")
for i, value in enumerate([MLP_ACCURACY * 100, test_acc * 100]):
    axes[1].text(i, value + 0.05, f"{value:.2f}%", ha="center")

plt.tight_layout()
plt.savefig("hasil_cnn.png", dpi=130)
print("\nGrafik disimpan sebagai: hasil_cnn.png")

# Simpan model terlatih agar tidak perlu training ulang
model.save("model_cnn_mnist.keras")
print("Model disimpan sebagai: model_cnn_mnist.keras")