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-engagement-mobilenet"
CLASSES = [
    "Oriented (Lecture-Focused)",
    "Diverted (Looking Away/Down)",
    "Obscured (Blocked/Empty)",
]
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_baseline.pth"
)

model = models.mobilenet_v2(weights=None)
model.classifier[1] = torch.nn.Linear(model.last_channel, 3)
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_engagement(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(3)}

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

    return confidences, heatmap


demo = gr.Interface(
    fn=predict_engagement,
    inputs=gr.Image(label="Upload Student Crop"),
    outputs=[
        gr.Label(num_top_classes=3, label="Engagement Prediction"),
        gr.Image(label="Grad-CAM Attention Heatmap"),
    ],
    title="Real-Time Classroom Engagement Telemetry",
    description=(
        "Upload a cropped image of a student to classify their "
        "physical engagement state using a custom MobileNetV2 model. "
        "The Grad-CAM heatmap shows which regions the CNN focuses on."
    ),
    examples=[],
)

demo.launch()