import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import gradio as gr
from huggingface_hub import snapshot_download
import os

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# === Model Definition ===
class CNNModel(nn.Module):
    def __init__(self, dropout_rate=0.5, hidden_size=512, use_batchnorm=True):
        super(CNNModel, self).__init__()
        self.use_batchnorm = use_batchnorm
        
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32) if use_batchnorm else nn.Identity()
        
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64) if use_batchnorm else nn.Identity()
        
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128) if use_batchnorm else nn.Identity()
        
        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.bn4 = nn.BatchNorm2d(256) if use_batchnorm else nn.Identity()
        
        self.conv5 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.bn5 = nn.BatchNorm2d(256) if use_batchnorm else nn.Identity()

        self.pool = nn.MaxPool2d(2, 2)
        self.dropout = nn.Dropout(dropout_rate)
        
        self.fc1 = nn.Linear(256 * 7 * 7, hidden_size)
        self.fc2 = nn.Linear(hidden_size, 1)

    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        x = self.pool(F.relu(self.bn4(self.conv4(x))))
        x = self.pool(F.relu(self.bn5(self.conv5(x))))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# === Load Model from Hugging Face ===
def load_model(repo_id="ZacToh/RandomSearchCNN"):
    download_dir = snapshot_download(repo_id)
    model_path = os.path.join(download_dir, "cnn_final_model.pth")

    model = CNNModel()
    checkpoint = torch.load(model_path, map_location=DEVICE)
    state_dict = checkpoint.get("model_state_dict", checkpoint)
    model.load_state_dict(state_dict, strict=False)

    model.to(DEVICE)
    model.eval()
    return model

model = load_model()

# === Image Transform ===
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# === Prediction Function ===
def predict(img: Image.Image) -> str:
    img_tensor = transform(img).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        output = model(img_tensor)
        prob = torch.sigmoid(output).item()
    
    label = "hf" if prob > 0.5 else "cc"
    confidence = prob * 100 if label == "hf" else (1 - prob) * 100

    return f"Prediction: {label.upper()} ({confidence:.2f}%)"

# === Gradio UI ===
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="CNN Classifier: CC vs HF",
    description="Upload an image to classify whether it belongs to the 'cc' or 'hf' category using a CNN model."
).launch()