import torch
import torch.nn as nn
from torchvision import transforms, models
import gradio as gr
from PIL import Image
import numpy as np
from huggingface_hub import InferenceClient
import os

# Copyright 2025 Muhammad Zeeshan Burki
# Licensed under the Apache License, Version 2.0...


class KidneyClassifier(nn.Module):
    def __init__(self, num_classes=4):
        super(KidneyClassifier, self).__init__()
        self.backbone = models.resnet50(pretrained=False)  
        num_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(num_features, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes)
        )
    
    def forward(self, x):
        return self.backbone(x)

def load_model():
    device = torch.device('cpu')
    model = KidneyClassifier(num_classes=4)
    model.load_state_dict(torch.load('best_kidney_model.pth', map_location=device))
    model.eval()
    return model, device

model, device = load_model()
class_names = ['Normal', 'Cyst', 'Tumor', 'Stone']

class MedicalAssistant:
    def __init__(self):
        try:
            self.client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
        except:
            self.client = None
        
        self.prevention_methods = {
            'Cyst': [
                "Stay hydrated by drinking plenty of water",
                "Maintain a balanced diet low in sodium", 
                "Regular medical check-ups",
                "Avoid excessive protein intake",
                "Monitor blood pressure regularly"
            ],
            'Tumor': [
                "Avoid smoking and tobacco products",
                "Maintain a healthy weight",
                "Eat a diet rich in fruits and vegetables",
                "Limit processed and red meat consumption",
                "Regular exercise and physical activity"
            ],
            'Stone': [
                "Drink plenty of water (8-10 glasses daily)",
                "Limit sodium and oxalate-rich foods",
                "Reduce animal protein consumption",
                "Maintain adequate calcium intake",
                "Avoid excessive vitamin C supplements"
            ],
            'Normal': [
                "Continue healthy lifestyle habits",
                "Regular hydration",
                "Balanced nutrition",
                "Regular exercise",
                "Routine medical check-ups"
            ]
        }
    
    def generate_summary(self, prediction, confidence):
        if self.client:
            prompt = f"""You are a medical AI assistant. A CT kidney scan shows: {prediction} with {confidence:.2%} confidence.
            
Provide a brief, compassionate explanation and next steps. Keep it under 150 words."""
            
            try:
                response = self.client.text_generation(
                    prompt,
                    max_new_tokens=150,
                    temperature=0.7
                )
                return response
            except:
                pass
        
        return f"The analysis suggests {prediction}. Please consult with your healthcare provider for proper diagnosis and treatment recommendations."
    
    def get_prevention_methods(self, prediction):
        return self.prevention_methods.get(prediction, [])

medical_assistant = MedicalAssistant()

def predict_kidney_condition(image):
    if image is None:
        return "Please upload an image first.", {}
    
    try:
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        image_tensor = transform(image).unsqueeze(0).to(device)
        
        with torch.no_grad():
            outputs = model(image_tensor)
            probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
            confidence, predicted = torch.max(probabilities, 0)
        
        predicted_class = class_names[predicted.item()]
        confidence_score = confidence.item()
        
        summary = medical_assistant.generate_summary(predicted_class, confidence_score)
        prevention_methods = medical_assistant.get_prevention_methods(predicted_class)
        
        prob_dict = {class_names[i]: float(probabilities[i]) for i in range(len(class_names))}
        
        result_text = f"""## 🏥 Analysis Results

**Predicted Condition:** {predicted_class}
**Confidence:** {confidence_score:.1%}

## 📋 AI Summary

{summary}

## 🛡️ Prevention & Management Tips

"""
        
        for i, method in enumerate(prevention_methods, 1):
            result_text += f"\n{i}. {method}"
        
        result_text += "\n\n⚠️ **Disclaimer:** This is an AI analysis tool for educational purposes only. Always consult healthcare professionals for medical advice."
        
        return result_text, prob_dict
        
    except Exception as e:
        return f"Error processing image: {str(e)}", {}
        
def create_app():
    with gr.Blocks(title="CT Kidney Analysis", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
# 🫘 CT Kidney Analysis System

Upload a CT kidney scan to get AI-powered analysis with medical insights.

**Detects:** Normal | Cyst | Tumor | Kidney Stone
        """)
        
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(
                    label="📤 Upload CT Kidney Scan",
                    type="pil",
                    height=350
                )
                
                analyze_btn = gr.Button(
                    "🔍 Analyze Scan",
                    variant="primary", 
                    size="lg"
                )
                
                gr.Markdown("### 📋 Instructions:")
                gr.Markdown("""
1. Upload a clear CT kidney scan image
2. Click 'Analyze Scan'  
3. Review results and recommendations
4. **Always consult medical professionals**
                """)
            
            with gr.Column():
                result_output = gr.Markdown(
                    value="👆 Upload an image and click 'Analyze Scan' to begin"
                )
                
                confidence_output = gr.Label(
                    label="📊 Confidence Scores",
                    num_top_classes=4
                )
        
        analyze_btn.click(
            fn=predict_kidney_condition,
            inputs=image_input,
            outputs=[result_output, confidence_output]
        )
        
        gr.Markdown("""
---
### ⚕️ Medical Disclaimer

This AI system is for **educational and research purposes only**. It should not replace professional medical advice, diagnosis, or treatment.

### 🔬 Model Information

- **Architecture:** ResNet-50 with Transfer Learning
- **Training Data:** CT Kidney Dataset (Normal, Cyst, Tumor, Stone)  
- **AI Assistant:** HuggingFace Zephyr-7B-Beta for summaries
- **Made By:** Muhammad Zeeshan Burki 
        """)
    
    return demo

if __name__ == "__main__":
    app = create_app()
    app.launch()