import gradio as gr
import torch
import torch.nn.functional as F
import numpy as np
from torchvision import models, transforms
from huggingface_hub import hf_hub_download
from PIL import Image
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image

MODEL_REPO = "ummanmm/classroom-reaction-resnet18"
CLASSES = [
    "Bored / Tired",
    "Confused",
    "Neutral",
    "Smiling / Amused",
    "Surprised",
]
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

model_path = hf_hub_download(
    repo_id=MODEL_REPO, filename="best_resnet18.pth"
)

model = models.resnet18(weights=None)
model.fc = torch.nn.Linear(model.fc.in_features, 5)
model.load_state_dict(
    torch.load(model_path, map_location=torch.device("cpu"))
)
model.eval()

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
])


def predict_reaction(image):
    if image is None:
        return {}, None

    img = Image.fromarray(image).convert("RGB")
    tensor = transform(img).unsqueeze(0)

    with torch.no_grad():
        outputs = model(tensor)
        probs = F.softmax(outputs[0], dim=0)

    confidences = {CLASSES[i]: float(probs[i]) for i in range(5)}

    rgb = np.array(img.resize((224, 224))).astype(np.float32) / 255.0
    cam = GradCAM(model=model, target_layers=[model.layer4[-1]])
    grayscale = cam(input_tensor=tensor, targets=None)[0, :]
    heatmap = show_cam_on_image(rgb, grayscale, use_rgb=True)

    return confidences, heatmap


theme = gr.themes.Soft(primary_hue="blue")

demo = gr.Interface(
    fn=predict_reaction,
    inputs=gr.Image(label="Upload Student Crop"),
    outputs=[
        gr.Label(num_top_classes=5, label="Reaction Prediction"),
        gr.Image(label="Grad-CAM Attention Heatmap"),
    ],
    title="Classroom Reaction — ResNet18",
    description=(
        "Upload a cropped image of a student to classify their facial reaction using a ResNet18 model. The Grad-CAM heatmap shows which regions the CNN focuses on."
    ),
    examples=[],
    theme=theme,
)

demo.launch()