app.py

import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
import cv2
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use("Agg")
from PIL import Image
import requests
import json
import fitz
import pytesseract
import re
import os
import pickle
import pandas as pd
from sklearn.preprocessing import LabelEncoder
from sklearn.impute import SimpleImputer
from huggingface_hub import hf_hub_download
import joblib
import shap

# ── Device ────────────────────────────────────────────────────────────────────
device = torch.device("cpu")
REPO_ID = "Dina-Raslan/ckd-retinal-model"

# ── UCI Column Definitions ────────────────────────────────────────────────────
NUMERIC_COLS     = ["age", "bp", "bgr", "bu", "sc", "sod", "pot", "hemo", "pcv", "wbcc", "rbcc"]
CATEGORICAL_COLS = ["sg", "al", "su", "rbc", "pc", "pcc", "ba", "htn", "dm", "cad", "appet", "pe", "ane"]
ALL_COLS         = NUMERIC_COLS + CATEGORICAL_COLS

# القيم الممكنة لكل عمود categorical - لازم تكون ثابتة علشان LabelEncoder يشتغل صح
LABEL_ENCODER_CLASSES = {
    "sg":    ["1.005", "1.010", "1.015", "1.020", "1.025"],
    "al":    ["0", "1", "2", "3", "4", "5"],
    "su":    ["0", "1", "2", "3", "4", "5"],
    "rbc":   ["abnormal", "normal"],
    "pc":    ["abnormal", "normal"],
    "pcc":   ["notpresent", "present"],
    "ba":    ["notpresent", "present"],
    "htn":   ["no", "yes"],
    "dm":    ["no", "yes"],
    "cad":   ["no", "yes"],
    "appet": ["good", "poor"],
    "pe":    ["no", "yes"],
    "ane":   ["no", "yes"],
}

def encode_categoricals(df):
    """Encode categorical columns consistently using fixed classes"""
    for col in CATEGORICAL_COLS:
        le = LabelEncoder()
        le.classes_ = np.array(LABEL_ENCODER_CLASSES[col])
        val = str(df[col].iloc[0]).strip()
        # لو القيمة مش موجودة في الـ classes خد أقرب قيمة
        if val not in LABEL_ENCODER_CLASSES[col]:
            val = LABEL_ENCODER_CLASSES[col][0]
        df[col] = le.transform([val])
    return df

# ── Retinal Model ─────────────────────────────────────────────────────────────
class RetinalModelV2(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        self.backbone = timm.create_model(
            "efficientnet_b3", pretrained=False,
            num_classes=0, global_pool="avg", drop_rate=0.3
        )
        self.feature_dim = self.backbone.num_features
        self.classifier = nn.Sequential(
            nn.Linear(self.feature_dim, 512), nn.BatchNorm1d(512),
            nn.ReLU(), nn.Dropout(0.5),
            nn.Linear(512, 256), nn.BatchNorm1d(256),
            nn.ReLU(), nn.Dropout(0.4),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        features = self.backbone(x)
        return self.classifier(features), features

# ── Fusion Model ──────────────────────────────────────────────────────────────
class CrossModalAttentionFusion(nn.Module):
    def __init__(self, retinal_dim=64, clinical_dim=24, hidden_dim=128, num_classes=2):
        super().__init__()
        self.retinal_proj = nn.Sequential(
            nn.Linear(retinal_dim, hidden_dim), nn.LayerNorm(hidden_dim),
            nn.ReLU(), nn.Dropout(0.3)
        )
        self.clinical_proj = nn.Sequential(
            nn.Linear(clinical_dim, hidden_dim), nn.LayerNorm(hidden_dim),
            nn.ReLU(), nn.Dropout(0.3)
        )
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=hidden_dim, num_heads=4, dropout=0.1, batch_first=True
        )
        self.self_attention = nn.MultiheadAttention(
            embed_dim=hidden_dim, num_heads=4, dropout=0.1, batch_first=True
        )
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim * 2, 128), nn.BatchNorm1d(128),
            nn.ReLU(), nn.Dropout(0.4),
            nn.Linear(128, 64), nn.ReLU(),
            nn.Dropout(0.3), nn.Linear(64, num_classes)
        )

    def forward(self, retinal, clinical):
        r = self.retinal_proj(retinal).unsqueeze(1)
        c = self.clinical_proj(clinical).unsqueeze(1)
        r_attended, _ = self.cross_attention(query=r, key=c, value=c)
        c_attended, _ = self.self_attention(query=c, key=c, value=c)
        fused = torch.cat([r_attended.squeeze(1), c_attended.squeeze(1)], dim=1)
        return self.classifier(fused)

# ── Load All Models ───────────────────────────────────────────────────────────
print("Loading models...")

retinal_model_path = hf_hub_download(repo_id=REPO_ID, filename="best_model.pth")
retinal_model = RetinalModelV2(num_classes=2)
retinal_model.load_state_dict(torch.load(retinal_model_path, map_location="cpu"), strict=False)
retinal_model.eval()
print("Retinal model loaded")

fusion_model_path = hf_hub_download(repo_id=REPO_ID, filename="best_fusion_model.pth")
fusion_model = CrossModalAttentionFusion(retinal_dim=64, clinical_dim=24, hidden_dim=128, num_classes=2)
fusion_model.load_state_dict(torch.load(fusion_model_path, map_location="cpu"))
fusion_model.eval()
print("Fusion model loaded")

pca_path = hf_hub_download(repo_id=REPO_ID, filename="pca_model.pkl")
with open(pca_path, "rb") as f:
    pca_model = pickle.load(f)
print("PCA model loaded")

scaler_path = hf_hub_download(repo_id=REPO_ID, filename="clinical_scaler.pkl")
with open(scaler_path, "rb") as f:
    clinical_scaler = pickle.load(f)
print("Clinical scaler loaded")

gb_path = hf_hub_download(repo_id=REPO_ID, filename="gb_clinical_model_joblib.pkl")
gb_model = joblib.load(gb_path)
print("GB clinical model loaded")

# SHAP explainer for the clinical model (created once, reused for every request)
shap_explainer = shap.TreeExplainer(gb_model)
print("SHAP explainer ready")

print("All models ready.")

# تحقق من اسماء الأعمدة في الـ scaler
if hasattr(clinical_scaler, 'feature_names_in_'):
    print("Scaler trained on:", list(clinical_scaler.feature_names_in_))
else:
    print("Scaler has no feature_names_in_ - will use numpy array directly")

OPTIMAL_THRESHOLD = 0.50

# ── Preprocessing ─────────────────────────────────────────────────────────────
def remove_black_border(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        x, y, w, h = cv2.boundingRect(max(contours, key=cv2.contourArea))
        img = img[y:y+h, x:x+w]
    return img

def lab_ace_enhancement(img):
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    l = clahe.apply(l)
    return cv2.cvtColor(cv2.merge([l, a, b]), cv2.COLOR_LAB2BGR)

def graham_preprocessing(img):
    return cv2.addWeighted(img, 4, cv2.GaussianBlur(img, (0, 0), 10), -4, 128)

def green_channel_enhancement(img):
    b, g, r = cv2.split(img)
    g = cv2.normalize(g, None, 0, 255, cv2.NORM_MINMAX)
    return cv2.merge([b, g, r])

def preprocess_retinal(img_pil, target_size=512):
    img = cv2.cvtColor(np.array(img_pil.convert("RGB")), cv2.COLOR_RGB2BGR)
    img = remove_black_border(img)
    img = cv2.resize(img, (target_size, target_size))
    img = lab_ace_enhancement(img)
    img = graham_preprocessing(img)
    img = green_channel_enhancement(img)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return img.astype(np.float32) / 255.0

# ── Clinical Feature Preparation ──────────────────────────────────────────────
def prepare_clinical_features(
    # من التحاليل
    sc, bu, hemo, bgr, sod, pot, pcv, wbcc, rbcc,
    sg, al, su, rbc, pc, pcc, ba,
    # مانيوال
    age, bp, htn, dm, cad, appet, pe, ane
):
    """
    تحضير الـ 24 feature بالترتيب الصح للـ scaler
    """
    row = {
        # Numeric
        "age":  float(age),
        "bp":   float(bp),
        "bgr":  float(bgr),
        "bu":   float(bu),
        "sc":   float(sc),
        "sod":  float(sod),
        "pot":  float(pot),
        "hemo": float(hemo),
        "pcv":  float(pcv),
        "wbcc": float(wbcc),
        "rbcc": float(rbcc),
        # Categorical
        "sg":    str(sg),
        "al":    str(al),
        "su":    str(su),
        "rbc":   str(rbc),
        "pc":    str(pc),
        "pcc":   str(pcc),
        "ba":    str(ba),
        "htn":   str(htn),
        "dm":    str(dm),
        "cad":   str(cad),
        "appet": str(appet),
        "pe":    str(pe),
        "ane":   str(ane),
    }

    df = pd.DataFrame([row])

    # Encode categoricals بطريقة صح (fixed classes)
    df = encode_categoricals(df)

    # رتب الأعمدة بالترتيب الصح
    df = df[ALL_COLS]

    # اعمل numpy array مباشرة علشان نتجنب مشكلة feature names
    X = df.values.astype(np.float64)

    # Scale
    scaled = clinical_scaler.transform(X)
    return scaled

def get_top_shap_features(clinical_feat, top_n=6):
    """
    Compute SHAP values for the clinical (GB) model's prediction and
    return the top_n features ranked by absolute impact, together with
    each feature's percentage share of total |SHAP| impact and the
    combined percentage share of all remaining features.
    clinical_feat must be the scaled numpy array in ALL_COLS order.

    Returns:
        top_features: list of (feature_name, shap_value, pct_impact) tuples
        other_pct: combined percentage impact of the remaining features
    """
    shap_values = shap_explainer.shap_values(clinical_feat)

    # Some SHAP/model versions return a list per class, others a single array
    if isinstance(shap_values, list):
        shap_row = shap_values[1][0]  # class "CKD" (positive class)
    else:
        shap_row = shap_values[0]

    abs_vals   = np.abs(shap_row)
    total_abs  = abs_vals.sum()

    # Avoid division by zero in the (unlikely) case all SHAP values are 0
    if total_abs == 0:
        pct_vals = np.zeros_like(abs_vals)
    else:
        pct_vals = (abs_vals / total_abs) * 100.0

    sorted_idx  = np.argsort(abs_vals)[::-1]
    top_idx     = sorted_idx[:top_n]
    other_idx   = sorted_idx[top_n:]

    top_features = [
        (ALL_COLS[i], float(shap_row[i]), float(pct_vals[i]))
        for i in top_idx
    ]
    other_pct = float(pct_vals[other_idx].sum()) if len(other_idx) else 0.0

    return top_features, other_pct

# ── Grad-CAM++ ────────────────────────────────────────────────────────────────
class GradCAMPlusPlus:
    def __init__(self, model, target_layer):
        self.model       = model
        self.gradients   = None
        self.activations = None
        target_layer.register_forward_hook(self._save_activation)
        target_layer.register_full_backward_hook(self._save_gradient)

    def _save_activation(self, module, input, output):
        self.activations = output.detach()

    def _save_gradient(self, module, grad_input, grad_output):
        self.gradients = grad_output[0].detach()

    def generate(self, img_tensor, class_idx=None):
        self.model.eval()
        img_tensor = img_tensor.unsqueeze(0)
        logits, _  = self.model(img_tensor)
        if class_idx is None:
            class_idx = logits.argmax(dim=1).item()
        self.model.zero_grad()
        logits[0, class_idx].backward()
        grads     = self.gradients
        acts      = self.activations
        grads_sq  = grads ** 2
        grads_cub = grads ** 3
        denom     = 2 * grads_sq + acts * grads_cub
        denom     = torch.where(denom != 0, denom, torch.ones_like(denom))
        alpha     = grads_sq / denom
        weights   = (alpha * F.relu(grads)).sum(dim=[2, 3], keepdim=True)
        cam       = (weights * acts).sum(dim=1, keepdim=True)
        cam       = F.relu(cam)
        cam       = F.interpolate(cam, size=(512, 512), mode="bilinear", align_corners=False)
        cam       = cam.squeeze().cpu().numpy()
        cam      -= cam.min()
        if cam.max() > 0:
            cam /= cam.max()
        return cam

target_layer = retinal_model.backbone.blocks[6][1].conv_pwl
gradcam      = GradCAMPlusPlus(retinal_model, target_layer)

def make_retinal_mask(size=512):
    mask = np.zeros((size, size), dtype=np.float32)
    cv2.circle(mask, (size//2, size//2), int(size * 0.47), 1.0, -1)
    return mask

# ── OCR ───────────────────────────────────────────────────────────────────────
LAB_PATTERNS = {
    # Blood Tests
    "sc":   [r"creatinine[\s\S]{0,40}?(\d+\.?\d*)\s*mg", r"s\.?\s*creat[\s\S]{0,30}?(\d+\.?\d*)"],
    "bu":   [r"(?:blood\s*urea|urea|bun)[\s\S]{0,30}?(\d+\.?\d*)\s*mg", r"b\.?u\.?n[\s\S]{0,20}?(\d+\.?\d*)"],
    "hemo": [r"h[ae]moglobin[\s\S]{0,30}?(\d+\.?\d*)\s*g", r"hgb[\s\S]{0,20}?(\d+\.?\d*)"],
    "bgr":  [r"(?:blood\s*glucose|glucose|random\s*blood\s*sugar|rbs)[\s\S]{0,30}?(\d+\.?\d*)\s*mg",
             r"fbs[\s\S]{0,20}?(\d+\.?\d*)"],
    "sod":  [r"sodium[\s\S]{0,30}?(\d+\.?\d*)\s*m?eq", r"na\+?[\s\S]{0,20}?(\d+\.?\d*)"],
    "pot":  [r"potassium[\s\S]{0,30}?(\d+\.?\d*)\s*m?eq", r"k\+?[\s\S]{0,20}?(\d+\.?\d*)"],
    "pcv":  [r"(?:packed\s*cell\s*volume|hematocrit|pcv|hct)[\s\S]{0,30}?(\d+\.?\d*)\s*%?"],
    "wbcc": [r"(?:white\s*blood\s*cell|wbc|leukocyte)[\s\S]{0,30}?(\d+[\d,]*\.?\d*)",
             r"wbcc[\s\S]{0,20}?(\d+[\d,]*\.?\d*)"],
    "rbcc": [r"(?:red\s*blood\s*cell|rbc\s*count|erythrocyte)[\s\S]{0,30}?(\d+\.?\d*)",
             r"rbcc[\s\S]{0,20}?(\d+\.?\d*)"],
    # Urine Tests  
    "sg":   [r"(?:specific\s*gravity|sp\.?\s*gr\.?)[\s\S]{0,30}?(1\.\d{3})"],
    "al":   [r"(?:albumin|protein)[\s\S]{0,30}?(\d)\s*\+?", r"albumin[\s\S]{0,20}?(\d)"],
    "su":   [r"(?:sugar|glucose)\s*(?:in\s*urine|urine)[\s\S]{0,30}?(\d)", r"glycosuria[\s\S]{0,20}?(\d)"],
    "bp":   [r"(?:blood\s*pressure|bp|diastolic)[\s\S]{0,30}?(\d{2,3})\s*mm",
             r"\d{2,3}\s*/\s*(\d{2,3})\s*mm"],
}

VALUE_RANGES = {
    "sc":   (0.1, 20),    "bu":   (1, 300),
    "hemo": (1, 25),      "bgr":  (30, 600),
    "sod":  (100, 170),   "pot":  (1, 10),
    "pcv":  (10, 60),     "wbcc": (2000, 30000),
    "rbcc": (1, 8),       "sg":   (1.000, 1.030),
    "al":   (0, 5),       "su":   (0, 5),
    "bp":   (50, 250),
}

def parse_lab_values(text):
    text_lower = text.lower()
    values = {}

    for key, patterns in LAB_PATTERNS.items():
        for pattern in patterns:
            match = re.search(pattern, text_lower)
            if match:
                try:
                    val_str = match.group(1).replace(",", "")
                    val = float(val_str)
                    lo, hi = VALUE_RANGES.get(key, (0, 9999))
                    if lo <= val <= hi:
                        values[key] = val
                        break
                except (ValueError, IndexError):
                    pass

    # استنتاج categorical من النص
    # RBC in urine
    if re.search(r"rbc\s*(?:in\s*urine|urine)[\s\S]{0,30}?abnormal", text_lower):
        values["rbc"] = "abnormal"
    elif re.search(r"rbc\s*(?:in\s*urine|urine)[\s\S]{0,30}?normal", text_lower):
        values["rbc"] = "normal"

    # Pus Cells
    if re.search(r"pus\s*cells?[\s\S]{0,30}?abnormal", text_lower):
        values["pc"] = "abnormal"
    elif re.search(r"pus\s*cells?[\s\S]{0,30}?normal", text_lower):
        values["pc"] = "normal"

    # Pus Cell Clumps
    if re.search(r"pus\s*cell\s*clumps?[\s\S]{0,30}?present", text_lower):
        values["pcc"] = "present"
    elif re.search(r"pus\s*cell\s*clumps?[\s\S]{0,30}?not\s*present", text_lower):
        values["pcc"] = "notpresent"

    # Bacteria
    if re.search(r"bacteri[\s\S]{0,30}?present", text_lower):
        values["ba"] = "present"
    elif re.search(r"bacteri[\s\S]{0,30}?not\s*present", text_lower):
        values["ba"] = "notpresent"

    # Specific Gravity - snap to nearest valid value
    if "sg" in values:
        valid_sg = [1.005, 1.010, 1.015, 1.020, 1.025]
        sg_val = values["sg"]
        nearest = min(valid_sg, key=lambda x: abs(x - sg_val))
        values["sg"] = f"{nearest:.3f}"

    # Albumin و Sugar -> snap to integer string
    for col in ["al", "su"]:
        if col in values:
            values[col] = str(int(round(values[col])))

    return values

def process_lab_file(lab_file):
    if lab_file is None:
        return ("No file uploaded.", 
                gr.update(), gr.update(), gr.update(), gr.update(), gr.update(),
                gr.update(), gr.update(), gr.update(), gr.update(),
                gr.update(), gr.update(), gr.update())

    path = lab_file.name if hasattr(lab_file, "name") else lab_file

    if path.lower().endswith(".pdf"):
        doc  = fitz.open(path)
        text = "".join(page.get_text() for page in doc)
    else:
        img    = Image.open(path)
        img_np = np.array(img.convert("RGB"))
        img_cv = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
        img_cv = cv2.resize(img_cv, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)
        gray   = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
        gray   = cv2.fastNlMeansDenoising(gray, h=10)
        _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        text   = pytesseract.image_to_string(thresh, config="--oem 3 --psm 6")

    values = parse_lab_values(text)

    if values:
        found = [k for k in values]
        status = f" Extracted {len(found)} values: {', '.join(found)}\nPlease verify and fill missing fields manually."
    else:
        status = "️ Could not extract values automatically.\nPlease enter all values manually."

    def upd(key, default):
        return gr.update(value=values[key]) if key in values else gr.update(value=default)

    def upd_cat(key, default):
        return gr.update(value=values[key]) if key in values else gr.update(value=default)

    return (
        status,
        upd("sc",   1.2),
        upd("bu",   30),
        upd("hemo", 13),
        upd("bgr",  100),
        upd("sod",  135),
        upd("pot",  4.0),
        upd("pcv",  39),
        upd("wbcc", 7500),
        upd("rbcc", 4.5),
        upd("bp",   80),
        upd_cat("sg",  "1.020"),
        upd_cat("al",  "0"),
        upd_cat("su",  "0"),
        upd_cat("rbc", "normal"),
        upd_cat("pc",  "normal"),
        upd_cat("pcc", "notpresent"),
        upd_cat("ba",  "notpresent"),
    )

# ── Main Prediction ───────────────────────────────────────────────────────────
def predict_ckd_risk(
    retinal_image,
    # Blood tests
    sc, bu, hemo, bgr, sod, pot, pcv, wbcc, rbcc,
    # Blood pressure
    bp,
    # Urine tests (categorical)
    sg, al, su, rbc, pc, pcc, ba,
    # Manual entry
    age, htn, dm, cad, appet, pe, ane
):
    if retinal_image is None:
        return None, "️ Please upload a retinal image.", "", ""

    # 1. Preprocess retinal image
    img_processed = preprocess_retinal(retinal_image)
    img_tensor    = torch.tensor(img_processed).permute(2, 0, 1)

    # 2. Extract retinal features
    with torch.no_grad():
        logits, retinal_features = retinal_model(img_tensor.unsqueeze(0))
        retinal_prob = torch.softmax(logits, dim=1)[0, 1].item()

    # 3. PCA on retinal features
    retinal_feat_np  = retinal_features.cpu().numpy()
    retinal_feat_pca = pca_model.transform(retinal_feat_np)

    # 4. Prepare clinical features (24 columns, correct order)
    clinical_feat = prepare_clinical_features(
        sc=sc, bu=bu, hemo=hemo, bgr=bgr, sod=sod, pot=pot,
        pcv=pcv, wbcc=wbcc, rbcc=rbcc,
        sg=sg, al=al, su=su, rbc=rbc, pc=pc, pcc=pcc, ba=ba,
        age=age, bp=bp,
        htn=htn, dm=dm, cad=cad, appet=appet, pe=pe, ane=ane
    )

    # 5. GB clinical prediction
    clinical_prob = gb_model.predict_proba(clinical_feat)[0, 1]

    # 5.b Top-6 SHAP features driving the clinical prediction (with % impact)
    top_shap_features, other_shap_pct = get_top_shap_features(clinical_feat, top_n=6)

    # 6. Fusion model prediction
    retinal_tensor  = torch.tensor(retinal_feat_pca, dtype=torch.float32)
    clinical_tensor = torch.tensor(clinical_feat,    dtype=torch.float32)

    with torch.no_grad():
        fusion_logits = fusion_model(retinal_tensor, clinical_tensor)
        fusion_prob   = torch.softmax(fusion_logits, dim=1)[0, 1].item()

    # 7. Final risk
    if fusion_prob >= 0.6:
        risk_label = " HIGH RISK"
        risk_color = "red"
    elif fusion_prob >= 0.35:
        risk_label = " MEDIUM RISK"
        risk_color = "orange"
    else:
        risk_label = " LOW RISK"
        risk_color = "green"

    # 8. Grad-CAM++
    cam          = gradcam.generate(img_tensor)
    retinal_mask = make_retinal_mask(512)
    cam          = cam * retinal_mask
    cam         -= cam.min()
    if cam.max() > 0:
        cam /= cam.max()

    img_rgb   = (img_processed * 255).astype(np.uint8)
    heatmap   = cv2.applyColorMap((cam * 255).astype(np.uint8), cv2.COLORMAP_JET)
    heatmap   = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
    heat_mask = (cam > 0.3).astype(np.float32)
    heat_mask = np.stack([heat_mask] * 3, axis=-1)
    overlay   = (img_rgb * (1 - 0.6 * heat_mask) + heatmap * 0.6 * heat_mask).astype(np.uint8)
    circ_mask = np.stack([retinal_mask] * 3, axis=-1)
    overlay   = (overlay * circ_mask).astype(np.uint8)

    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    axes[0].imshow(img_rgb);         axes[0].set_title("Preprocessed Retinal Image"); axes[0].axis("off")
    axes[1].imshow(cam, cmap="jet"); axes[1].set_title("Grad-CAM++ Heatmap");         axes[1].axis("off")
    axes[2].imshow(overlay);         axes[2].set_title("Overlay");                    axes[2].axis("off")
    plt.suptitle(
        f"CKD Risk: {risk_label}  |  Fusion Score: {fusion_prob:.3f}",
        fontsize=13, color=risk_color, fontweight="bold"
    )
    plt.tight_layout()
    gradcam_path = "/tmp/gradcam_result.png"
    plt.savefig(gradcam_path, dpi=120)
    plt.close()

    clinical_summary = f"""
╔══════════════════════════════════════╗
║         CLINICAL VALUES              ║
╠══════════════════════════════════════╣
  BLOOD TESTS
  Creatinine (sc):      {sc} mg/dL
  Urea/BUN (bu):        {bu} mg/dL
  Hemoglobin (hemo):    {hemo} g/dL
  Blood Glucose (bgr):  {bgr} mg/dL
  Sodium (sod):         {sod} mEq/L
  Potassium (pot):      {pot} mEq/L
  Packed Cell Vol (pcv):{pcv} %
  WBC Count (wbcc):     {wbcc} cells/cumm
  RBC Count (rbcc):     {rbcc} millions/cumm
  Blood Pressure (bp):  {bp} mmHg
  URINE TESTS
  Specific Gravity (sg):{sg}
  Albumin (al):         {al}
  Sugar (su):           {su}
  RBC in Urine (rbc):   {rbc}
  Pus Cells (pc):       {pc}
  Pus Cell Clumps (pcc):{pcc}
  Bacteria (ba):        {ba}
  PATIENT INFO
  Age:                  {age} years
  Hypertension (htn):   {htn}
  Diabetes (dm):        {dm}
  Coronary Art. (cad):  {cad}
  Appetite (appet):     {appet}
  Pedal Edema (pe):     {pe}
  Anemia (ane):         {ane}
╠══════════════════════════════════════╣
  AI PREDICTION BREAKDOWN
  Retinal Score:        {retinal_prob:.4f}
  Clinical Score:       {clinical_prob:.4f}
  Fusion Score:         {fusion_prob:.4f}
  Risk Level:           {risk_label}
╠══════════════════════════════════════╣
  TOP 6 CLINICAL FEATURES (SHAP)
{chr(10).join(f"  {i+1}. {name:<8s} {'+' if val >= 0 else ''}{val:.4f}  ({pct:.1f}% impact)" for i, (name, val, pct) in enumerate(top_shap_features))}
  Other 18 factors combined:  {other_shap_pct:.1f}% impact
╚══════════════════════════════════════╝
"""

    risk_display = f"## {risk_label}\n\n**Fusion Score:** {fusion_prob:.4f}\n\n**Retinal:** {retinal_prob:.4f} | **Clinical:** {clinical_prob:.4f}"

    return gradcam_path, risk_display, clinical_summary, f"Done. Risk: {risk_label}"

# ── Chatbot ───────────────────────────────────────────────────────────────────
GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")

SYSTEM_PROMPT = """You are a CKD Health Coach, a specialized AI assistant for patients
with Chronic Kidney Disease (CKD) or those at risk.
- Provide lifestyle and dietary recommendations specific to CKD
- Offer fluid intake and exercise advice
- Answer basic kidney health questions
- NEVER diagnose or replace a doctor
- For severe symptoms, recommend emergency care immediately"""

EMERGENCY_KEYWORDS = ["chest pain", "difficulty breathing", "severe swelling",
                      "cannot breathe", "unconscious", "bleeding", "fainted", "heart attack"]

def check_emergency(message):
    return any(kw in message.lower() for kw in EMERGENCY_KEYWORDS)

def chat_with_groq(message, history):
    if check_emergency(message):
        return " This sounds like a medical emergency. Please call emergency services immediately."

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for item in history:
        if isinstance(item, dict):
            messages.append({"role": item["role"], "content": item["content"]})
        elif isinstance(item, (list, tuple)) and len(item) == 2:
            messages.append({"role": "user",      "content": item[0]})
            messages.append({"role": "assistant", "content": item[1]})
    messages.append({"role": "user", "content": message})

    if not GROQ_API_KEY:
        return "Chatbot unavailable: GROQ_API_KEY not set in Space secrets."

    try:
        response = requests.post(
            "https://api.groq.com/openai/v1/chat/completions",
            headers={"Authorization": f"Bearer {GROQ_API_KEY}", "Content-Type": "application/json"},
            json={"model": "llama-3.3-70b-versatile", "messages": messages,
                  "max_tokens": 512, "temperature": 0.7},
            timeout=30
        )
        return response.json()["choices"][0]["message"]["content"]
    except Exception as e:
        return f"Error: {str(e)}"

# ── Gradio UI ─────────────────────────────────────────────────────────────────
with gr.Blocks(title="CKD-MultiVision") as demo:
    gr.Markdown("""
    #  CKD-MultiVision
    ### Multimodal AI System for Early Chronic Kidney Disease Detection
    *Retinal Image + Clinical Lab Values → AI Risk Assessment*
    """)

    with gr.Tabs():

        # ── Tab 1: Risk Prediction ─────────────────────────────────────────
        with gr.TabItem(" Risk Prediction"):

            # Lab file upload
            with gr.Row():
                with gr.Column():
                    gr.Markdown("###  Upload Lab Report (Optional)")
                    lab_file_input = gr.File(
                        label="Upload Lab Report (PDF or Image)",
                        file_types=[".pdf", ".png", ".jpg", ".jpeg"]
                    )
                    extract_btn    = gr.Button(" Extract Values from File", variant="secondary")
                    extract_status = gr.Textbox(label="Extraction Status", lines=3, interactive=False)

            gr.Markdown("---")

            with gr.Row():
                # Left: Retinal Image
                with gr.Column(scale=1):
                    gr.Markdown("### ️ Retinal Fundus Image")
                    retinal_input = gr.Image(
                        label="Upload Retinal Image (JPG/PNG)",
                        type="pil", height=300
                    )

                # Middle: Blood Tests
                with gr.Column(scale=1):
                    gr.Markdown("###  Blood Tests")
                    sc_input   = gr.Number(label="Serum Creatinine - sc (mg/dL)",        value=1.2,  minimum=0.1, maximum=20)
                    bu_input   = gr.Number(label="Blood Urea/BUN - bu (mg/dL)",          value=30,   minimum=1,   maximum=300)
                    hemo_input = gr.Number(label="Hemoglobin - hemo (g/dL)",             value=13,   minimum=1,   maximum=25)
                    bgr_input  = gr.Number(label="Blood Glucose Random - bgr (mg/dL)",   value=100,  minimum=30,  maximum=600)
                    sod_input  = gr.Number(label="Sodium - sod (mEq/L)",                 value=135,  minimum=100, maximum=170)
                    pot_input  = gr.Number(label="Potassium - pot (mEq/L)",              value=4.0,  minimum=1,   maximum=10)
                    pcv_input  = gr.Number(label="Packed Cell Volume/Hematocrit - pcv (%)", value=39, minimum=10, maximum=60)
                    wbcc_input = gr.Number(label="WBC Count - wbcc (cells/cumm)",        value=7500, minimum=2000, maximum=30000)
                    rbcc_input = gr.Number(label="RBC Count - rbcc (millions/cumm)",     value=4.5,  minimum=1,   maximum=8)
                    bp_input   = gr.Number(label="Blood Pressure Diastolic - bp (mmHg)", value=80,   minimum=50,  maximum=200)

                # Right: Urine Tests + Manual
                with gr.Column(scale=1):
                    gr.Markdown("###  Urine Tests")
                    sg_input  = gr.Dropdown(label="Specific Gravity - sg",
                                            choices=["1.005","1.010","1.015","1.020","1.025"],
                                            value="1.020")
                    al_input  = gr.Dropdown(label="Albumin in Urine - al (0=None, 5=Heavy)",
                                            choices=["0","1","2","3","4","5"], value="0")
                    su_input  = gr.Dropdown(label="Sugar in Urine - su (0=None, 5=Heavy)",
                                            choices=["0","1","2","3","4","5"], value="0")
                    rbc_input = gr.Radio(label="RBC in Urine - rbc",
                                         choices=["normal","abnormal"], value="normal")
                    pc_input  = gr.Radio(label="Pus Cells - pc",
                                         choices=["normal","abnormal"], value="normal")
                    pcc_input = gr.Radio(label="Pus Cell Clumps - pcc",
                                         choices=["notpresent","present"], value="notpresent")
                    ba_input  = gr.Radio(label="Bacteria in Urine - ba",
                                         choices=["notpresent","present"], value="notpresent")

                    gr.Markdown("###  Patient Info (Manual)")
                    age_input   = gr.Number(label="Age (years)",                  value=45, minimum=2, maximum=90)
                    htn_input   = gr.Radio(label="Hypertension - htn",            choices=["no","yes"], value="no")
                    dm_input    = gr.Radio(label="Diabetes Mellitus - dm",         choices=["no","yes"], value="no")
                    cad_input   = gr.Radio(label="Coronary Artery Disease - cad", choices=["no","yes"], value="no")
                    appet_input = gr.Radio(label="Appetite - appet",              choices=["good","poor"], value="good")
                    pe_input    = gr.Radio(label="Pedal Edema (Swelling) - pe",   choices=["no","yes"], value="no")
                    ane_input   = gr.Radio(label="Anemia - ane",                  choices=["no","yes"], value="no")

            predict_btn = gr.Button(" Run CKD Risk Assessment", variant="primary", size="lg")

            with gr.Row():
                with gr.Column():
                    gradcam_output  = gr.Image(label="Retinal Analysis + Grad-CAM++")
                with gr.Column():
                    risk_output     = gr.Markdown(label="Risk Result")
                    clinical_output = gr.Textbox(label="Full Report", lines=20, interactive=False)
                    status_output   = gr.Textbox(label="Status", interactive=False)

        # ── Tab 2: CKD Health Coach ────────────────────────────────────────
        with gr.TabItem(" CKD Health Coach"):
            gr.Markdown("### Ask about diet, lifestyle, exercise, and kidney health.")
            chatbot = gr.Chatbot(height=450)
            chat_input = gr.Textbox(placeholder="Ask a question about CKD...", label="Your Message")
            send_btn   = gr.Button("Send", variant="primary")
            gr.Examples(
                examples=["What foods should I avoid with CKD?",
                          "How much water should I drink daily?",
                          "Can I exercise with kidney disease?",
                          "What are warning signs I should watch for?"],
                inputs=chat_input
            )

        # ── Tab 3: About ───────────────────────────────────────────────────
        with gr.TabItem("About"):
            gr.Markdown("""
            ## About CKD-MultiVision
            **Input Fields (24 total — matching UCI CKD Dataset):**
            | # | Field | Source | Description |
            |---|-------|--------|-------------|
            | 1 | age | Manual | Patient age in years |
            | 2 | bp | Lab | Blood pressure diastolic (mmHg) |
            | 3 | bgr | Lab | Blood glucose random (mg/dL) |
            | 4 | bu | Lab | Blood urea (mg/dL) |
            | 5 | sc | Lab | Serum creatinine (mg/dL) |
            | 6 | sod | Lab | Sodium (mEq/L) |
            | 7 | pot | Lab | Potassium (mEq/L) |
            | 8 | hemo | Lab | Hemoglobin (g/dL) |
            | 9 | pcv | Lab | Packed cell volume (%) |
            | 10 | wbcc | Lab | WBC count (cells/cumm) |
            | 11 | rbcc | Lab | RBC count (millions/cumm) |
            | 12 | sg | Lab | Specific gravity |
            | 13 | al | Lab | Albumin in urine (0-5) |
            | 14 | su | Lab | Sugar in urine (0-5) |
            | 15 | rbc | Lab | RBC in urine |
            | 16 | pc | Lab | Pus cells |
            | 17 | pcc | Lab | Pus cell clumps |
            | 18 | ba | Lab | Bacteria |
            | 19 | htn | Manual | Hypertension |
            | 20 | dm | Manual | Diabetes mellitus |
            | 21 | cad | Manual | Coronary artery disease |
            | 22 | appet | Manual | Appetite |
            | 23 | pe | Manual | Pedal edema |
            | 24 | ane | Manual | Anemia |
            **Pipeline:**
            - Retinal Image → EfficientNet-B3 → Feature Vector (1536) → PCA (64 dim)
            - Clinical Values (24) → StandardScaler → Gradient Boosting + Fusion Model
            - Cross-Modal Attention Fusion → Final Risk Level
            *Research prototype — not for clinical use.*
            """)

    # ── Event Handlers ─────────────────────────────────────────────────────
    extract_btn.click(
        fn=process_lab_file,
        inputs=[lab_file_input],
        outputs=[
            extract_status,
            sc_input, bu_input, hemo_input, bgr_input, sod_input,
            pot_input, pcv_input, wbcc_input, rbcc_input, bp_input,
            sg_input, al_input, su_input, rbc_input, pc_input, pcc_input, ba_input,
        ]
    )

    predict_btn.click(
        fn=predict_ckd_risk,
        inputs=[
            retinal_input,
            # Blood tests
            sc_input, bu_input, hemo_input, bgr_input,
            sod_input, pot_input, pcv_input, wbcc_input, rbcc_input,
            bp_input,
            # Urine tests
            sg_input, al_input, su_input, rbc_input, pc_input, pcc_input, ba_input,
            # Manual
            age_input, htn_input, dm_input, cad_input, appet_input, pe_input, ane_input,
        ],
        outputs=[gradcam_output, risk_output, clinical_output, status_output]
    )

    def respond(message, history):
        if not message.strip():
            return history, ""
        bot_response = chat_with_groq(message, history)
        history.append({"role": "user",      "content": message})
        history.append({"role": "assistant", "content": bot_response})
        return history, ""

    send_btn.click(fn=respond, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input])
    chat_input.submit(fn=respond, inputs=[chat_input, chatbot], outputs=[chatbot, chat_input])

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())