import gradio as gr
import torchvision.transforms as transforms
from torchvision.transforms import InterpolationMode
import torch
from huggingface_hub import hf_hub_download

from .model import Model

# Load Model
model_path = hf_hub_download(
    repo_id="itserr/exvoto_classifier_convnext_base_224",
    filename="model.pt"
)

model = Model('convnext_base')
ckpt = torch.load(model_path, map_location=torch.device("cpu"))  # Ensure compatibility
model.load_state_dict(ckpt['model'])

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)
model.eval()


# Image Transformations
transform = transforms.Compose([
    transforms.Resize(size=(224,224), interpolation=InterpolationMode.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Classification Function
def classify_img(img, threshold):
    classification_threshold = threshold
    img_tensor = transform(img).unsqueeze(0).to(device)

    with torch.no_grad():
        pred = model(img_tensor)
        score = torch.sigmoid(pred).item()
        
    # Determine Prediction
    if score >= classification_threshold:
        label = "✅ This is an **Ex-Voto** image!"
    else:
        label = "❌ This is **NOT** an Ex-Voto image."
    
    # Format Confidence Score
    confidence = f"The probability that the image is an ex-voto is: {score:.2%}"
    return label, confidence


# # **🎨 Customized Interface**
demo = gr.Interface(
    fn=classify_img,
    inputs=[
        gr.Image(type="pil"),
        gr.Slider(minimum=0.5, maximum=1.0, value=0.7, step=0.1, label="Classification Threshold")
    ],
    outputs=[
        gr.Textbox(label="Prediction", interactive=False),
        gr.Textbox(label="Confidence Score", interactive=False),
    ],
    title="🖼️✟ Ex-Voto Image Classifier",
    description="📸 **Upload an image** to check if it's an **Ex-Voto** painting!",
    theme="soft",
    allow_flagging="never",
    live=False,  # Avoids auto-updating; requires a button click
)

# Launch App
demo.launch()