"""
Tugas 1 (Week 2): Pipeline CNN MNIST Lengkap
Memenuhi seluruh poin instruksi:
(a) Model Summary
(b) Loss & Accuracy curve training vs validation
(c) Confusion Matrix pada test set
(d) Perbandingan accuracy MLP vs CNN dan analisis efisiensi parameter
"""

import os
import sys

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

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from sklearn.metrics import classification_report, confusion_matrix
import tensorflow as tf

# Tambahkan root directory ke path
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():
    SEED = 42
    EPOCHS = 10
    BATCH_SIZE = 128
    VAL_SPLIT = 0.1

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

    # Baseline MLP Terbaik (Model C dari Tugas 1 & 2)
    MLP_ACCURACY = 0.9824
    MLP_PARAMS = 244_554
    MLP_EPOCH_TIME = 2.50

    print("=" * 70)
    print("TUGAS 1: PIPELINE CNN MNIST LENGKAP")
    print("=" * 70)

    # 1. Load Data
    print("\n[1/5] 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 ({x_train.shape[1]}x{x_train.shape[2]}x{x_train.shape[3]})")
    print(f"  -> Data testing  : {x_test.shape[0]:,} gambar")

    # 2. Model Summary (Poin A)
    print("\n[2/5] Membangun Arsitektur CNN MNIST...")
    model = create_cnn_mnist_model()
    print("\n" + "-" * 70)
    print("(a) MODEL SUMMARY:")
    print("-" * 70)
    model.summary()

    # 3. Training Model
    print("\n[3/5] Memulai proses training (10 Epochs)...")
    timer_cb = EpochTimerCallback()

    history = model.fit(
        x_train,
        y_train,
        validation_split=VAL_SPLIT,
        epochs=EPOCHS,
        batch_size=BATCH_SIZE,
        verbose=1,
        callbacks=[timer_cb],
    )

    avg_epoch_time = float(np.mean(timer_cb.epoch_times)) if timer_cb.epoch_times else 0.0

    # 4. Evaluasi & Prediksi pada Test Set
    print("\n[4/5] Evaluasi pada Test Set (10,000 data)...")
    test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0)
    y_pred_probs = model.predict(x_test, verbose=0)
    y_pred = np.argmax(y_pred_probs, axis=1)

    print(f"  -> Test Accuracy : {test_acc * 100:.2f}%")
    print(f"  -> Test Loss     : {test_loss:.4f}")

    # 5. Visualisasi & Penyimpanan
    figures_dir = os.path.join("results", "figures")
    ensure_dir(figures_dir)

    # (b) Loss & Accuracy Curve
    print("\n[5/5] Membuat visualisasi...")
    curve_path = os.path.join(figures_dir, "loss_accuracy_cnn_lengkap.png")
    fig, axes = plt.subplots(1, 2, figsize=(13, 5))
    epochs_range = range(1, EPOCHS + 1)

    # Loss Curve
    axes[0].plot(epochs_range, history.history["loss"], "b-o", label="Training Loss")
    axes[0].plot(epochs_range, history.history["val_loss"], "r--s", label="Validation Loss")
    axes[0].set_title("(b1) Loss Curve: Training vs Validation", fontsize=12, fontweight="bold")
    axes[0].set_xlabel("Epoch", fontsize=10)
    axes[0].set_ylabel("Cross-Entropy Loss", fontsize=10)
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)

    # Accuracy Curve
    axes[1].plot(epochs_range, history.history["accuracy"], "b-o", label="Training Accuracy")
    axes[1].plot(epochs_range, history.history["val_accuracy"], "r--s", label="Validation Accuracy")
    axes[1].set_title("(b2) Accuracy Curve: Training vs Validation", fontsize=12, fontweight="bold")
    axes[1].set_xlabel("Epoch", fontsize=10)
    axes[1].set_ylabel("Accuracy", fontsize=10)
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(curve_path, dpi=150)
    plt.close()
    print(f"  -> (b) Kurva Loss & Accuracy tersimpan: {curve_path}")

    # (c) Confusion Matrix pada Test Set
    cm_path = os.path.join(figures_dir, "confusion_matrix_cnn.png")
    cm = confusion_matrix(y_test, y_pred)

    plt.figure(figsize=(9, 7.5))
    sns.heatmap(
        cm,
        annot=True,
        fmt="d",
        cmap="Blues",
        cbar=True,
        xticklabels=range(10),
        yticklabels=range(10),
        annot_kws={"size": 10}
    )
    plt.title("(c) Confusion Matrix CNN pada MNIST Test Set (10.000 Gambar)", fontsize=13, fontweight="bold", pad=12)
    plt.xlabel("Predicted Digit", fontsize=11, labelpad=8)
    plt.ylabel("True Digit", fontsize=11, labelpad=8)
    plt.tight_layout()
    plt.savefig(cm_path, dpi=150)
    plt.close()
    print(f"  -> (c) Confusion Matrix tersimpan: {cm_path}")

    # (d) Perbandingan Kuantitatif MLP vs CNN
    n_params_cnn = model.count_params()
    diff_acc = (test_acc - MLP_ACCURACY) * 100
    param_saving = (1 - n_params_cnn / MLP_PARAMS) * 100
    error_mlp = (1 - MLP_ACCURACY) * 100
    error_cnn = (1 - test_acc) * 100
    rel_error_red = (1 - error_cnn / error_mlp) * 100

    print("\n" + "=" * 70)
    print("(d) PERBANDINGAN ACCURACY & EFISIENSI: MLP vs CNN")
    print("=" * 70)
    print(f"{'Metrik Evaluasi':<28}{'MLP Terbaik':>18}{'CNN Sederhana':>18}")
    print("-" * 70)
    print(f"{'Test Accuracy':<28}{MLP_ACCURACY * 100:>17.2f}%{test_acc * 100:>17.2f}%")
    print(f"{'Error Rate':<28}{error_mlp:>17.2f}%{error_cnn:>17.2f}%")
    print(f"{'Jumlah Parameter':<28}{MLP_PARAMS:>18,}{n_params_cnn:>18,}")
    print(f"{'Waktu Training / Epoch':<28}{MLP_EPOCH_TIME:>16.2f}s{avg_epoch_time:>16.2f}s")
    print("-" * 70)
    print(f"-> Improvement Akurasi : +{diff_acc:.2f}% poin persentase")
    print(f"-> Pengurangan Error   : {rel_error_red:.1f}% penurunan error relatif")
    print(f"-> Efisiensi Parameter : CNN hemat {param_saving:.1f}% bobot dibanding MLP")
    print("=" * 70)

    # Analisis Singkat Tertulis
    print("\n" + "=" * 70)
    print("ANALISIS: EFISIENSI PARAMETER & PERFORMA (MLP vs CNN)")
    print("=" * 70)
    print(
        f"1. Jumlah Parameter:\n"
        f"   Model CNN memiliki {n_params_cnn:,} parameter, sedangkan MLP terbaik memiliki {MLP_PARAMS:,} parameter.\n"
        f"   CNN menggunakan {param_saving:.1f}% parameter lebih sedikit (sekitar separuh dari MLP).\n\n"
        f"2. Mana yang Lebih Efisien?\n"
        f"   CNN jauh lebih efisien. Meskipun memiliki parameter yang jauh lebih sedikit, CNN menghasilkan\n"
        f"   akurasi {test_acc * 100:.2f}% (lebih tinggi dari MLP {MLP_ACCURACY * 100:.2f}%).\n\n"
        f"3. Mengapa demikian?\n"
        f"   - Parameter Sharing: Filter konvolusi berukuran kecil (3x3) digunakan secara berulang di seluruh citra,\n"
        f"     sehingga tidak membutuhkan koneksi bobot penuh antar setiap piksel seperti Dense layer pada MLP.\n"
        f"   - Preservasi Spasial: CNN memproses data dalam representasi 2D asli (lebar x tinggi x channel),\n"
        f"     menjaga hubungan antar-piksel tetangga, tepi, dan kontur angka tulisan tangan."
    )
    print("=" * 70 + "\n")

    # Simpan bobot terbaru
    models_dir = os.path.join("results", "models")
    ensure_dir(models_dir)
    model.save(os.path.join(models_dir, "model_cnn_mnist.keras"))
    print(f"[INFO] Model terbaru disimpan di: results/models/model_cnn_mnist.keras")


if __name__ == "__main__":
    main()