# app.py - Fixed for 4-feature BioSentinel Model
import gradio as gr
import joblib
import numpy as np
from typing import Tuple

class BioSentinelPredictor:
    def __init__(self):
        try:
            self.model = joblib.load("model.joblib")
            self.model_loaded = True
            print(f"✅ Model loaded successfully!")
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            self.model_loaded = False
        
        # Only 4 features to match your trained model
        self.feature_columns = [
            "pH",
            "Hardness", 
            "Solids",
            "Chloramines"
        ]
        
        # Realistic ranges for water quality parameters
        self.feature_ranges = {
            "pH": (3.0, 9.5, 7.0),
            "Hardness": (47.0, 323.0, 196.0),
            "Solids": (320.0, 61227.0, 22000.0),
            "Chloramines": (0.35, 13.13, 7.0)
        }

    def predict(self, ph, hardness, solids, chloramines) -> Tuple[str, str, str]:
        if not self.model_loaded:
            return "❌ Model Error", "Could not load the trained model. Please check if model.joblib exists.", ""
        
        try:
            # Create input array with exactly 4 features (matching your model)
            features = np.array([[ph, hardness, solids, chloramines]])
            
            # Debug: Print input values
            print(f"🔍 Input values: pH={ph}, Hardness={hardness}, Solids={solids}, Chloramines={chloramines}")
            
            # Make prediction
            prediction = self.model.predict(features)[0]
            print(f"🎯 Model prediction: {prediction}")
            
            # Get confidence score if available
            confidence_text = ""
            proba_values = None
            if hasattr(self.model, 'predict_proba'):
                try:
                    proba = self.model.predict_proba(features)[0]
                    proba_values = proba
                    confidence = max(proba)
                    confidence_text = f"\nConfidence: {confidence:.1%}"
                    print(f"📊 Probabilities: {proba}")
                except Exception as e:
                    print(f"Confidence error: {e}")
            
            # Create dynamic detailed analysis based on actual parameters
            detailed_analysis = self._create_dynamic_analysis(ph, hardness, solids, chloramines, prediction, proba_values)
            
            # Interpret results with more nuanced responses
            if prediction == 1:
                result = "✅ SAFE TO DRINK"
                base_description = "Water quality analysis indicates this sample is generally SAFE for consumption."
            else:
                result = "⚠️ NOT SAFE"
                base_description = "Water quality analysis indicates potential safety concerns with this sample."
            
            # Add specific concerns based on parameters
            concerns = self._identify_concerns(ph, hardness, solids, chloramines)
            
            final_description = f"{base_description}\n\n{detailed_analysis}{confidence_text}"
            
            # Parameter breakdown
            feature_analysis = self._create_parameter_breakdown(ph, hardness, solids, chloramines)
            
            return result, final_description, feature_analysis
            
        except Exception as e:
            error_msg = str(e)
            print(f"❌ Prediction error: {error_msg}")
            return "❌ Prediction Failed", f"Error during prediction: {error_msg}", "Please check your input values and try again."

    def _create_dynamic_analysis(self, ph, hardness, solids, chloramines, prediction, proba_values):
        """Create dynamic analysis based on actual parameter values"""
        
        analysis_parts = []
        
        # pH Analysis
        if ph < 6.5:
            analysis_parts.append("🔴 **pH too acidic** - May cause corrosion and health issues")
        elif ph > 8.5:
            analysis_parts.append("🔴 **pH too alkaline** - May cause scaling and taste issues")
        elif 6.5 <= ph <= 8.5:
            analysis_parts.append("🟢 **pH within safe range** - Good for consumption")
        else:
            analysis_parts.append("🟡 **pH borderline** - Monitor closely")
        
        # Hardness Analysis
        if hardness > 300:
            analysis_parts.append("🟡 **High mineral content** - Very hard water")
        elif hardness > 180:
            analysis_parts.append("🟢 **Moderate mineral content** - Moderately hard water")
        else:
            analysis_parts.append("🟢 **Low to moderate hardness** - Acceptable levels")
        
        # Solids Analysis
        if solids > 1000:
            analysis_parts.append("🔴 **High dissolved solids** - Exceeds WHO guidelines")
        elif solids > 500:
            analysis_parts.append("🟡 **Elevated dissolved solids** - Above recommended levels")
        else:
            analysis_parts.append("🟢 **Dissolved solids acceptable** - Within safe range")
        
        # Chloramines Analysis
        if chloramines > 4.0:
            analysis_parts.append("🔴 **High chloramines** - Exceeds EPA maximum")
        elif chloramines > 2.0:
            analysis_parts.append("🟡 **Moderate chloramines** - Acceptable but elevated")
        else:
            analysis_parts.append("🟢 **Low chloramines** - Within safe limits")
        
        # Risk level based on prediction and confidence
        if prediction == 1:
            if proba_values is not None and max(proba_values) > 0.8:
                risk_text = "🟢 **Low Risk** - High confidence in safety assessment"
            elif proba_values is not None and max(proba_values) > 0.6:
                risk_text = "🟡 **Moderate Confidence** - Generally safe but monitor"
            else:
                risk_text = "🟡 **Low Confidence** - Additional testing recommended"
        else:
            if proba_values is not None and max(proba_values) > 0.8:
                risk_text = "🔴 **High Risk** - Strong indication of safety concerns"
            elif proba_values is not None and max(proba_values) > 0.6:
                risk_text = "🟡 **Moderate Risk** - Some safety concerns identified"
            else:
                risk_text = "🟡 **Uncertain Risk** - Requires professional testing"
        
        # Recommendations
        if prediction == 0:
            if ph < 6.5 or ph > 8.5:
                recommendations = "💡 **Recommended Actions:** pH adjustment, professional water treatment"
            elif solids > 1000:
                recommendations = "💡 **Recommended Actions:** Filtration system, reverse osmosis treatment"
            elif chloramines > 4.0:
                recommendations = "💡 **Recommended Actions:** Carbon filtration, contact water utility"
            else:
                recommendations = "💡 **Recommended Actions:** Professional water testing, consider treatment options"
        else:
            recommendations = "💡 **Recommended Actions:** Regular monitoring, maintain current water source quality"
        
        return f"{risk_text}\n\n" + "\n".join(analysis_parts) + f"\n\n{recommendations}"

    def _identify_concerns(self, ph, hardness, solids, chloramines):
        """Identify specific parameter concerns"""
        concerns = []
        
        if ph < 6.5 or ph > 8.5:
            concerns.append("pH out of safe range")
        if hardness > 300:
            concerns.append("very hard water")
        if solids > 500:
            concerns.append("high dissolved solids")
        if chloramines > 4.0:
            concerns.append("excessive chloramines")
            
        return concerns

    def _create_parameter_breakdown(self, ph, hardness, solids, chloramines) -> str:
        """Create detailed parameter analysis"""
        
        analysis = "📊 **Parameter Analysis:**\n\n"
        values = [ph, hardness, solids, chloramines]
        units = ["", "(mg/L)", "(ppm)", "(ppm)"]
        
        # WHO/EPA standard ranges for reference
        safe_ranges = {
            "pH": (6.5, 8.5),
            "Hardness": (0, 300),
            "Solids": (0, 500),  # WHO guideline for TDS
            "Chloramines": (0, 4.0)  # EPA maximum
        }
        
        for i, (feature, value) in enumerate(zip(self.feature_columns, values)):
            unit = units[i]
            
            # Determine status based on WHO/EPA guidelines
            if feature in safe_ranges:
                min_safe, max_safe = safe_ranges[feature]
                if min_safe <= value <= max_safe:
                    status = "✅ Within Safe Range"
                    color = "🟢"
                elif value < min_safe:
                    status = "⚠️ Below Safe Range" 
                    color = "🟡"
                else:
                    status = "❌ Above Safe Range"
                    color = "🔴"
            else:
                status = "ℹ️ Measured"
                color = "⚪"
            
            analysis += f"{color} **{feature}**: {value:.2f} {unit} - {status}\n"
        
        analysis += "\n📋 **Guidelines Used**: WHO & EPA Water Quality Standards"
        return analysis

# Initialize the predictor
predictor = BioSentinelPredictor()

def create_interface():
    with gr.Blocks(
        title="🧬 BioSentinel Water Quality Predictor",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
        }
        """
    ) as interface:
        
        # Header
        gr.Markdown("""
        # 🧬 BioSentinel - Water Quality Predictor
        
        **AI-Powered Water Safety Assessment System**
        
        Enter water quality parameters below for instant safety analysis. This system uses machine learning 
        to evaluate water potability based on critical physicochemical properties.
        """)
        
        # Main interface
        with gr.Row():
            # Input Section
            with gr.Column(scale=1):
                gr.Markdown("### 🔬 Water Quality Parameters")
                
                ph_input = gr.Slider(
                    minimum=3.0, maximum=9.5, value=7.0, step=0.1,
                    label="🌊 pH Level",
                    info="Acidity/Alkalinity measure (6.5-8.5 ideal)"
                )
                
                hardness_input = gr.Slider(
                    minimum=47.0, maximum=323.0, value=196.0, step=1.0,
                    label="💎 Hardness (mg/L)",
                    info="Calcium & magnesium content (0-300 acceptable)"
                )
                
                solids_input = gr.Slider(
                    minimum=320.0, maximum=61227.0, value=22000.0, step=100.0,
                    label="⚪ Total Dissolved Solids (ppm)",
                    info="Dissolved mineral content (<500 ideal)"
                )
                
                chloramines_input = gr.Slider(
                    minimum=0.35, maximum=13.13, value=7.0, step=0.1,
                    label="🧪 Chloramines (ppm)",
                    info="Disinfection byproducts (<4.0 safe)"
                )
                
                # Analysis Button
                predict_button = gr.Button(
                    "🔍 Analyze Water Quality", 
                    variant="primary", 
                    size="lg",
                    scale=1
                )
            
            # Results Section  
            with gr.Column(scale=1):
                gr.Markdown("### 🎯 Analysis Results")
                
                # Main result
                result_output = gr.Textbox(
                    label="Safety Assessment",
                    placeholder="Click 'Analyze Water Quality' to get results...",
                    lines=2,
                    max_lines=2
                )
                
                # Detailed analysis
                description_output = gr.Textbox(
                    label="Detailed Analysis",
                    placeholder="Detailed safety analysis will appear here...",
                    lines=4,
                    max_lines=6
                )
                
                # Parameter breakdown
                feature_output = gr.Textbox(
                    label="Parameter Breakdown",
                    placeholder="Individual parameter analysis will show here...",
                    lines=8,
                    max_lines=12
                )
        
        # Connect prediction function
        inputs = [ph_input, hardness_input, solids_input, chloramines_input]
        outputs = [result_output, description_output, feature_output]
        
        predict_button.click(
            fn=predictor.predict,
            inputs=inputs,
            outputs=outputs
        )
        
        # Quick Examples Section
        with gr.Row():
            gr.Markdown("### 📋 Quick Test Examples")
        
        with gr.Row():
            safe_btn = gr.Button("✅ Safe Water Sample", variant="secondary")
            unsafe_btn = gr.Button("❌ Unsafe Water Sample", variant="secondary")
            neutral_btn = gr.Button("⚖️ Borderline Sample", variant="secondary")
        
        # Example functions with more varied data
        def load_safe_example():
            return [7.2, 150.0, 250.0, 2.5]  # Really good quality water
        
        def load_unsafe_example():
            return [4.5, 50.0, 45000.0, 11.0]  # Really poor quality water
        
        def load_borderline_example():
            return [8.8, 320.0, 800.0, 4.5]  # Borderline case with multiple issues
        
        # Connect example buttons
        safe_btn.click(load_safe_example, outputs=inputs)
        unsafe_btn.click(load_unsafe_example, outputs=inputs)
        neutral_btn.click(load_borderline_example, outputs=inputs)
        
        # Footer
        gr.Markdown("""
        ---
        ### ⚠️ Important Disclaimer
        
        This tool provides **AI-based predictions for educational and research purposes only**. 
        For official water quality certification and regulatory compliance, always consult:
        - Certified water testing laboratories
        - Local health departments  
        - Environmental protection agencies
        
        ### 🔬 About This Model
        
        - **Technology**: Machine Learning classification model
        - **Parameters**: pH, Hardness, Total Dissolved Solids, Chloramines
        - **Standards**: Based on WHO and EPA water quality guidelines
        - **Purpose**: Educational demonstration of AI in water quality assessment
        
        **Built with ❤️ for water quality research and education**
        """)
    
    return interface

# Create the interface
demo = create_interface()

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