import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision.models as models
from skimage.segmentation import watershed
from scipy import ndimage as ndi
from skimage.feature import peak_local_max
import gradio as gr

# -----------------------------
# Config
# -----------------------------
model_path = "fusion_best.pth"
class_names = ["512Glioma", "512Meningioma", "512Normal", "512Pituitary"]
img_size = (224, 224)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# -----------------------------
# Preprocessing
# -----------------------------
def preprocess_image(img):
    img_resized = cv2.resize(np.array(img), img_size)
    img_gray = cv2.cvtColor(img_resized, cv2.COLOR_BGR2GRAY)
    img_norm = (img_gray - np.mean(img_gray)) / (np.std(img_gray) + 1e-8)
    img_norm = np.clip(img_norm, -3, 3)
    img_norm = ((img_norm - img_norm.min()) / (img_norm.max() - img_norm.min()) * 255).astype(np.uint8)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    img_clahe = clahe.apply(img_norm)
    img_denoised = cv2.fastNlMeansDenoising(img_clahe, h=10, templateWindowSize=7, searchWindowSize=21)
    ret, thresh = cv2.threshold(img_denoised, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    distance = ndi.distance_transform_edt(thresh)
    coordinates = peak_local_max(distance, labels=thresh, footprint=np.ones((3,3)))
    markers = np.zeros_like(distance, dtype=int)
    for i, coord in enumerate(coordinates, 1):
        markers[coord[0], coord[1]] = i
    labels = watershed(-distance, markers, mask=thresh)
    watershed_mask = (labels > 0).astype(np.uint8)
    img_final_gray = cv2.bitwise_and(img_denoised, img_denoised, mask=watershed_mask*255)
    img_final_rgb = np.stack([img_final_gray]*3, axis=-1)
    return img_final_rgb

# -----------------------------
# Model
# -----------------------------
class FusionNet(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        mobilenet = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)
        self.mobilenet = mobilenet.features
        self.mobilenet_out = mobilenet.last_channel

        densenet = models.densenet201(weights=models.DenseNet201_Weights.IMAGENET1K_V1)
        self.densenet = densenet.features
        self.densenet_out = densenet.classifier.in_features

        self.gap = nn.AdaptiveAvgPool2d((1,1))
        self.fc = nn.Sequential(
            nn.Linear(self.mobilenet_out + self.densenet_out, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_classes)
        )

        # Freeze backbones
        for p in self.mobilenet.parameters():
            p.requires_grad = False
        for p in self.densenet.parameters():
            p.requires_grad = False

    def forward(self, x):
        x1 = self.mobilenet(x)
        x1_gap = self.gap(x1).view(x1.size(0), -1)

        x2 = self.densenet(x)
        self.feature_map_densenet = x2.detach()  # store feature map for Grad-CAM (detached!)
        x2_gap = self.gap(x2).view(x2.size(0), -1)

        fused = torch.cat([x1_gap, x2_gap], dim=1)
        out = self.fc(fused)
        return out

# Load model
model = FusionNet(len(class_names)).to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()

# -----------------------------
# Grad-CAM DenseNet201 only
# -----------------------------
def generate_gradcam_densenet(feature_map, original_img):
    # Compute channel-wise mean for Grad-CAM without backward (memory efficient)
    cam = feature_map.mean(dim=1).squeeze().cpu().numpy()
    cam = np.maximum(cam, 0)
    cam = cv2.resize(cam, (original_img.shape[1], original_img.shape[0]))
    cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
    heatmap = cv2.applyColorMap(np.uint8(255*cam), cv2.COLORMAP_JET)
    overlay = cv2.addWeighted(original_img, 0.6, heatmap, 0.4, 0)
    return overlay

# -----------------------------
# Prediction + Grad-CAM
# -----------------------------
def predict_with_gradcam(img):
    img_np = np.array(img)
    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    processed_img = preprocess_image(img_bgr)
    input_tensor = torch.from_numpy(processed_img / 255.0).permute(2,0,1).float().unsqueeze(0).to(device)
    mean = torch.tensor([0.485,0.456,0.406],device=device).view(1,3,1,1)
    std = torch.tensor([0.229,0.224,0.225],device=device).view(1,3,1,1)
    input_tensor = (input_tensor - mean)/std

    # Forward pass
    with torch.no_grad():  # no gradients saved → memory optimized
        output = model(input_tensor)
        probs = torch.nn.functional.softmax(output[0], dim=0)
        pred_idx = torch.argmax(output[0]).item()
        pred_class = class_names[pred_idx]
        confidence = probs[pred_idx].item()*100

    # Grad-CAM (DenseNet only, memory efficient)
    gradcam_densenet = generate_gradcam_densenet(model.feature_map_densenet, processed_img)

    return f"Prediction: {pred_class}\nConfidence: {confidence:.2f}%", gradcam_densenet

# -----------------------------
# Gradio Interface
# -----------------------------
demo = gr.Interface(
    fn=predict_with_gradcam,
    inputs=gr.Image(type="pil", label="Upload Brain MRI Image"),
    outputs=[
        gr.Textbox(label="Prediction Result"),
        gr.Image(label="Grad-CAM DenseNet201")
    ],
    title="Brain Tumor Classification",
    description="Upload a brain MRI image to classify and grad cam visualization."
)

if __name__ == "__main__":
    demo.launch()