import gradio as gr
import torch
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import random

# Simple mock model for demonstration
class MockSARModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.features = torch.nn.Sequential(
            torch.nn.Conv2d(1, 64, 3, padding=1),
            torch.nn.ReLU(),
            torch.nn.AdaptiveAvgPool2d((1, 1)),
            torch.nn.Flatten(),
            torch.nn.Linear(64, 1)
        )
    
    def forward(self, x):
        return self.features(x)

# Initialize model
model = MockSARModel()
model.eval()

# Transform
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.Grayscale(num_output_channels=1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485], std=[0.229])
])

def predict_oil_slick(image):
    """Predict oil slick in SAR image"""
    if image is None:
        return "Please upload an image", 0.0
    
    try:
        # Convert and analyze image
        img_array = np.array(image.convert('L'))
        
        # Simple heuristic: dark patches might indicate oil slicks
        dark_ratio = np.sum(img_array < 100) / img_array.size
        edge_variance = np.var(img_array)
        
        # Generate confidence score
        confidence = max(0.1, min(0.9, dark_ratio * 2 + edge_variance / 10000 + random.uniform(-0.1, 0.1)))
        
        # Prediction
        prediction = "Oil Slick Detected" if confidence > 0.5 else "No Oil Slick"
        
        return f"{prediction} (Confidence: {confidence:.3f})", confidence
        
    except Exception as e:
        return f"Error: {str(e)}", 0.0

# Create interface
with gr.Blocks(title="SAR Oil Slick Detection") as demo:
    gr.Markdown("""
    # 🛰️ SAR Oil Slick Detection System
    
    Interactive demo for detecting oil slicks in SAR satellite imagery.
    Part of an end-to-end maritime monitoring pipeline.
    
    **Model Repository**: [MeghanaK25/sar-oil-slick-detection](https://huggingface.co/MeghanaK25/sar-oil-slick-detection)
    """)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload SAR Image")
            analyze_btn = gr.Button("🔍 Analyze Image", variant="primary")
            
        with gr.Column():
            result_text = gr.Textbox(label="Result", lines=2)
            confidence_score = gr.Number(label="Confidence Score", precision=3)
    
    gr.Markdown("""
    ### Pipeline Context
    This model is part of an end-to-end maritime monitoring system:
    1. AIS anomaly detection flags suspicious ship behavior
    2. Sentinel-1 SAR imagery is fetched for flagged areas
    3. This model detects oil slicks with confidence scores
    4. Alerts are issued when oil slicks overlap with anomalies
    
    **Note**: This is a demonstration version simulating the production model.
    """)
    
    # Connect functionality
    analyze_btn.click(
        predict_oil_slick,
        inputs=[image_input],
        outputs=[result_text, confidence_score]
    )

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