import gradio as gr
import torch
import torch.nn as nn
from torchvision import transforms, models
from PIL import Image
import numpy as np

# === 使用 grad-cam (套件名稱已修正) ===
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image

# 設定設備
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ==================== 1. 載入肺炎模型 (ResNet50) ====================
# 使用 weights=None 取代已棄用的 pretrained=False
model_pneumonia = models.resnet50(weights=None)
num_ftrs = model_pneumonia.fc.in_features
model_pneumonia.fc = nn.Linear(num_ftrs, 2)  # 兩類：NORMAL vs PNEUMONIA

model_path_pneumonia = "best_pneumonia_model.pth"

# 載入模型權重
try:
    # 建議使用 weights_only=False 以確保舊版模型權重能正確讀取
    model_pneumonia.load_state_dict(torch.load(model_path_pneumonia, map_location=device, weights_only=False))
    print(f"成功載入模型權重：{model_path_pneumonia}")
except Exception as e:
    print(f"模型載入失敗，請確認檔案是否存在：{e}")

model_pneumonia.to(device)
model_pneumonia.eval()

# 設定 Grad-CAM 目標層 (ResNet50 的最後一層捲積層)
target_layers = [model_pneumonia.layer4[-1]]

# ==================== 2. 影像前處理 ====================
img_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225]),
])

# ==================== 3. 肺炎預測函數 ====================
def predict_pneumonia(image):
    if image is None:
        return "請上傳影像", None, "無影像"
    
    # 轉換圖片格式
    pil_img = Image.fromarray(image).convert("RGB")
    input_tensor = img_transform(pil_img).unsqueeze(0).to(device)
    
    # 模型推論
    with torch.no_grad():
        output = model_pneumonia(input_tensor)
        probs = torch.softmax(output, dim=1).cpu().numpy()[0]
        pred_class = np.argmax(probs)
        confidence = probs[pred_class]
    
    label = "PNEUMONIA" if pred_class == 1 else "NORMAL"
    result_text = f"肺炎檢測結果：{label} (信心度：{confidence:.1%})"
    
    # --- Grad-CAM 視覺化 ---
    # 注意：GradCAM 在計算時需要梯度，所以不放在 no_grad 中
    cam = GradCAM(model=model_pneumonia, target_layers=target_layers)
    targets = [ClassifierOutputTarget(pred_class)]
    
    # 產生熱圖
    grayscale_cam = cam(input_tensor=input_tensor, targets=targets)[0]
    
    # 將原始圖片縮放到 224x224 以符合熱圖尺寸
    input_float_img = np.array(pil_img.resize((224, 224))).astype(np.float32) / 255.0
    visualization = show_cam_on_image(input_float_img, grayscale_cam, use_rgb=True)
    
    # --- 風險評估建議 ---
    if label == "PNEUMONIA" and confidence > 0.7:
        risk_assessment = "🚨 **高風險**：高度懷疑肺炎，建議立即就醫並進行進一步檢查。"
    elif label == "PNEUMONIA":
        risk_assessment = "⚠️ **中等風險**：疑似肺炎徵象，建議儘快諮詢醫療專業人員。"
    else:
        risk_assessment = "✅ **正常**：目前影像未顯示明顯肺炎特徵，若有呼吸道症狀仍請留意。"
    
    return result_text, visualization, risk_assessment

# ==================== 4. Gradio 介面 ====================
with gr.Blocks(title="肺炎檢測 AI 輔助系統") as demo:
    gr.Markdown("# 🏥 肺炎檢測 AI 輔助系統")
    gr.Markdown("這是一個基於深度學習的胸部 X 光影像分析工具。請上傳一張 X 光片，模型將分析是否存在肺炎徵兆，並顯示模型關注的區域。")
    
    with gr.Row():
        with gr.Column(scale=1):
            img_input = gr.Image(label="1. 上傳胸部 X 光影像", type="numpy")
            btn = gr.Button("🔍 開始檢測", variant="primary")
        
        with gr.Column(scale=1):
            pneumonia_out = gr.Textbox(label="檢測結果")
            cam_out = gr.Image(label="模型關注區域（紅色代表重點關注部位）")
            risk_out = gr.Markdown(label="專業建議與風險評估")
    
    # 設定按鈕邏輯
    btn.click(
        fn=predict_pneumonia,
        inputs=[img_input],
        outputs=[pneumonia_out, cam_out, risk_out]
    )
    
    gr.Markdown("---")
    gr.Markdown("⚠️ **免責聲明**：本工具僅作為研究及技術展示用途，**不可**取代專業醫師的醫療診斷。如果您感到身體不適，請務必尋求正式醫療協助。")

# ==================== 5. 啟動服務 ====================
if __name__ == "__main__":
    # 部署在 Hugging Face 時 share=True 是可選的，建議維持預設
    demo.launch()