"""
app.py - Gradio Demo for EfficientNet-B4 + CBAM (Diabetic Retinopathy Classification)
Deployed as a HuggingFace Space linked to chrisnguyenx/EfficientNet-P3
"""

import os
import json
import torch
import torch.nn as nn
import numpy as np
import cv2
from PIL import Image
from torchvision import models, transforms
from huggingface_hub import hf_hub_download
import gradio as gr

# ========================== MODEL DEFINITION ==========================

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc = nn.Sequential(
            nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False),
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc(self.avg_pool(x))
        max_out = self.fc(self.max_pool(x))
        return self.sigmoid(avg_out + max_out)


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        return self.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1)))


class CBAM(nn.Module):
    def __init__(self, in_planes, ratio=16, kernel_size=7):
        super().__init__()
        self.ca = ChannelAttention(in_planes, ratio)
        self.sa = SpatialAttention(kernel_size)

    def forward(self, x):
        x = x * self.ca(x)
        x = x * self.sa(x)
        return x


class EfficientNetB4_CBAM(nn.Module):
    def __init__(self, num_classes=5, drop_rate=0.3, cbam_ratio=16):
        super().__init__()
        backbone = models.efficientnet_b4(weights=None)
        self.features = backbone.features
        in_planes = 1792
        self.cbam = CBAM(in_planes, ratio=cbam_ratio)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            nn.Dropout(p=drop_rate),
            nn.Linear(in_planes, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.cbam(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)


# ========================== PREPROCESSING ==========================

CROP_TOLERANCE = 12
BEN_SIGMA = 10
TARGET_SIZE = (224, 224)
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def crop_fundus_circle(img, tolerance=CROP_TOLERANCE):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    _, mask = cv2.threshold(gray, tolerance, 255, cv2.THRESH_BINARY)
    coords = cv2.findNonZero(mask)
    if coords is None:
        return img
    x, y, w, h = cv2.boundingRect(coords)
    return img[y: y + h, x: x + w]


def auto_detect_border(img, thresh=0.05):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if img.ndim == 3 else img
    return float((gray < CROP_TOLERANCE).mean()) > thresh


def letterbox_resize(img, target_size=TARGET_SIZE):
    h, w = img.shape[:2]
    th, tw = target_size
    scale = min(tw / w, th / h)
    nw, nh = int(w * scale), int(h * scale)
    resized = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_CUBIC)
    canvas = np.zeros((th, tw, 3), dtype=np.uint8)
    pad_y = (th - nh) // 2
    pad_x = (tw - nw) // 2
    canvas[pad_y: pad_y + nh, pad_x: pad_x + nw] = resized
    return canvas


def ben_graham_transform(img, sigma_x=BEN_SIGMA):
    blur = cv2.GaussianBlur(img, (0, 0), sigmaX=sigma_x)
    enhanced = cv2.addWeighted(img, 4, blur, -4, 128)
    return np.clip(enhanced, 0, 255).astype(np.uint8)


def preprocess_image(pil_image):
    """Full pipeline: PIL Image -> preprocessed BGR numpy array."""
    img_rgb = np.array(pil_image.convert("RGB"))
    img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)

    if auto_detect_border(img_bgr):
        img_bgr = crop_fundus_circle(img_bgr)

    img_bgr = letterbox_resize(img_bgr, TARGET_SIZE)
    img_bgr = ben_graham_transform(img_bgr)
    return img_bgr


def to_tensor(img_bgr):
    """BGR numpy -> normalized PyTorch tensor."""
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(img_rgb)
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
    ])
    return transform(pil_img)


# ========================== LOAD MODEL ==========================

REPO_ID = "chrisnguyenx/EfficientNet-P3"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

CLASS_NAMES = {
    0: "No DR",
    1: "Mild",
    2: "Moderate",
    3: "Severe",
    4: "Proliferative DR",
}

CLINICAL_ADVICE = {
    0: {
        "title": "Mat Binh Thuong (No DR)",
        "emoji": "✅",
        "color": "#10b981",
        "advice": "Chua phat hien ton thuong vong mac tieu duong. Khuyen nghi kham mat dinh ky 12 thang/lan.",
        "urgency": "Binh thuong",
    },
    1: {
        "title": "Benh Nhe (Mild DR)",
        "emoji": "🔵",
        "color": "#3b82f6",
        "advice": "Xuat hien vi phinh mach nho. Khuyen nghi tai kham sau 6-12 thang va kiem soat duong huyet.",
        "urgency": "Theo doi dinh ky",
    },
    2: {
        "title": "Benh Trung Binh (Moderate DR)",
        "emoji": "🟡",
        "color": "#f59e0b",
        "advice": "Ton thuong xuat huyet/xuat tiet muc do vua. Can kham bac si nhan khoa trong 3-6 thang.",
        "urgency": "Kham chuyen khoa",
    },
    3: {
        "title": "Benh Nang (Severe DR)",
        "emoji": "🔴",
        "color": "#ef4444",
        "advice": "Ton thuong nghiem trong o nhieu goc phan tu vong mac. CAN chuyen kham gap trong 2-4 tuan.",
        "urgency": "Can can thiep som",
    },
    4: {
        "title": "Tang Sinh Nguy Hiem (Proliferative DR)",
        "emoji": "🟣",
        "color": "#8b5cf6",
        "advice": "Tang sinh tan mach nguy co gay mo mat vinh vien! CAN DIEU TRI KHAN CAP.",
        "urgency": "KHAN CAP",
    },
}


def load_model():
    """Download weights from HF Hub and load model."""
    weights_path = hf_hub_download(
        repo_id=REPO_ID,
        filename="efficientnet_b4_cbam_fold1.pth",
    )

    model = EfficientNetB4_CBAM(num_classes=5, drop_rate=0.3, cbam_ratio=16)

    checkpoint = torch.load(weights_path, map_location=DEVICE, weights_only=False)
    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
        state_dict = checkpoint["model_state_dict"]
    else:
        state_dict = checkpoint

    model.load_state_dict(state_dict)
    model.to(DEVICE)
    model.eval()
    return model


print("Loading model from HuggingFace Hub...")
model = load_model()
print(f"Model loaded on {DEVICE}")


# ========================== INFERENCE ==========================

def predict(image):
    """Main prediction function for Gradio."""
    if image is None:
        return None, "Please upload a retinal fundus image.", ""

    # Preprocess
    img_bgr = preprocess_image(image)
    tensor = to_tensor(img_bgr).unsqueeze(0).to(DEVICE)

    # Inference
    with torch.no_grad():
        outputs = model(tensor)
        probs = torch.softmax(outputs, dim=1)[0].cpu().numpy()

    pred_class = int(np.argmax(probs))
    confidence = float(probs[pred_class]) * 100.0

    # Build results
    label_results = {CLASS_NAMES[i]: float(probs[i]) for i in range(5)}

    clinical = CLINICAL_ADVICE[pred_class]
    clinical_text = f"""
### {clinical['emoji']} {clinical['title']}

**Confidence:** {confidence:.1f}%
**Urgency:** {clinical['urgency']}

**Clinical Advice:**
{clinical['advice']}
"""

    # Return preprocessed image for visualization
    preprocessed_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    preprocessed_pil = Image.fromarray(preprocessed_rgb)

    return label_results, clinical_text, preprocessed_pil


# ========================== GRADIO UI ==========================

with gr.Blocks(
    title="DR Classification - EfficientNet-B4 + CBAM",
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="purple",
    ),
) as demo:
    gr.Markdown(
        """
        # 🔬 Diabetic Retinopathy Classification
        ### EfficientNet-B4 + CBAM Attention | 5-Class ICDR Standard

        Upload a **retinal fundus image** to classify the severity of Diabetic Retinopathy.

        | Class | Description |
        |-------|-------------|
        | **0 - No DR** | No visible retinopathy |
        | **1 - Mild** | Microaneurysms only |
        | **2 - Moderate** | More than just microaneurysms |
        | **3 - Severe** | Extensive hemorrhages |
        | **4 - Proliferative DR** | Neovascularization / vitreous hemorrhage |
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                type="pil",
                label="Upload Fundus Image",
                height=350,
            )
            predict_btn = gr.Button(
                "🔍 Analyze Image",
                variant="primary",
                size="lg",
            )

        with gr.Column(scale=1):
            output_label = gr.Label(
                label="Classification Probabilities",
                num_top_classes=5,
            )
            output_clinical = gr.Markdown(label="Clinical Guidance")

    with gr.Row():
        output_preprocessed = gr.Image(
            label="Preprocessed Image (Ben Graham Enhanced)",
            height=250,
        )

    predict_btn.click(
        fn=predict,
        inputs=input_image,
        outputs=[output_label, output_clinical, output_preprocessed],
    )

    gr.Examples(
        examples=[],
        inputs=input_image,
        label="Example Images (upload your own fundus images)",
    )

    gr.Markdown(
        """
        ---
        **Model:** [`chrisnguyenx/EfficientNet-P3`](https://huggingface.co/chrisnguyenx/EfficientNet-P3)
        | **Architecture:** EfficientNet-B4 + CBAM
        | **Metrics:** QWK=0.854, F1=0.722
        """
    )

demo.launch()