import gradio as gr
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import albumentations as A
from albumentations.pytorch import ToTensorV2
from huggingface_hub import hf_hub_download
import io
import requests

# Your UNET Model Definition
class UNET(nn.Module):
    def __init__(self, dropout_rate=0.1, ch=32):
        super(UNET, self).__init__()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        def conv_block(in_channels, out_channels):
            return nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 3, padding=1),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True),
                nn.Dropout2d(p=dropout_rate),
                nn.Conv2d(out_channels, out_channels, 3, padding=1),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True),
                nn.Dropout2d(p=dropout_rate)
            )
        
        self.encoder1 = conv_block(3, ch)
        self.encoder2 = conv_block(ch, ch*2)
        self.encoder3 = conv_block(ch*2, ch*4)
        self.encoder4 = conv_block(ch*4, ch*8)
        self.bottle_neck = conv_block(ch*8, ch*16)

        self.upsample1 = nn.ConvTranspose2d(ch*16, ch*8, kernel_size=2, stride=2)
        self.decoder1 = conv_block(ch*16, ch*8)
        self.upsample2 = nn.ConvTranspose2d(ch*8, ch*4, kernel_size=2, stride=2)
        self.decoder2 = conv_block(ch*8, ch*4)
        self.upsample3 = nn.ConvTranspose2d(ch*4, ch*2, kernel_size=2, stride=2)
        self.decoder3 = conv_block(ch*4, ch*2)
        self.upsample4 = nn.ConvTranspose2d(ch*2, ch, kernel_size=2, stride=2)
        self.decoder4 = conv_block(ch*2, ch)
        self.final = nn.Conv2d(ch, 1, kernel_size=1)

    def forward(self, x):
        c1 = self.encoder1(x)
        c2 = self.encoder2(self.pool(c1))
        c3 = self.encoder3(self.pool(c2))
        c4 = self.encoder4(self.pool(c3))
        c5 = self.bottle_neck(self.pool(c4))

        u6 = self.upsample1(c5)
        u6 = torch.cat([c4, u6], dim=1)
        c6 = self.decoder1(u6)
        u7 = self.upsample2(c6)
        u7 = torch.cat([c3, u7], dim=1)
        c7 = self.decoder2(u7)
        u8 = self.upsample3(c7)
        u8 = torch.cat([c2, u8], dim=1)
        c8 = self.decoder3(u8)
        u9 = self.upsample4(c8)
        u9 = torch.cat([c1, u9], dim=1)
        c9 = self.decoder4(u9)
        return self.final(c9)

# Global variables
model = None
device = torch.device('cpu')  # HF Spaces use CPU
transform = A.Compose([
    A.Resize(384, 384),
    A.Normalize(mean=(0,0,0), std=(1,1,1), max_pixel_value=255),
    ToTensorV2()
])

def load_model():
    """Load model from your HF repository"""
    global model
    try:
        print("📥 Downloading model from Hugging Face...")
        
        # Download your model from HF
        model_path = hf_hub_download(
            repo_id="ibrahim313/unet-adam-diceloss",
            filename="pytorch_model.bin"
        )
        
        # Load model
        model = UNET(ch=32)
        model.load_state_dict(torch.load(model_path, map_location=device))
        model.eval()
        
        print("✅ Model loaded successfully!")
        return "✅ Model loaded from ibrahim313/unet-adam-diceloss"
        
    except Exception as e:
        print(f"❌ Error loading model: {e}")
        return f"❌ Error: {e}"

def predict_polyp(image, threshold=0.5):
    """Predict polyp in uploaded image"""
    if model is None:
        return None, "❌ Model not loaded! Please wait for model to load.", None
    
    if image is None:
        return None, "❌ Please upload an image first!", None
    
    try:
        # Convert image to numpy array
        if isinstance(image, Image.Image):
            original_image = np.array(image.convert('RGB'))
        else:
            original_image = np.array(image)
        
        # Preprocess image
        transformed = transform(image=original_image)
        input_tensor = transformed['image'].unsqueeze(0).float()
        
        # Make prediction
        with torch.no_grad():
            prediction = model(input_tensor)
            prediction = torch.sigmoid(prediction)
            prediction = (prediction > threshold).float()
        
        # Convert to numpy
        pred_mask = prediction.squeeze().cpu().numpy()
        
        # Calculate metrics
        polyp_pixels = np.sum(pred_mask)
        total_pixels = pred_mask.shape[0] * pred_mask.shape[1]
        polyp_percentage = (polyp_pixels / total_pixels) * 100
        
        # Create visualization
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        
        # Original image
        axes[0].imshow(original_image)
        axes[0].set_title('🖼️ Original Image', fontsize=14)
        axes[0].axis('off')
        
        # Predicted mask
        axes[1].imshow(pred_mask, cmap='gray')
        axes[1].set_title('🎭 Predicted Mask', fontsize=14)
        axes[1].axis('off')
        
        # Overlay
        axes[2].imshow(original_image)
        axes[2].imshow(pred_mask, cmap='Reds', alpha=0.6)
        axes[2].set_title('🔍 Detection Overlay', fontsize=14)
        axes[2].axis('off')
        
        # Add main title with results
        if polyp_pixels > 100:
            main_title = f"🚨 POLYP DETECTED! Coverage: {polyp_percentage:.2f}%"
            title_color = 'red'
        else:
            main_title = f"✅ No Polyp Detected - Coverage: {polyp_percentage:.2f}%"
            title_color = 'green'
        
        fig.suptitle(main_title, fontsize=16, fontweight='bold', color=title_color)
        plt.tight_layout()
        
        # Save plot to image
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        buf.seek(0)
        result_image = Image.open(buf)
        plt.close()
        
        # Create detailed results text
        if polyp_pixels > 100:
            status_emoji = "🚨"
            status_text = "POLYP DETECTED"
            recommendation = "⚠️ **Recommendation:** Medical review recommended"
        else:
            status_emoji = "✅"
            status_text = "NO POLYP DETECTED"
            recommendation = "✅ **Recommendation:** Continue routine monitoring"
        
        results_text = f"""
## {status_emoji} **{status_text}**

### 📊 **Analysis Results:**
- **Polyp Coverage:** {polyp_percentage:.3f}%
- **Detected Pixels:** {int(polyp_pixels):,} / {total_pixels:,}
- **Detection Threshold:** {threshold}

### 🏥 **Clinical Assessment:**
{recommendation}

### 🔬 **Technical Details:**
- **Model:** U-Net (32 channels)
- **Input Size:** 384×384 pixels
- **Architecture:** Encoder-Decoder with skip connections
"""
        
        return result_image, results_text, pred_mask
        
    except Exception as e:
        error_msg = f"❌ **Error processing image:** {str(e)}"
        return None, error_msg, None

def load_example_image(image_num):
    """Load example images from your HF space"""
    try:
        if image_num == 1:
            # Image 1: cju0qoxqj9q6s0835b43399p4.jpg
            image_path = hf_hub_download(
                repo_id="ibrahim313/unet-adam-diceloss",
                filename="cju0qoxqj9q6s0835b43399p4.jpg",
                repo_type="space"
            )
        else:
            # Image 2: cju0roawvklrq0799vmjorwfv.jpg  
            image_path = hf_hub_download(
                repo_id="ibrahim313/unet-adam-diceloss", 
                filename="cju0roawvklrq0799vmjorwfv.jpg",
                repo_type="space"
            )
        
        # Load and return the image
        image = Image.open(image_path)
        return image
        
    except Exception as e:
        print(f"Error loading example image {image_num}: {e}")
        return None

# Load model when app starts
print("🚀 Starting Polyp Detection App...")
load_status = load_model()
print(load_status)

# Create Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(), title="🏥 Polyp Detection AI") as demo:
    
    # Header
    gr.HTML("""
    <div style="text-align: center; padding: 30px; background: linear-gradient(90deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 10px; margin-bottom: 20px;">
        <h1 style="margin: 0; font-size: 2.5em;">🏥 AI Polyp Detection System</h1>
        <p style="margin: 10px 0 0 0; font-size: 1.2em;">Advanced Medical Imaging with Deep Learning</p>
        <p style="margin: 5px 0 0 0; opacity: 0.9;">Upload colonoscopy images for intelligent polyp detection</p>
    </div>
    """)
    
    # Model info
    gr.HTML(f"""
    <div style="background: black; padding: 15px; border-radius: 8px; border-left: 4px solid #0ea5e9; margin-bottom: 20px;">
        <strong>🔬 Model:</strong> ibrahim313/unet-adam-diceloss<br>
        <strong>📏 Architecture:</strong> U-Net with 32 base channels<br>
        <strong>🎯 Dataset:</strong> Trained on Kvasir-SEG (1000 polyp images)<br>
        <strong>📸 Examples:</strong> 2 test colonoscopy images included<br>
        <strong>⚡ Status:</strong> {load_status}
    </div>
    """)
    
    # Main interface
    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML("<h3>📤 Upload Image</h3>")
            input_image = gr.Image(
                label="Drop colonoscopy image here",
                type="pil",
                height=300
            )
            
            threshold_slider = gr.Slider(
                minimum=0.1,
                maximum=0.9,
                value=0.5,
                step=0.1,
                label="🎯 Detection Sensitivity",
                info="Higher = more sensitive detection"
            )
            
            analyze_btn = gr.Button(
                "🔍 Analyze for Polyps",
                variant="primary",
                size="lg"
            )
            
            gr.HTML("<br>")
            
            # Quick examples
            gr.HTML("<h4>📸 Try Sample Images:</h4>")
            gr.HTML("<p style='font-size: 0.9em; color: #666; margin: 5px 0;'>Click to load colonoscopy test images</p>")
            with gr.Row():
                example1_btn = gr.Button("🖼️ Test Image 1", size="sm", variant="secondary")
                example2_btn = gr.Button("🖼️ Test Image 2", size="sm", variant="secondary")
        
        with gr.Column(scale=2):
            gr.HTML("<h3>📊 Detection Results</h3>")
            output_image = gr.Image(
                label="Analysis Results",
                height=400
            )
            
            results_text = gr.Markdown(
                value="Upload an image and click 'Analyze for Polyps' to see results.",
                label="Detailed Analysis"
            )
    
    # Event handlers
    analyze_btn.click(
        fn=predict_polyp,
        inputs=[input_image, threshold_slider],
        outputs=[output_image, results_text, gr.State()]
    )
    
    # Example button handlers
    example1_btn.click(
        fn=lambda: load_example_image(1),
        inputs=[],
        outputs=[input_image]
    )
    
    example2_btn.click(
        fn=lambda: load_example_image(2), 
        inputs=[],
        outputs=[input_image]
    )
    
    # Footer
    gr.HTML("""
    <div style="text-align: center; padding: 20px; margin-top: 40px; border-top: 2px solid #e5e7eb; background: #f9fafb;">
        <p style="margin: 0; color: #dc2626; font-weight: bold;">
            ⚠️ MEDICAL DISCLAIMER
        </p>
        <p style="margin: 5px 0; color: #4b5563;">
            This AI system is for research and educational purposes only.<br>
            Always consult qualified medical professionals for clinical decisions.
        </p>
        <p style="margin: 10px 0 0 0; color: #6b7280; font-size: 0.9em;">
            🔬 Powered by PyTorch | 🤗 Hosted on Hugging Face | 📊 Gradio Interface
        </p>
    </div>
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch()