import torch
import torch.nn as nn
from torchvision.models import resnet34
from torchvision.transforms import transforms
from PIL import Image

class FireDetectionModel(nn.Module):
    def __init__(self, num_classes=2):
        super(FireDetectionModel, self).__init__()
        self.model = resnet34(pretrained=False)
        self.model.fc = nn.Linear(512, num_classes)
        
    def forward(self, x):
        return self.model(x)
    
    def predict(self, image):
        """Predict fire detection from input image"""
        # Preprocessing transform
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        # Convert to tensor and add batch dimension
        if isinstance(image, Image.Image):
            image_tensor = transform(image).unsqueeze(0)
        else:
            image_tensor = image
            
        # Inference
        self.eval()
        with torch.no_grad():
            outputs = self.forward(image_tensor)
            probabilities = torch.softmax(outputs, dim=1)
            predicted_class = torch.argmax(probabilities, dim=1)
            confidence = torch.max(probabilities, dim=1)[0]
            
        class_names = ['Non-Fire', 'Fire']
        return {
            'prediction': class_names[predicted_class.item()],
            'confidence': confidence.item(),
            'probabilities': {
                'Non-Fire': probabilities[0][0].item(),
                'Fire': probabilities[0][1].item()
            }
        }