"""
Tugas 2 (Week 2): Eksplorasi Kedalaman Arsitektur CNN
Memodifikasi CNN dengan menambahkan satu blok Conv2D + MaxPooling ketiga.
Mengevaluasi:
1. Apakah model yang lebih dalam menghasilkan akurasi lebih tinggi?
2. Bagaimana perbandingan waktu training per epoch?
3. Catatan temuan dalam satu paragraf analitis.
"""

import os
import sys
import time

# 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 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, create_deep_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)

    print("=" * 75)
    print("TUGAS 2: EKSPLORASI KEDALAMAN ARSITEKTUR CNN (2 BLOK vs 3 BLOK)")
    print("=" * 75)

    # 1. Load Data
    print("\nMemuat dataset MNIST...")
    (x_train, y_train), (x_test, y_test) = load_mnist_for_cnn()

    architectures = [
        {
            "name": "CNN 2 Blok (Baseline)",
            "builder": create_cnn_mnist_model,
            "desc": "Conv(32) + MaxPool -> Conv(64) + MaxPool -> Dense(64)"
        },
        {
            "name": "CNN 3 Blok (Lebih Dalam)",
            "builder": create_deep_cnn_mnist_model,
            "desc": "Conv(32) + MaxPool -> Conv(64) + MaxPool -> Conv(128) + MaxPool -> Dense(64)"
        }
    ]

    results = []
    histories = {}

    for arch in architectures:
        model_name = arch["name"]
        print("\n" + "=" * 75)
        print(f"PELATIHAN: {model_name}")
        print(f"Konfigurasi: {arch['desc']}")
        print("=" * 75)

        model = arch["builder"]()
        model.summary()

        timer_cb = EpochTimerCallback()
        start_train_time = time.time()

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

        total_train_time = time.time() - start_train_time
        avg_epoch_time = float(np.mean(timer_cb.epoch_times)) if timer_cb.epoch_times else 0.0

        # Evaluasi pada test set
        test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0)
        n_params = model.count_params()

        results.append({
            "name": model_name,
            "params": n_params,
            "test_acc": test_acc * 100,
            "test_loss": test_loss,
            "epoch_time": avg_epoch_time,
            "total_time": total_train_time,
            "model_obj": model
        })

        histories[model_name] = history

    # Simpan model 3 blok
    models_dir = os.path.join("results", "02_cnn", "models")
    ensure_dir(models_dir)
    results[1]["model_obj"].save(os.path.join(models_dir, "model_cnn_3blocks.keras"))

    # Tabel Perbandingan Kuantitatif
    print("\n\n" + "=" * 75)
    print("HASIL PERBANDINGAN: KEDALAMAN CNN PADA MNIST")
    print("=" * 75)
    print(f"{'Arsitektur Model':<28}{'Parameter':>14}{'Test Acc':>14}{'Waktu/Epoch':>18}")
    print("-" * 75)
    for r in results:
        print(f"{r['name']:<28}{r['params']:>14,}{r['test_acc']:>13.2f}%{r['epoch_time']:>16.2f}s")
    print("-" * 75)

    diff_acc = results[1]["test_acc"] - results[0]["test_acc"]
    time_increase = ((results[1]["epoch_time"] - results[0]["epoch_time"]) / results[0]["epoch_time"]) * 100
    param_diff = results[1]["params"] - results[0]["params"]

    print(f"Selisih Akurasi (3 Blok vs 2 Blok) : {diff_acc:+.2f}% poin")
    print(f"Perubahan Waktu Training / Epoch  : {time_increase:+.1f}%")
    print(f"Selisih Jumlah Parameter          : {param_diff:+,} parameter")
    print("=" * 75)

    # Visualisasi
    figures_dir = os.path.join("results", "02_cnn", "figures")
    ensure_dir(figures_dir)
    fig_path = os.path.join(figures_dir, "perbandingan_kedalaman_cnn.png")

    fig, axes = plt.subplots(1, 2, figsize=(13, 5))
    epochs_range = range(1, EPOCHS + 1)

    # Plot Accuracy Perbandingan
    for name, hist in histories.items():
        axes[0].plot(epochs_range, hist.history["val_accuracy"], label=f"{name} (Val Acc)", linewidth=2)
    axes[0].set_title("Validation Accuracy per Epoch", fontsize=12, fontweight="bold")
    axes[0].set_xlabel("Epoch", fontsize=10)
    axes[0].set_ylabel("Accuracy", fontsize=10)
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)

    # Bar chart perbandingan Waktu vs Akurasi
    labels = [r["name"] for r in results]
    accs = [r["test_acc"] for r in results]
    times = [r["epoch_time"] for r in results]

    x = np.arange(len(labels))
    width = 0.35

    ax2 = axes[1].twinx()
    p1 = axes[1].bar(x - width/2, accs, width, label="Test Accuracy (%)", color="#3498db")
    p2 = ax2.bar(x + width/2, times, width, label="Waktu / Epoch (s)", color="#e67e22")

    axes[1].set_ylabel("Accuracy (%)", color="#3498db", fontsize=10)
    ax2.set_ylabel("Detik per Epoch", color="#e67e22", fontsize=10)
    axes[1].set_xticks(x)
    axes[1].set_xticklabels(labels, fontsize=10)
    axes[1].set_ylim(97, 100)
    axes[1].set_title("Trade-off: Test Accuracy vs Waktu Training", fontsize=12, fontweight="bold")

    # Nilai di atas bar
    for bar in p1:
        axes[1].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.05, f"{bar.get_height():.2f}%", ha='center', va='bottom', fontsize=9)
    for bar in p2:
        ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.2, f"{bar.get_height():.2f}s", ha='center', va='bottom', fontsize=9)

    plt.tight_layout()
    plt.savefig(fig_path, dpi=150)
    plt.close()
    print(f"\n[INFO] Grafik perbandingan tersimpan di: {fig_path}")

    # Catatan Analisis Satu Paragraf (Sesuai Instruksi Tugas 2)
    print("\n" + "=" * 75)
    print("CATATAN TEMUAN (SATU PARAGRAF ANALISIS):")
    print("=" * 75)
    paragraf = (
        f"Penambahan blok ketiga (Conv2D 128 filter + MaxPooling2D) pada citra MNIST (28x28 piksel) "
        f"menghasilkan akurasi test sebesar {results[1]['test_acc']:.2f}% dibandingkan dengan {results[0]['test_acc']:.2f}% "
        f"pada baseline 2 blok (selisih {diff_acc:+.2f}% poin), namun meningkatkan waktu komputasi training "
        f"sebesar {time_increase:+.1f}% per epoch ({results[1]['epoch_time']:.2f}s vs {results[0]['epoch_time']:.2f}s) "
        f"serta menambah kompleksitas parameter. Pada dataset sederhana dengan resolusi spasial kecil seperti MNIST, "
        f"tiga kali operasi pooling mereduksi resolusi spasial representasi secara drastis (hingga 3x3 piksel), "
        f"sehingga arsitektur yang lebih dalam mengalami fenomena diminishing returns di mana penambahan layer "
        f"tidak memberikan peningkatan performa yang sebanding dengan biaya komputasi yang dikeluarkan; "
        f"oleh karena itu, arsitektur CNN 2 blok tetap menjadi pilihan paling optimal dan efisien untuk MNIST."
    )
    print(paragraf)
    print("=" * 75 + "\n")


if __name__ == "__main__":
    main()