"""
Eye Screen backend — deploy on Hugging Face Spaces (Gradio SDK).
Pipeline: image quality gate -> pretrained DR classifier -> Grad-CAM -> calibrated confidence.
This is the "real model" half of the app. Google AI Studio frontend calls this via its API.
"""

import spaces
import cv2
import numpy as np
import torch
import torch.nn.functional as F
import timm
import gradio as gr
from PIL import Image
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image

LABELS = ["No DR", "Mild", "Moderate", "Severe", "Proliferative DR"]
TEMPERATURE = 1.8  # softmax temperature for calibrated confidence; tune on a val split if time allows

# --- Model ---
# Pretrained on APTOS/EyePACS/IDRID/DDR fundus data (timm hub format).
# If this exact repo id doesn't load cleanly, swap for another fundus-pretrained
# timm/HF checkpoint (e.g. search "diabetic retinopathy" on huggingface.co/models)
# and keep everything else (quality gate, Grad-CAM, calibration) unchanged.
MODEL_ID = "hf_hub:ClementP/FundusDRGrading-efficientnet_b0"
model = timm.create_model(MODEL_ID, pretrained=True)
model.eval()

# last conv block for Grad-CAM (EfficientNet: conv_head is the final conv layer)
target_layers = [model.conv_head] if hasattr(model, "conv_head") else [list(model.children())[-3]]
cam = GradCAM(model=model, target_layers=target_layers)

IMG_SIZE = 224


def check_quality(img_bgr):
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    blur_score = cv2.Laplacian(gray, cv2.CV_64F).var()
    brightness = gray.mean()

    issues = []
    if blur_score < 100:
        issues.append("Image appears blurry/out of focus")
    if brightness < 40:
        issues.append("Image is underexposed (too dark)")
    if brightness > 210:
        issues.append("Image is overexposed (too bright)")

    return issues, {"blur_score": round(float(blur_score), 1), "brightness": round(float(brightness), 1)}


def preprocess(img_pil):
    img = img_pil.resize((IMG_SIZE, IMG_SIZE))
    arr = np.array(img).astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    arr = (arr - mean) / std
    tensor = torch.tensor(arr.transpose(2, 0, 1), dtype=torch.float32).unsqueeze(0)
    return tensor, arr


@spaces.GPU(duration=30)
def predict(image):
    if image is None:
        return "No image uploaded", "", None, {}

    img_pil = Image.fromarray(image).convert("RGB")
    img_bgr = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

    issues, metrics = check_quality(img_bgr)
    if issues:
        msg = "RETAKE REQUESTED: " + "; ".join(issues)
        return msg, "N/A — quality check failed before classification", None, metrics

    tensor, norm_img = preprocess(img_pil)

    with torch.no_grad():
        logits = model(tensor)
        calibrated_probs = F.softmax(logits / TEMPERATURE, dim=1)[0]

    pred_idx = int(torch.argmax(calibrated_probs))
    confidence = float(calibrated_probs[pred_idx]) * 100
    severity = LABELS[pred_idx]

    # Grad-CAM heatmap
    grayscale_cam = cam(input_tensor=tensor, targets=None)[0]
    display_img = np.array(img_pil.resize((IMG_SIZE, IMG_SIZE))).astype(np.float32) / 255.0
    overlay = show_cam_on_image(display_img, grayscale_cam, use_rgb=True)

    result_text = f"{severity}  ({confidence:.1f}% confidence)"
    disclaimer = "Screening aid only — not a diagnosis. Consult an ophthalmologist for confirmation."
    return result_text, disclaimer, overlay, metrics


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Fundus photo"),
    outputs=[
        gr.Textbox(label="Result"),
        gr.Textbox(label="Note"),
        gr.Image(label="Grad-CAM heatmap"),
        gr.JSON(label="Quality metrics (debug)"),
    ],
    title="Eye Screen — DR Grading API",
    description="Upload a retinal fundus image. Rejects low-quality photos before grading.",
)

if __name__ == "__main__":
    demo.launch()