import gradio as gr
import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import resnet34
from PIL import Image
import json
import os

# Load model information
try:
    with open('model_info.json', 'r') as f:
        model_info = json.load(f)
except:
    model_info = {
        "model_type": "ResNet-34",
        "final_test_accuracy": 0.8801,
        "training_epochs": 20
    }

# Initialize model
model = resnet34(weights=None)
model.fc = nn.Linear(512, 2)

# Load model weights
model_loaded = False
try:
    state_dict = torch.load('fire_detection_model.pth', map_location='cpu')
    # Remove 'model.' prefix if it exists
    new_state_dict = {}
    for key, value in state_dict.items():
        if key.startswith('model.'):
            new_key = key[6:]
            new_state_dict[new_key] = value
        else:
            new_state_dict[key] = value
    
    model.load_state_dict(new_state_dict, strict=False)
    model.eval()
    model_loaded = True
    print("✅ Model loaded successfully!")
except Exception as e:
    print(f"❌ Error loading model: {e}")
    model_loaded = False

# Image preprocessing
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

class_names = ['Fire', 'Non-Fire']

def predict_fire(image):
    try:
        if not model_loaded:
            return "❌ Model not loaded"
        
        if image is None:
            return "Please upload an image"
        
        # Convert to RGB if needed
        if hasattr(image, 'mode') and image.mode != 'RGB':
            image = image.convert('RGB')
        
        # Apply transforms
        input_tensor = transform(image).unsqueeze(0)
        
        # Make prediction
        with torch.no_grad():
            outputs = model(input_tensor)
            probabilities = torch.nn.functional.softmax(outputs, dim=1)
            predicted_class = torch.argmax(probabilities, dim=1).item()
            confidence = probabilities[0][predicted_class].item()
        
        # Create result
        if predicted_class == 0:  # Fire
            result = f"🔥 **FIRE DETECTED!** 🔥\n\nConfidence: {confidence:.2%}"
        else:
            result = f"✅ **No Fire Detected** ✅\n\nConfidence: {confidence:.2%}"
        
        return result
        
    except Exception as e:
        return f"Error processing image: {str(e)}"

# Create simple Gradio interface
demo = gr.Interface(
    fn=predict_fire,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Markdown(label="Result"),
    title="🔥 Fire Detection System",
    description=f"""
    Upload an image to detect fire using ResNet-34 model.
    
    **Model Information:**
    - Architecture: {model_info['model_type']}
    - Accuracy: {model_info['final_test_accuracy']:.2%}
    - Status: {'✅ Loaded' if model_loaded else '❌ Error'}
    """,
    examples=None,
    cache_examples=False
)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )