# Required imports for the Gradio interface
import os, random, joblib, io
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm.notebook import tqdm
from PIL import Image

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms, models

import optuna
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget

import gradio as gr
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import classification_report, confusion_matrix

# Define device, IMG_SIZE, and get_model function if not globally available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
IMG_SIZE = (224, 224) # Assuming this is the image size used for preprocessing

def get_model(model_name, pretrained=True):
    weights = 'DEFAULT' if pretrained else None
    if model_name == "resnet50":
        model = models.resnet50(weights=weights)
        model.fc = nn.Linear(model.fc.in_features, 2)
    elif model_name == "efficientnet_b0":
        model = models.efficientnet_b0(weights=weights)
        model.classifier[1] = nn.Linear(model.classifier[1].in_features, 2)
    elif model_name == "densenet121":
        model = models.densenet121(weights=weights)
        model.classifier = nn.Linear(model.classifier.in_features, 2)
    else:
        raise ValueError("Model name not recognized")
    return model

# Define the path where your model pipeline is saved
MODEL_PIPELINE_PATH = "/content/drive/MyDrive/pneumonia_pipeline_kfold.pth"

# 1. تحميل الأنبوب المتكامل من المسار المحدد
pipeline = torch.load(MODEL_PIPELINE_PATH, map_location=device, weights_only=False)

# 2. استخراج المكونات من الأنبوب
model_name = pipeline['model_name']
class_names = pipeline['class_names']
preprocess = pipeline['transforms']

# 3. بناء الموديل وتحميل الأوزان
model = get_model(model_name, pretrained=False).to(device)
model.load_state_dict(pipeline['model_state_dict'])
model.eval()

# 4. تحديد طبقة الـ CAM تلقائياً
if model_name == "resnet50":
    target_layers = [model.layer4[-1]]
elif model_name == "efficientnet_b0":
    target_layers = [model.features[-1][0]]
elif model_name == "densenet121":
    target_layers = [model.features.norm5]

cam = GradCAM(model=model, target_layers=target_layers)

def predict_and_explain(image):
    if image is None: return None, None
    img_rgb = image.convert("RGB")
    input_tensor = preprocess(img_rgb).unsqueeze(0).to(device)
    with torch.no_grad():
        output = model(input_tensor)
        probs = torch.softmax(output, dim=1)[0]
    confidences = {class_names[i]: float(probs[i]) for i in range(len(class_names))}
    grayscale_cam = cam(input_tensor=input_tensor, targets=[ClassifierOutputTarget(probs.argmax().item())], aug_smooth=True, eigen_smooth=True)[0, :]
    img_np = np.array(img_rgb.resize(IMG_SIZE)) / 255.0
    cam_image = show_cam_on_image(img_np, grayscale_cam, use_rgb=True)
    return confidences, cam_image

# إطلاق واجهة Gradio مع إضافة أسماء الفريق في الوصف
gr.Interface(
    fn=predict_and_explain,
    inputs=gr.Image(type="pil", label="ارفع صورة الأشعة هنا"),
    outputs=[gr.Label(label="التشخيص"), gr.Image(label="التفسير البصري")],
    title=f"نظام تشخيص الالتهاب الرئوي ({model_name})",
    description="""
    ### 🦅 تم تطوير وتدريب هذا النموذج بواسطة فريق النسور:
    * **المهندس عبدالله الجحافي**
    * **المهندس أيمن الوهبي**
    * **المهندس واثق المريسي**
    * **المهندس علي المحرابي**
    """,
    theme=gr.themes.Soft()
).launch(debug=True, share=True)