import sys
import os

os.environ.setdefault("HF_HUB_DISABLE_XET", "1")

try:
    import gradio as gr
    from PIL import Image
    import re
    import torch
    from transformers import pipeline
except ModuleNotFoundError as exc:
    package_map = {
        "gradio": "gradio",
        "transformers": "transformers",
        "PIL": "Pillow",
        "torch": "torch",
    }
    missing = exc.name or "unknown"
    suggested = package_map.get(missing, missing)
    print(f"[error] Required dependency '{missing}' is missing.")
    print("[error] Install requirements first:")
    print("  python3 -m pip install -r requirements.txt")
    print(f"  python3 -m pip install {suggested}")
    print("Then rerun: python3 app.py")
    sys.exit(1)

caption_model_name = "ydshieh/vit-gpt2-coco-en"
captioner = None
stop_words = {
    "a",
    "an",
    "and",
    "as",
    "at",
    "before",
    "but",
    "by",
    "for",
    "from",
    "how",
    "in",
    "into",
    "of",
    "on",
    "or",
    "that",
    "the",
    "this",
    "to",
    "with",
    "without",
    "is",
    "are",
    "was",
    "were",
    "be",
    "it",
    "its",
    "itself",
    "they",
    "their",
    "there",
}

def get_captioner():
    global captioner

    if captioner is None:
        print("Loading caption model. This can take a moment...", flush=True)
        captioner = pipeline(
            "image-to-text",
            model=caption_model_name,
            dtype=torch.float32,
            device=-1,
        )
        print("Caption model loaded.", flush=True)

    return captioner


def extract_labels_from_caption(caption, max_labels=6):
    tokens = re.findall(r"[a-zA-Z][a-zA-Z']+", caption.lower())
    tokens = [token.strip("'") for token in tokens]
    cleaned = [token for token in tokens if len(token) > 2 and token not in stop_words]
    if not cleaned:
        return {"other": 1.0}

    ordered_unique = []
    seen = set()
    for token in cleaned:
        if token in seen:
            continue
        seen.add(token)
        ordered_unique.append(token)

    top_terms = ordered_unique[:max_labels]
    raw_scores = [1.0 / (idx + 1) for idx in range(len(top_terms))]
    total = sum(raw_scores)
    return {term: score / total for term, score in zip(top_terms, raw_scores)}


def predict_from_image(image):
    if image is None:
        return "Please upload an image", {"Error": 1.0}

    try:
        caption = get_captioner()(image)[0]["generated_text"]
        labels = extract_labels_from_caption(caption)
        return caption, labels
    except Exception as exc:
        error_message = f"{exc.__class__.__name__}: {exc}"
        print(f"[error] Prediction failed: {error_message}", flush=True)
        return error_message, {"Error": 1.0}


# Gradio interface

demo = gr.Interface(
    fn=predict_from_image,
    inputs=gr.Image(label="Upload Image", type="pil", sources=["upload"]),
    outputs=[
        gr.Textbox(label="Caption", lines=2),
        gr.Label(label="Generated Labels"),
    ],
    title="Image Classifier",
    flagging_mode="never"
)

# Launch the interface
print("Starting Gradio app at http://127.0.0.1:7860", flush=True)
demo.launch(show_error=True)