import gradio as gr
import torch
from torchvision import transforms
from PIL import Image
import json, random, os

# --- 1️⃣ Load models ---
model_paths = {
    "Classifier Only (M1)": "M1.pt",
    "Classifier + L1-4 (M3)": "M3.pt",
    "Classifier + L1-4 (M4)": "M4.pt"
}
models = {name: torch.jit.load(path, map_location='cpu').eval()
          for name, path in model_paths.items()}

# --- 2️⃣ Load label map ---
with open("label_map.json") as f:
    label_map = json.load(f)
classes = [label_map[k] for k in sorted(label_map, key=int)]

# --- 3️⃣ Load test data ---
TEST_DIR = "test"
TEST_FILE = "test.txt"
with open(TEST_FILE, "r") as f:
    test_entries = [line.strip().split() for line in f.readlines()]
test_data = [(os.path.join(TEST_DIR, fn), int(cls)) for fn, cls in test_entries]

# --- 4️⃣ Preprocessing ---
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# --- 5️⃣ Helper functions ---
def get_random_test_image():
    """Pick a random test image and its label."""
    img_path, cls_id = random.choice(test_data)
    true_label = label_map[str(cls_id)]
    img = Image.open(img_path).convert("RGB")
    return img, f"Ground Truth: {true_label}"

def classify_bird(img, model_name, mode):
    """Classify based on current input mode."""
    if img is None:
        return {}
    model = models[model_name]
    x = preprocess(img).unsqueeze(0)
    with torch.no_grad():
        logits = model(x)
        probs = torch.softmax(logits, dim=1)[0]
    top5_prob, top5_idx = torch.topk(probs, 5)
    return {classes[i]: float(top5_prob[j]) for j, i in enumerate(top5_idx)}

# --- 6️⃣ Interface layout ---
with gr.Blocks(theme=gr.themes.Soft(), title="CUB-200 Bird Classifier") as demo:
    gr.Markdown("""
    # 🐦 CUB-200 Bird Classifier
    Identify bird species using fine-tuned ResNet101 models.  
    Choose input mode and model, then classify!
    """)

    # --- Input controls ---
    with gr.Row():
        input_mode = gr.Radio(
            choices=["Upload Image", "Use Random Test Image"],
            value="Upload Image",
            label="Input Mode",
            interactive=True
        )
        model_choice = gr.Dropdown(
            choices=list(models.keys()),
            value="Classifier + L1-4 (M4)",
            label="Select Model"
        )

    gr.Markdown("---")

    # --- Main visualization area ---
    with gr.Row(equal_height=True):
        with gr.Column():
            image_display = gr.Image(type="pil", label="Bird Image", interactive=True)
            regenerate_btn = gr.Button("🔁 Regenerate Random Image", visible=False)
            ground_truth = gr.Textbox(label="Ground Truth", interactive=False, visible=False)
        predictions = gr.Label(num_top_classes=5, label="Top-5 Predictions")

    # --- Bottom control ---
    classify_btn = gr.Button("🔍 Classify")

    # --- Logic ---
    def handle_mode_change(mode):
        """Handle mode switch and optionally load a random image."""
        if mode == "Use Random Test Image":
            img, label = get_random_test_image()
            return (
                gr.update(value=img, interactive=False),
                gr.update(visible=True),
                gr.update(value=label, visible=True)
            )
        else:
            return (
                gr.update(value=None, interactive=True),
                gr.update(visible=False),
                gr.update(value="", visible=False)
            )

    def regenerate_random_image():
        img, label = get_random_test_image()
        return gr.update(value=img), gr.update(value=label)

    input_mode.change(
        fn=handle_mode_change,
        inputs=input_mode,
        outputs=[image_display, regenerate_btn, ground_truth]
    )

    regenerate_btn.click(
        fn=regenerate_random_image,
        inputs=None,
        outputs=[image_display, ground_truth]
    )

    classify_btn.click(
        fn=classify_bird,
        inputs=[image_display, model_choice, input_mode],
        outputs=predictions
    )

demo.launch()