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

import os
import sys
import time

# Sembunyikan log informasi TensorFlow
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import matplotlib
matplotlib.use("Agg")  # simpan plot sebagai file
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

# 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_cnn
from src.models import create_cnn_mnist_model
from src.utils import EpochTimerCallback, ensure_dir


def main():
    # 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)
    MLP_ACCURACY = 0.9764
    MLP_PARAMS = 242_762
    MLP_EPOCH_TIME = 2.50

    # 1. Load dan preprocessing data
    print("Memuat dataset MNIST untuk CNN...")
    (x_train, y_train), (x_test, y_test) = load_mnist_for_cnn()

    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 (channel={x_train.shape[3]})\n")

    # 2. Bangun model CNN
    model = create_cnn_mnist_model()
    model.summary()

    # 3. Training dengan pencatatan waktu per epoch
    timer_cb = EpochTimerCallback()

    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=[timer_cb],
    )

    # 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(timer_cb.epoch_times)) if timer_cb.epoch_times else 0.0

    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 ke folder results/figures
    results_fig_dir = os.path.join("results", "figures")
    ensure_dir(results_fig_dir)
    fig_path = os.path.join(results_fig_dir, "hasil_cnn.png")

    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(fig_path, dpi=130)
    print(f"\n[INFO] Grafik disimpan di: {fig_path}")

    # 7. Simpan model terlatih ke folder results/models
    results_model_dir = os.path.join("results", "models")
    ensure_dir(results_model_dir)
    model_path = os.path.join(results_model_dir, "model_cnn_mnist.keras")
    model.save(model_path)
    print(f"[INFO] Model disimpan di: {model_path}")


if __name__ == "__main__":
    main()