import streamlit as st
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
import os

# ----------------- CLASS LABELS -----------------
CLASSES = [
    "airplane", "automobile", "bird", "cat", "deer",
    "dog", "frog", "horse", "ship", "truck"
]

# ----------------- CNN MODEL (same as training) -----------------
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_layer = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),

            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )

        self.fc_layer = nn.Sequential(
            nn.Linear(64 * 8 * 8, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )

    def forward(self, x):
        x = self.conv_layer(x)
        x = x.view(x.size(0), -1)
        x = self.fc_layer(x)
        return x

# ----------------- LOAD TRAINED MODEL -----------------
model = CNN()
import os
MODEL_PATH = os.path.join(os.getcwd(), "model.pth")
model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')))

model.eval()

# ----------------- IMAGE TRANSFORMS -----------------
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
# ----------------- GRAD CAM UTILS -----------------

# ----------------- CORRECT GRAD-CAM IMPLEMENTATION -----------------

# Store activations + gradients
activations = None
gradients = None

# Save forward activations
def save_activation(module, input, output):
    global activations
    activations = output

# Save backward gradients
def save_gradient(module, grad_input, grad_output):
    global gradients
    gradients = grad_output[0]


def generate_gradcam(model, image_tensor):
    global activations, gradients

    # Last conv layer
    last_conv_layer = model.conv_layer[3]  # Conv2d(32 → 64)

    # Hook for forward activations
    forward_handle = last_conv_layer.register_forward_hook(save_activation)

    # Hook for backward gradients
    backward_handle = last_conv_layer.register_backward_hook(save_gradient)

    # Forward pass
    output = model(image_tensor)
    pred_class = output.argmax(dim=1)

    # Backward pass
    model.zero_grad()
    output[0, pred_class].backward()

    # Remove hooks
    forward_handle.remove()
    backward_handle.remove()

    # Process gradients + activations
    pooled_grads = torch.mean(gradients, dim=[0, 2, 3])
    activation_maps = activations[0]

    # Weight channels
    for i in range(len(pooled_grads)):
        activation_maps[i, :, :] *= pooled_grads[i]

    heatmap = torch.mean(activation_maps, dim=0).detach().cpu().numpy()

    # Normalize
    heatmap = np.maximum(heatmap, 0)
    heatmap = heatmap / np.max(heatmap)

    return heatmap

def overlay_heatmap(img, heatmap):
    heatmap = np.uint8(255 * heatmap)
    heatmap = Image.fromarray(heatmap).resize(img.size, Image.BILINEAR)
    heatmap = np.array(heatmap)

    # Colorize heatmap
    heatmap_color = plt.cm.jet(heatmap)[:, :, :3] * 255
    heatmap_color = heatmap_color.astype(np.uint8)

    # Overlay with original
    img_np = np.array(img)
    superimposed = (0.6 * heatmap_color + 0.4 * img_np).astype(np.uint8)

    return Image.fromarray(superimposed)

# ----------------- STREAMLIT UI -----------------
st.title("🖼️ CIFAR-10 Image Classifier (CNN)")

uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])

if uploaded_image:
    image = Image.open(uploaded_image).convert("RGB")
    st.image(image, caption="Uploaded Image", width=250)

    img_tensor = transform(image).unsqueeze(0)

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

    # Apply softmax to get probabilities
    probs = F.softmax(outputs, dim=1)[0]

    # Get top-3 predictions
    top3_prob, top3_idx = torch.topk(probs, 3)

    st.write("### 🔍 Top Predictions:")
    for i in range(3):
        st.write(f"**{CLASSES[top3_idx[i]]}: {top3_prob[i].item()*100:.2f}%**")

    # ----------------- BAR CHART -----------------
    st.write("### 📊 Probability Distribution")

    fig, ax = plt.subplots()
    ax.bar(CLASSES, probs.tolist())
    plt.xticks(rotation=45)
    st.pyplot(fig)
    # ----------------- GENERATE GRAD-CAM -----------------
    heatmap = generate_gradcam(model, img_tensor)
    cam_image = overlay_heatmap(image, heatmap)

    st.write("### 🔥 Grad-CAM Heatmap")
    st.image(cam_image, caption="Where the model is looking", use_column_width=True)