import os
import numpy as np
import cv2
import gradio as gr
import onnxruntime as ort
from huggingface_hub import hf_hub_download

# ==============================================================================
# 1. PENGUNDUHAN & INISIALISASI MODEL ONNX FROM HUB
# ==============================================================================
REPO_ID = "ASomeoneWhoInterestedWithAI/LookThemV7_Caltech256-ONNX"

print("📡 Mengunduh sasis LookThem V7-W dari Hugging Face Hub...")
onnx_path = hf_hub_download(repo_id=REPO_ID, filename="LookThemV7_Caltech256.onnx")
onnx_data_path = hf_hub_download(repo_id=REPO_ID, filename="LookThemV7_Caltech256.onnx.data")

print("🧠 Membakar Graf ke ONNX Runtime Session...")
session = ort.InferenceSession(onnx_path)
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# ==============================================================================
# 2. STRUKTUR MAP RESMI 257 LABEL SEMANTIK CALTECH-256
# ==============================================================================
# ==============================================================================
# 2. STRUKTUR MAP SEMANTIK CALTECH-256 (KALIBRASI MUTLAK DATASET ASLI)
# ==============================================================================
CLASSES = ['ak47', 'american-flag', 'backpack', 'baseball-bat', 'baseball-glove', 'basketball-hoop', 'bat', 'bathtub', 'bear', 'beer-mug', 'billiards', 'binoculars', 
           'birdbath', 'blimp', 'bonsai-101', 'boom-box', 'bowling-ball', 'bowling-pin', 'boxing-glove', 'brain-101', 'breadmaker', 'buddha-101', 'bulldozer', 'butterfly', 
           'cactus', 'cake', 'calculator', 'camel', 'cannon', 'canoe', 'car-tire', 'cartman', 'cd', 'centipede', 'cereal-box', 'chandelier-101', 'chess-board', 'chimp', 
           'chopsticks', 'cockroach', 'coffee-mug', 'coffin', 'coin', 'comet', 'computer-keyboard', 'computer-monitor', 'computer-mouse', 'conch', 'cormorant', 'covered-wagon', 
           'cowboy-hat', 'crab-101', 'desk-globe', 'diamond-ring', 'dice', 'dog', 'dolphin-101', 'doorknob', 'drinking-straw', 'duck', 'dumb-bell', 'eiffel-tower', 'electric-guitar-101', 
           'elephant-101', 'elk', 'ewer-101', 'eyeglasses', 'fern', 'fighter-jet', 'fire-extinguisher', 'fire-hydrant', 'fire-truck', 'fireworks', 'flashlight', 'floppy-disk', 
           'football-helmet', 'french-horn', 'fried-egg', 'frisbee', 'frog', 'frying-pan', 'galaxy', 'gas-pump', 'giraffe', 'goat', 'golden-gate-bridge', 'goldfish', 
           'golf-ball', 'goose', 'gorilla', 'grand-piano-101', 'grapes', 'grasshopper', 'guitar-pick', 'hamburger', 'hammock', 'harmonica', 'harp', 'harpsichord', 'hawksbill-101', 
           'head-phones', 'helicopter-101', 'hibiscus', 'homer-simpson', 'horse', 'horseshoe-crab', 'hot-air-balloon', 'hot-dog', 'hot-tub', 'hourglass', 'house-fly', 
           'human-skeleton', 'hummingbird', 'ibis-101', 'ice-cream-cone', 'iguana', 'ipod', 'iris', 'jesus-christ', 'joy-stick', 'kangaroo-101', 'kayak', 'ketch-101', 
           'killer-whale', 'knife', 'ladder', 'laptop-101', 'lathe', 'leopards-101', 'license-plate', 'lightbulb', 'light-house', 'lightning', 'llama-101', 'mailbox', 
           'mandolin', 'mars', 'mattress', 'megaphone', 'menorah-101', 'microscope', 'microwave', 'minaret', 'minotaur', 'motorbikes-101', 'mountain-bike', 'mushroom', 
           'mussels', 'necktie', 'octopus', 'ostrich', 'owl', 'palm-pilot', 'palm-tree', 'paperclip', 'paper-shredder', 'pci-card', 'penguin', 'people', 'pez-dispenser', 
           'photocopier', 'picnic-table', 'playing-card', 'porcupine', 'pram', 'praying-mantis', 'pyramid', 'raccoon', 'radio-telescope', 'rainbow', 'refrigerator', 
           'revolver-101', 'rifle', 'rotary-phone', 'roulette-wheel', 'saddle', 'saturn', 'school-bus', 'scorpion-101', 'screwdriver', 'segway', 'self-propelled-lawn-mower', 
           'sextant', 'sheet-music', 'skateboard', 'skunk', 'skyscraper', 'smokestack', 'snail', 'snake', 'sneaker', 'snowmobile', 'soccer-ball', 'socks', 'soda-can', 'spaghetti', 
           'speed-boat', 'spider', 'spoon', 'stained-glass', 'starfish-101', 'steering-wheel', 'stirrups', 'sunflower-101', 'superman', 'sushi', 'swan', 'swiss-army-knife', 
           'sword', 'syringe', 'tambourine', 'teapot', 'teddy-bear', 'teepee', 'telephone-box', 'tennis-ball', 'tennis-court', 'tennis-racket', 'theodolite', 'toaster', 
           'tomato', 'tombstone', 'top-hat', 'touring-bike', 'tower-pisa', 'traffic-light', 'treadmill', 'triceratops', 'tricycle', 'trilobite-101', 'tripod', 't-shirt', 
           'tuning-fork', 'tweezer', 'umbrella-101', 'unicorn', 'vcr', 'video-projector', 'washing-machine', 'watch-101', 'waterfall', 'watermelon', 'welding-mask', 
           'wheelbarrow', 'windmill', 'wine-bottle', 'xylophone', 'yarmulke', 'yo-yo', 'zebra', 'airplanes-101', 'car-side-101', 'faces-easy-101', 'greyhound', 
           'tennis-shoes', 'toad', 'clutter']

# ==============================================================================
# 3. FUNGSI PIPELINE INFERENCE (PREPROCESSING & RUN)
# ==============================================================================
def predict_image(img):
    if img is None:
        return "Silakan masukkan gambar terlebih dahulu."
        
    # --- PREPROCESSING (Sinkron dengan Transformasi PyTorch V7-W) ---
    if len(img.shape) == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
    elif img.shape[2] == 4:
        img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
        
    img_resized = cv2.resize(img, (96, 96))
    img_normalized = img_resized.astype(np.float32) / 255.0
    
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    img_normalized = (img_normalized - mean) / std
    
    img_final = img_normalized.transpose(2, 0, 1)
    input_tensor = np.expand_dims(img_final, axis=0)
    
    # --- TEMBAK RUNTIME INFERENCE ---
    raw_outputs = session.run([output_name], {input_name: input_tensor})
    logits = raw_outputs[0][0]
    
    # --- PASCA-PROSES (SOFTMAX STABIL) ---
    exp_logits = np.exp(logits - np.max(logits))
    probabilities = exp_logits / np.sum(exp_logits)
    
    # Mengambil Top 5 prediksi tertinggi untuk dashboard Gradio
    top_5_indices = np.argsort(probabilities)[::-1][:5]
    
    results = {}
    for idx in top_5_indices:
        # Penanganan darurat jika indeks melompat di luar batas list
        if idx < len(CLASSES):
            class_name = CLASSES[idx]
        else:
            class_name = f"Unknown Class Idx-{idx}"
            
        confidence = float(probabilities[idx])
        results[class_name] = confidence
        
    return results

# ==============================================================================
# 4. ARSITEKTUR ANTARMUKA GRADIO (UI DESIGN)
# ==============================================================================
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🏎️ LookThem V7 (96x96) - Caltech-256 Demo
        ### Model Kustom Mini 0.17 MB (Graf) + External Data | Akurasi Pengujian Semantik 38.03% (*From Scratch*)!
        Modul ini berjalan menggunakan **ONNX Runtime Engine** berbasis pelacakan Dynamo. Semua dropout latihan telah dilepas!
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="numpy", label="Input Gambar Caltech-256")
            btn_submit = gr.Button("Analisis Karakteristik Gambar 🚀", variant="primary")
            
        with gr.Column():
            output_labels = gr.Label(num_top_classes=5, label="Top 5 Prediksi Probabilitas Semantik")
            
    btn_submit.click(fn=predict_image, inputs=input_img, outputs=output_labels)

if __name__ == "__main__":
    #demo.launch()
    # PASTIKAN DI BAGIAN PALING BAWAH APP.PY SEPERTI INI:
    demo.launch(theme=gr.themes.Soft())