!pip install -q transformers accelerate bitsandbytes gradio torch pillow

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    BlipProcessor,
    BlipForConditionalGeneration,
    BitsAndBytesConfig
)
import gradio as gr
from PIL import Image
import re
from typing import List, Tuple

# Configuration for 4-bit quantization
quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True
)

class RiverPollutionAnalyzer:
    def __init__(self):
        try:
            # Initialize BLIP for image captioning
            self.blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
            self.blip_model = BlipForConditionalGeneration.from_pretrained(
                "Salesforce/blip-image-captioning-base",
                torch_dtype=torch.float16,
                device_map="auto"
            )
            
            # Initialize FLAN-T5-XL for text analysis
            self.tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-xl")
            self.model = AutoModelForSeq2SeqLM.from_pretrained(
                "google/flan-t5-xl",
                device_map="auto",
                quantization_config=quant_config
            )
            
        except Exception as e:
            raise RuntimeError(f"Model loading failed: {str(e)}")

        self.pollutants = [
            "plastic waste", "chemical foam", "industrial discharge",
            "sewage water", "oil spill", "organic debris",
            "construction waste", "medical waste", "floating trash",
            "algal bloom", "toxic sludge", "agricultural runoff"
        ]

        self.severity_descriptions = {
            1: "Minimal pollution - Slightly noticeable",
            2: "Minor pollution - Small amounts visible",
            3: "Moderate pollution - Clearly visible",
            4: "Significant pollution - Affecting water quality",
            5: "Heavy pollution - Obvious environmental impact",
            6: "Severe pollution - Large accumulation",
            7: "Very severe pollution - Major ecosystem impact",
            8: "Extreme pollution - Dangerous levels",
            9: "Critical pollution - Immediate action needed",
            10: "Disaster level - Ecological catastrophe"
        }

    def analyze_image(self, image):
        """Two-step analysis: BLIP captioning + FLAN-T5 analysis"""
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)
        
        try:
            # Step 1: Generate image caption with BLIP
            inputs = self.blip_processor(image, return_tensors="pt").to(self.blip_model.device, torch.float16)
            caption = self.blip_model.generate(**inputs, max_new_tokens=100)[0]
            caption = self.blip_processor.decode(caption, skip_special_tokens=True)
            
            # Step 2: Analyze caption with FLAN-T5
            prompt = f"""Analyze this river scene: '{caption}'
1. List visible pollutants from: {self.pollutants}
2. Estimate severity (1-10)

Respond EXACTLY as:
Pollutants: [comma separated list]
Severity: [number]"""
            
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
            outputs = self.model.generate(**inputs, max_new_tokens=200)
            analysis = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            pollutants, severity = self._parse_response(analysis)
            return self._format_analysis(pollutants, severity)
            
        except Exception as e:
            return f"⚠️ Analysis failed: {str(e)}"

    # [Keep all your existing parsing/formatting methods unchanged]
    def _parse_response(self, analysis: str) -> Tuple[List[str], int]:
        """Same parsing logic as before"""
        # ... (unchanged from your original code) ...
    
    def _calculate_severity(self, pollutants: List[str]) -> int:
        """Same severity calculation"""
        # ... (unchanged from your original code) ...
    
    def _format_analysis(self, pollutants: List[str], severity: int) -> str:
        """Same formatting"""
        # ... (unchanged from your original code) ...

    def analyze_chat(self, message: str) -> str:
        """Same chat handler"""
        # ... (unchanged from your original code) ...

# Initialize with error handling
try:
    analyzer = RiverPollutionAnalyzer()
    model_status = "✅ Models loaded successfully"
except Exception as e:
    analyzer = None
    model_status = f"❌ Model loading failed: {str(e)}"

# Gradio Interface (unchanged layout from your original)
css = """
/* [Keep your existing CSS] */
"""

with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    with gr.Column(elem_classes="header"):
        gr.Markdown("# 🌍 River Pollution Analyzer")
        gr.Markdown(f"### {model_status}")

    with gr.Row(elem_classes="side-by-side"):
        # Left Panel
        with gr.Column(elem_classes="left-panel"):
            with gr.Group():
                image_input = gr.Image(type="pil", label="Upload River Image", height=300)
                analyze_btn = gr.Button("🔍 Analyze Pollution", variant="primary")
            
            with gr.Group(elem_classes="analysis-box"):
                gr.Markdown("### 📊 Analysis report")
                analysis_output = gr.Markdown()

        # Right Panel
        with gr.Column(elem_classes="right-panel"):
            with gr.Group(elem_classes="chat-container"):
                chatbot = gr.Chatbot(label="Pollution Analysis Q&A", height=400)
                with gr.Row():
                    chat_input = gr.Textbox(
                        placeholder="Ask about pollution sources...",
                        label="Your Question",
                        container=False,
                        scale=5
                    )
                    chat_btn = gr.Button("💬 Ask", variant="secondary", scale=1)
                clear_btn = gr.Button("🧹 Clear Chat History", size="sm")

    # Connect functions
    analyze_btn.click(
        analyzer.analyze_image if analyzer else lambda x: "Model not loaded",
        inputs=image_input,
        outputs=analysis_output
    )

    # [Keep all other UI event handlers unchanged]

    # Update examples to use local files
    gr.Examples(
        examples=[
            ["examples/polluted_river1.jpg"],
            ["examples/polluted_river2.jpg"]
        ],
        inputs=image_input,
        outputs=analysis_output,
        fn=analyzer.analyze_image if analyzer else lambda x: "Model not loaded",
        cache_examples=True,
        label="Try example images:"
    )

# Launch with queue for stability
demo.queue(max_size=3).launch()