import os
import gc
import numpy as np
import tensorflow as tf
import cv2
import gradio as gr
from PIL import Image
from tensorflow.keras.models import load_model, Model

# ==========================================
# 1. SETUP KELAS & VARIABEL GLOBAL
# ==========================================
class_names = ['Infiltrat Pneumonia', 'Normal']
last_conv_layer_name = 'block5_conv3'

# Dictionary pemetaan Model Dropdown -> Nama File
MODEL_DICT = {
    "Eksperimen 1 (LR 1e-1)": "model_vgg16_1.keras",
    "Eksperimen 2 (LR 1e-2)": "model_vgg16_2.keras",
    "Eksperimen 3 (LR 1e-3)": "model_vgg16_3.keras",
    "Eksperimen 4 (LR 1e-4)": "model_vgg16_4.keras",
    "Eksperimen 5 (LR 1e-4 + Custom Brightness/Contrast)": "model_vgg16_5.keras",
    "Eksperimen 6 (LR 1e-1 + Custom Brightness/Contrast)": "model_vgg16_6.keras"
}

# Variabel Global untuk Lazy Loading (Menghemat RAM HuggingFace)
active_model_name = None
active_model = None
active_grad_model = None

# Warna Brand
BRAND_BLUE = "#3F64A9"
BRAND_RED = "#A93F3F"
BRAND_WHITE = "#FFFFFF"

# ==========================================
# 2. CORE LOGIC (Model Switcher & Grad-CAM)
# ==========================================
def load_selected_model(selected_name):
    global active_model_name, active_model, active_grad_model
    
    # Jika model yang diminta sudah aktif, langsung gunakan (tidak perlu load ulang)
    if selected_name == active_model_name and active_model is not None:
        return active_model, active_grad_model
        
    print(f"[INFO] Beralih ke model: {selected_name}...")
    
    # Bersihkan memori dari model sebelumnya
    if active_model is not None:
        del active_model
        del active_grad_model
        gc.collect()
        tf.keras.backend.clear_session()
        
    file_name = MODEL_DICT[selected_name]
    
    if not os.path.exists(file_name):
        raise FileNotFoundError(f"Model {file_name} tidak ditemukan. Pastikan file sudah terunggah.")
        
    active_model = load_model(file_name)
    active_grad_model = Model(
        inputs=active_model.inputs,
        outputs=[active_model.get_layer(last_conv_layer_name).output, active_model.output]
    )
    active_model_name = selected_name
    
    return active_model, active_grad_model

def make_gradcam_heatmap(img_array, grad_model):
    with tf.GradientTape() as tape:
        last_conv_layer_output, preds = grad_model(img_array)
        pred_index = tf.argmax(preds[0])
        class_channel = preds[:, pred_index]
        grads = tape.gradient(class_channel, last_conv_layer_output)
        
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
    last_conv_layer_output = last_conv_layer_output[0]
    
    heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
    heatmap = tf.squeeze(heatmap)
    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
    
    return heatmap.numpy()

def create_superimposed_image(original_img, heatmap, alpha=0.4):
    if np.max(heatmap) == 0:
        return original_img
        
    heatmap = np.uint8(255 * heatmap)
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
    heatmap = cv2.resize(heatmap, (original_img.shape[1], original_img.shape[0]))
    
    superimposed_img = cv2.addWeighted(heatmap, alpha, original_img, 1 - alpha, 0)
    return superimposed_img

def predict_and_explain(image, selected_model_name):
    if image is None:
        return "<div style='color:red; padding:10px;'>Mohon unggah citra terlebih dahulu.</div>", None
        
    try:
        current_model, current_grad_model = load_selected_model(selected_model_name)
    except Exception as e:
        return f"<div style='color:red; padding:10px;'>Error: {str(e)}</div>", None

    # Preprocessing
    img_pil = Image.fromarray(image).resize((224, 224))
    img_array = np.array(img_pil)
    original_img_visual = img_array.copy()
    
    img_array = img_array.astype('float32') / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    # Prediction
    preds = current_model.predict(img_array)
    score = preds[0]
    predicted_class_index = int(np.argmax(score))
    predicted_label = class_names[predicted_class_index]
    confidence = float(np.max(score))

    # Grad-CAM
    heatmap = make_gradcam_heatmap(img_array, current_grad_model)
    gradcam_result = create_superimposed_image(original_img_visual, heatmap)

    # Output Formatting
    if predicted_label == "Infiltrat Pneumonia":
        header_color = BRAND_RED 
        status_icon = "⚠️"
        analysis_text = "AI mendeteksi pola <b>opasitas paru</b> yang mengindikasikan keberadaan Infiltrat Pneumonia. Segera lakukan peninjauan klinis lebih lanjut."
    else: 
        header_color = BRAND_BLUE 
        status_icon = "✅"
        analysis_text = "Paru-paru tampak bersih. Tidak ditemukan indikasi visual opasitas atau kelainan paru."

    result_md = f"""
    <div style="background-color: {header_color}; padding: 20px; border-radius: 12px; text-align: center; margin-bottom: 15px; box-shadow: 0 4px 12px rgba(0,0,0,0.15);">
        <h2 style="margin:0; font-size: 32px; color: #FFFFFF !important; font-weight: 700; font-family: 'SF Pro Display', 'Inter', sans-serif;">{status_icon} {predicted_label}</h2>
        <p style="margin:8px 0 0 0; font-size: 18px; color: #FFFFFF !important; opacity: 0.95; font-family: 'SF Pro Display', 'Inter', sans-serif;">Tingkat Kepercayaan: {confidence:.2%}</p>
        <p style="margin:4px 0 0 0; font-size: 14px; color: #FFFFFF !important; opacity: 0.8;">Menggunakan {selected_model_name}</p>
    </div>
    
    <div style="padding: 15px; border: 1px solid #E5E7EB; border-radius: 12px; background-color: white;">
        <p style="font-weight: 600; color: #374151; margin-bottom: 8px;">Analisis Klinis AI:</p>
        <p style="color: #4B5563; margin-bottom: 10px; line-height: 1.5;">{analysis_text}</p>
        <p style="font-size: 12px; color: #9CA3AF;">*Perhatikan area berwarna merah/hangat pada gambar Heatmap di samping sebagai fokus deteksi utama.*</p>
    </div>
    """
    return result_md, gradcam_result

# ==========================================
# 3. UI/UX & STYLING
# ==========================================
aidia_css = f"""
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&display=swap');
* {{
    font-family: 'SF Pro Display', 'Inter', -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif !important;
}}
.gradio-container {{
    max-width: 100% !important;
    padding: 0 40px !important;
    margin: 0 !important;
    background-color: #F8F9FA;
}}
.header-container {{
    text-align: center;
    padding: 40px 20px;
    background-color: {BRAND_WHITE};
    border-bottom: 4px solid {BRAND_BLUE};
    margin-bottom: 30px;
    border-radius: 0 0 20px 20px;
    box-shadow: 0 4px 20px rgba(0,0,0,0.05);
}}
.brand-title {{
    color: {BRAND_BLUE};
    font-weight: 800;
    font-size: 4rem; 
    letter-spacing: -1.5px;
    margin-bottom: 10px;
    line-height: 1.1;
}}
.brand-tagline {{
    color: #6B7280;
    font-weight: 400;
    font-size: 1.25rem;
    letter-spacing: 0.2px;
}}
.image-box {{
    border: 2px solid #E5E7EB;
    border-radius: 16px;
    overflow: hidden;
    background-color: white;
    box-shadow: 0 4px 6px rgba(0,0,0,0.05);
    transition: transform 0.2s;
}}
button.primary {{
    background-color: {BRAND_BLUE} !important;
    color: white !important;
    border: none !important;
    font-weight: 600;
    font-size: 1.1rem;
    padding: 12px 24px;
    border-radius: 8px;
    transition: all 0.3s ease;
}}
button.primary:hover {{
    background-color: #2c4a80 !important;
    box-shadow: 0 8px 20px rgba(63, 100, 169, 0.3);
    transform: translateY(-2px);
}}
.footer-text {{
    text-align: center;
    color: #9CA3AF;
    font-size: 0.9rem;
    margin-top: 50px;
    padding: 30px;
    border-top: 1px solid #E5E7EB;
}}
"""

theme = gr.themes.Soft(
    primary_hue="blue",
    neutral_hue="slate",
    radius_size="lg",
    font=['SF Pro Display', 'Inter', 'sans-serif']
).set(
    button_primary_background_fill=BRAND_BLUE,
    button_primary_text_color="white",
    block_title_text_color=BRAND_BLUE
)

# ==========================================
# 4. BUILDING THE APP
# ==========================================
with gr.Blocks(theme=theme, css=aidia_css, title="PneumoScan") as demo:
    
    # --- HEADER ---
    with gr.Row():
        gr.HTML(f"""
        <div class="header-container">
            <h1 class="brand-title">PneumoScan</h1>
            <p class="brand-tagline">Empowering Diagnostics. Leading Medicine. Building the Future.</p>
        </div>
        """)
        
    # --- MAIN CONTENT ---
    with gr.Row():
        
        # KOLOM INPUT (KIRI)
        with gr.Column(scale=5):
            gr.Markdown("### ⚙️ Konfigurasi Model")
            model_dropdown = gr.Dropdown(
                choices=list(MODEL_DICT.keys()), 
                value="Eksperimen 4 (LR 1e-4)", # Default Model
                label="Pilih Arsitektur Model",
                interactive=True
            )
            
            gr.Markdown("### 📥 Upload Citra Medis")
            input_image = gr.Image(
                label="Chest X-Ray Input", 
                type="numpy", 
                height=450, 
                elem_classes="image-box"
            )
            
            with gr.Row():
                clear_btn = gr.ClearButton(components=[input_image], value="Reset", size="sm")
                submit_btn = gr.Button("Analisis Diagnosis", variant="primary", size="lg")
            
            gr.Markdown("""
            <div style="background-color: white; padding: 20px; border-radius: 12px; border: 1px solid #F3F4F6; margin-top: 20px;">
                <strong style="color: #374151;">Petunjuk Penggunaan:</strong>
                <ul style="margin-top: 10px; padding-left: 20px; color: #6B7280; font-size: 0.95rem;">
                    <li>Pilih skenario model yang ingin diuji dari menu <i>dropdown</i>.</li>
                    <li>Upload foto Rontgen Dada (X-Ray) berformat JPG, PNG, atau JPEG.</li>
                    <li>Klik <b>Analisis Diagnosis</b> (Penggantian model pertama kali memakan waktu ±3 detik).</li>
                </ul>
            </div>
            """)

        # KOLOM OUTPUT (KANAN)
        with gr.Column(scale=6):
            gr.Markdown("### 🩺 Hasil Diagnosis & Visualisasi")
            output_text = gr.HTML(label="Diagnostic Report")
            output_image = gr.Image(
                label="Explainable AI (Grad-CAM)", 
                height=550,
                elem_classes="image-box", 
                show_label=True,
                show_download_button=True
            )

    # --- FOOTER ---
    gr.HTML("""
    <div class="footer-text">
        <p style="font-weight: 600; color: #4B5563;">Created by Muhammad Aqil</p>
        <p>PneumoScan by AIDIA Hub © 2026. Intelligent Clarity for Better Healthcare.</p>
        <p style="font-size: 12px; margin-top:10px; opacity: 0.8;">Disclaimer: This tool is for educational and research purposes only. Always consult a certified radiologist.</p>
    </div>
    """)

    # --- INTERACTION ---
    submit_btn.click(
        fn=predict_and_explain, 
        inputs=[input_image, model_dropdown], 
        outputs=[output_text, output_image]
    )

# Launch
demo.launch()