import gradio as gr
import torch
import timm
import torchvision.transforms as T
from PIL import Image
import numpy as np
import pandas as pd
import json, os

# ------------------------------
# CONFIGURATION
# ------------------------------
DEVICE = "cpu"
NUM_CLASSES = 200

MODEL_PATHS = {
    "EfficientNet-B3": "models/BestEfficientNetB3_model.pth",
    "EfficientNetV2-S (in21k)": "models/BestEfficientNetV2S_model.pth"
}

MODEL_NAMES = {
    "EfficientNet-B3": "tf_efficientnet_b3_ns",
    "EfficientNetV2-S (in21k)": "tf_efficientnetv2_s.in21k"
}

# ------------------------------
# LOAD CLASS NAMES (supports JSON or TXT)
# ------------------------------
CLASS_NAMES = []

if os.path.exists("class_names.json"):
    try:
        with open("class_names.json", "r") as f:
            data = json.load(f)
            if isinstance(data, dict):
                # some JSONs store indices as keys
                CLASS_NAMES = [v for _, v in sorted(data.items(), key=lambda x: int(x[0]))]
            elif isinstance(data, list):
                CLASS_NAMES = data
        print(f"✅ Loaded {len(CLASS_NAMES)} class names from class_names.json")
    except Exception as e:
        print(f"⚠️ Error reading class_names.json: {e}")
elif os.path.exists("class_names.txt"):
    with open("class_names.txt", "r") as f:
        CLASS_NAMES = [line.strip() for line in f.readlines()]
    print(f"✅ Loaded {len(CLASS_NAMES)} class names from class_names.txt")
else:
    CLASS_NAMES = [f"Class {i}" for i in range(NUM_CLASSES)]
    print("⚠️ No class_names file found — using generic labels")

# ------------------------------
# MODEL LOADING + CACHING
# ------------------------------
@torch.inference_mode()
def load_model(model_name):
    model = timm.create_model(MODEL_NAMES[model_name],
                              pretrained=False,
                              num_classes=NUM_CLASSES)
    model.load_state_dict(torch.load(MODEL_PATHS[model_name],
                                     map_location=DEVICE))
    model.to(DEVICE).eval()
    return model

_loaded_models = {}
def get_model(model_name):
    if model_name not in _loaded_models:
        print(f"🔄 Loading {model_name} ...")
        _loaded_models[model_name] = load_model(model_name)
    return _loaded_models[model_name]

# ------------------------------
# PREDICTION FUNCTION
# ------------------------------
def predict(image, model_choice):
    if image is None:
        return pd.DataFrame(), "Please upload an image."

    model = get_model(model_choice)

    # Ensure RGB
    try:
        image = image.convert("RGB")
    except Exception as e:
        return pd.DataFrame(), f"❌ Error reading image: {e}"

    # Resize dynamically
    size = 300 if "b3" in model_choice.lower() else 384
    preprocess = T.Compose([
        T.Resize((size, size)),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225])
    ])
    img_tensor = preprocess(image).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        logits = model(img_tensor)
        probs = torch.nn.functional.softmax(logits, dim=1)[0]

    top5 = torch.topk(probs, 5)
    indices = top5.indices.cpu().numpy().tolist()
    scores = top5.values.cpu().numpy().tolist()

    # Map to readable names
    if CLASS_NAMES and len(CLASS_NAMES) > max(indices):
        names = [CLASS_NAMES[i] for i in indices]
    else:
        names = [f"Class {i}" for i in indices]

    # Make a DataFrame (this avoids 'undefined' issue)
    df = pd.DataFrame({
        "Class": names,
        "Confidence (%)": [round(s * 100, 2) for s in scores]
    })

    return df, f"✅ Prediction complete using {model_choice}"

# ------------------------------
# GRADIO INTERFACE
# ------------------------------
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown(
        """
        #Bird Species Classifier  
        Upload a bird image or use your camera to identify its species.  
        Supports multiple EfficientNet models and displays top-5 predictions.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            img_in = gr.Image(type="pil",
                              label="Upload or Capture Image",
                              sources=["upload", "webcam"])
            model_sel = gr.Dropdown(
                choices=list(MODEL_PATHS.keys()),
                value="EfficientNetV2-S",
                label="Select Model")
            btn = gr.Button("🔍 Classify")

        with gr.Column(scale=1):
            df_out = gr.Dataframe(label="Top-5 Predictions",
                                  interactive=False)
            status = gr.Textbox(label="Status", interactive=False)

    btn.click(fn=predict, inputs=[img_in, model_sel],
              outputs=[df_out, status])

    gr.Markdown(
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)