import numpy as np
import gradio as gr
import glob
import imageio
from PIL import Image
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.models import resnet18, ResNet18_Weights
import albumentations as A
import tqdm
import random
from sklearn.metrics import f1_score
import torchvision.models as models
import numpy as np
import torch
import torch.nn as nn
import albumentations as A
import imageio
from PIL import Image
import gradio as gr
import os
import glob
import cv2



# 創建模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")



# 設置模型路徑
model_path = r"MRIy14.pt" 

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

class SimpleCNN(nn.Module):
    def __init__(self, num_classes):
        super(SimpleCNN, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        

        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        

        self.conv3 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1)
        self.relu3 = nn.ReLU()
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        

        self.fc = nn.Linear(32 * 28 * 28, num_classes)

    def forward(self, x):
        # 保存各層輸出以供熱力圖生成
        x1 = self.conv1(x)
        x1_relu = self.relu1(x1)
        x1_pool = self.pool1(x1_relu)

        x2 = self.conv2(x1_pool)
        x2_relu = self.relu2(x2)
        x2_pool = self.pool2(x2_relu)

        x3 = self.conv3(x2_pool)
        x3_relu = self.relu3(x3)
        x3_pool = self.pool3(x3_relu)

        x_flatten = x3_pool.view(x3_pool.size(0), -1)
        out = self.fc(x_flatten)

        # 返回最終輸出和用於熱力圖的中間特徵圖 (例如 x3_relu 或 x3_pool)
        return out, x3_relu  # 或者 return out, x3_pool



# 創建模型
model = torch.load(model_path, map_location='cpu')
model = model.to(device)
model.eval()

transform = A.Compose([
    A.Resize(224, 224),  
])





'''

def predict(image, model=model, transform=transform):
    #print(f"Input image path: {image}")
    im = np.array(image)

    im = Image.fromarray(image).convert('L')

    im = np.array(im)
    
    im = A.Resize(224, 224)(image=im)['image']

    im = np.stack([im]*3, axis=0)[None, ...] #1, 3, 224, 224
    im = im/np.max(im)

    im_d = torch.from_numpy(im).to(device).float()


    with torch.no_grad(): outputs = model(im_d)

    _, predicted = torch.max(outputs, 1)
    prediction = predicted.cpu().numpy()
    
    
    transformed_image = transform(image=im)['image']
    #transformed_image = np.transpose(transformed_image, (2, 1, 0))[None, ...]
    transformed_image = np.stack([transformed_image] * 3, axis=0)[None, ...]

    transformed_image = torch.from_numpy(transformed_image.astype(np.float32) / np.max(transformed_image)).to(device)

    with torch.no_grad():
        output = model(transformed_image.float())
        prediction = torch.argmax(output, dim=1).item()
    print(f"Prediction: {prediction}")

    print(f"Prediction: {prediction}")
    if prediction == 1:
        return "Cancer"
    else:
        return "Healthy"






demo = gr.Interface(predict, gr.Image(), outputs=["text"])

demo.launch()
'''






import torch.nn.functional as F

def generate_heatmap(model, last_conv_layer, input_tensor, pred_class):
    """
    使用 Grad-CAM 生成更精确的热力图。
    """
    # 获取最后卷积层的特征图和模型输出
    model.eval()
    features = []
    def hook_function(module, input, output):
        features.append(output)

    hook = last_conv_layer.register_forward_hook(hook_function)
    model_output = model(input_tensor)
    hook.remove()

    # 获取预测类别的得分
    class_score = model_output[:, pred_class]

    # 计算梯度
    model.zero_grad()
    class_score.backward(retain_graph=True)

    # 获取梯度和特征图
    gradients = features[0].grad  # 获取最后卷积层的梯度
    pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])  # 全局平均池化

    # 获取特征图，并应用梯度加权
    features = features[0]
    for i in range(features.shape[1]):
        features[:, i, :, :] *= pooled_gradients[i]

    # 创建热力图
    heatmap = torch.mean(features, dim=1).squeeze()
    heatmap = np.maximum(heatmap.detach().cpu().numpy(), 0)
    heatmap /= np.max(heatmap)
    return heatmap



def overlay_heatmap(image, heatmap, intensity=0.5, threshold=0.5):
    """
    优化热力图叠加逻辑。
    """
    heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]))
    heatmap = np.uint8(255 * heatmap)
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
    
    # 只在高信心区域应用热力图
    mask = heatmap > np.max(heatmap) * threshold
    superimposed_img = image.copy()
    superimposed_img[mask] = superimposed_img[mask] * (1 - intensity) + heatmap[mask] * intensity
    
    superimposed_img = np.clip(superimposed_img, 0, 255).astype(np.uint8)
    return superimposed_img



def predict(image):
    """
    預測並生成熱力圖。
    """
    # 處理圖片
    processed_image = transform(image=image)['image']
    processed_image = np.transpose(processed_image, (2, 0, 1)).astype(np.float32)
    input_tensor = torch.from_numpy(processed_image[None, ...]).to(device)  # 更有效的轉換方法

    # 預測
    with torch.no_grad():
        outputs, last_conv_output = model(input_tensor)
        _, predicted = torch.max(outputs, 1)
        prediction = predicted.item()

    # 生成熱力圖
    heatmap = generate_heatmap(last_conv_output, prediction)
    result_image = overlay_heatmap(np.array(image), heatmap)

    return prediction, result_image




# 使用 Gradio
demo = gr.Interface(fn=predict, inputs=gr.Image(), outputs=["text", "image"])
demo.launch()