app.py

import gradio as gr
import joblib
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Headless backend for Matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import os
import warnings

warnings.filterwarnings('ignore')

# -------------------------------------------------------------------------
# Load Model and Scaler
# -------------------------------------------------------------------------
MODEL_PATH = "lightgbm_split2_best_model2.pkl"
SCALER_PATH = "robust_scaler.pkl"

try:
    model = joblib.load(MODEL_PATH)
    scaler = joblib.load(SCALER_PATH)
    print("Model and Scaler loaded successfully!")
except Exception as e:
    print(f"Error loading model or scaler: {e}")
    raise e

# Feature Definitions
FEATURE_NAMES = [
    'Ammonia (mg/l)', 'Biochemical Oxygen Demand (mg/l)', 
    'Dissolved Oxygen (mg/l)', 'Orthophosphate (mg/l)', 
    'pH (ph units)', 'Temperature (cel)', 'Nitrogen (mg/l)', 
    'Nitrate (mg/l)'
]

# Medians and IQRs from the RobustScaler
MEDIANS = [0.066, 2.7, 10.2, 0.144, 7.78, 11.46, 4.98, 4.5]
IQRS = [0.438, 1.34, 0.93, 0.247, 0.39, 5.1, 6.01, 4.25]
IMPORTANCES = [3614, 3008, 873, 3838, 959, 707, 1371, 602]

CLASSES = ['Excellent', 'Good', 'Fair', 'Marginal', 'Poor']

# -------------------------------------------------------------------------
# Explainable AI & Charting Helpers
# -------------------------------------------------------------------------
def plot_probabilities(probs):
    """Generates a vertical bar chart of prediction confidence by class."""
    colors = ['#10b981', '#34d399', '#f59e0b', '#f97316', '#ef4444']  # Excellent -> Poor
    
    fig, ax = plt.subplots(figsize=(6, 4), facecolor='none')
    ax.set_facecolor('none')
    
    bars = ax.bar(CLASSES, probs * 100, color=colors, width=0.5, edgecolor='none')
    
    # Customize grid and spines
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_color('#cbd5e1')
    ax.spines['bottom'].set_color('#cbd5e1')
    ax.tick_params(colors='#cbd5e1', labelsize=10)
    ax.grid(axis='y', linestyle='--', alpha=0.15, color='#cbd5e1')
    
    # Values on top of bars
    for bar in bars:
        height = bar.get_height()
        ax.annotate(f'{height:.1f}%',
                    xy=(bar.get_x() + bar.get_width() / 2, height),
                    xytext=(0, 4),
                    textcoords="offset points",
                    ha='center', va='bottom', color='#f8fafc', fontsize=9, fontweight='bold')
        
    ax.set_title("Class Confidence (%)", color='#06b6d4', fontsize=12, fontweight='bold', pad=12)
    ax.set_ylabel("Probability (%)", color='#cbd5e1', fontsize=10)
    ax.set_ylim(0, 115)
    
    plt.tight_layout()
    return fig

def plot_feature_importance():
    """Generates a horizontal bar chart of global feature importances."""
    sorted_idx = np.argsort(IMPORTANCES)
    sorted_features = [FEATURE_NAMES[i] for i in sorted_idx]
    sorted_importances = [IMPORTANCES[i] for i in sorted_idx]
    
    fig, ax = plt.subplots(figsize=(6, 4.2), facecolor='none')
    ax.set_facecolor('none')
    
    # Horizontal bar plot
    bars = ax.barh(sorted_features, sorted_importances, color='#3b82f6', edgecolor='none', height=0.55)
    
    # Customize grid and spines
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_color('#cbd5e1')
    ax.spines['bottom'].set_color('#cbd5e1')
    ax.tick_params(colors='#cbd5e1', labelsize=10)
    ax.grid(axis='x', linestyle='--', alpha=0.15, color='#cbd5e1')
    
    # Value annotations on the right of the bars
    for bar in bars:
        width = bar.get_width()
        ax.annotate(f' {int(width)}',
                    xy=(width, bar.get_y() + bar.get_height() / 2),
                    xytext=(3, 0),
                    textcoords="offset points",
                    ha='left', va='center', color='#cbd5e1', fontsize=9)
        
    ax.set_title("Global Feature Importance (LightGBM)", color='#06b6d4', fontsize=12, fontweight='bold', pad=12)
    ax.set_xlabel("Split Importance Score", color='#cbd5e1', fontsize=10)
    
    plt.tight_layout()
    return fig

def explain_prediction(inputs):
    """Calculates local feature contribution rankings based on deviations from normal medians."""
    impacts = []
    
    for i in range(8):
        val = inputs[i]
        med = MEDIANS[i]
        iqr = IQRS[i]
        
        # Calculate scaled deviation (impact)
        if FEATURE_NAMES[i] == 'Dissolved Oxygen (mg/l)':
            # Low DO is harmful
            impact = (med - val) / iqr
            if val < med:
                status = f"Low ({val:.2f} vs normal {med:.2f} mg/l)"
                sig = "⚠️ Reduces aquatic life support"
                color = "#ef4444"
            else:
                status = f"Optimal ({val:.2f} mg/l)"
                sig = "✅ High oxygenation, excellent health"
                color = "#10b981"
        elif FEATURE_NAMES[i] == 'pH (ph units)':
            # Extreme pH is harmful
            impact = abs(val - med) / iqr
            if val < 6.5:
                status = f"Acidic ({val:.2f} vs neutral {med:.2f})"
                sig = "⚠️ Acidic conditions degrade WQI"
                color = "#ef4444"
            elif val > 8.5:
                status = f"Alkaline ({val:.2f} vs neutral {med:.2f})"
                sig = "⚠️ Alkaline levels can be toxic"
                color = "#ef4444"
            else:
                status = f"Optimal ({val:.2f})"
                sig = "✅ Balanced neutral pH"
                color = "#10b981"
        else:
            # High values of chemical pollutants are harmful
            impact = (val - med) / iqr
            if val > med:
                status = f"High ({val:.2f} vs normal {med:.2f} mg/l)"
                sig = f"⚠️ Pollutant elevation (x{val/med:.1f})" if med > 0 else "⚠️ Elevated level"
                color = "#ef4444" if impact > 1 else "#f97316"
            else:
                status = f"Low/Normal ({val:.2f} mg/l)"
                sig = "✅ Safe background concentration"
                color = "#10b981"
                
        impacts.append({
            'feature': FEATURE_NAMES[i],
            'val': val,
            'impact': impact,
            'status': status,
            'sig': sig,
            'color': color
        })
        
    # Sort features by impact score descending
    sorted_impacts = sorted(impacts, key=lambda x: x['impact'], reverse=True)
    
    explanation_html = """
    <div style='margin-top: 15px;'>
        <h4 style='color: var(--primary-cyan); margin-bottom: 12px; font-size: 1.15rem; font-weight: 600;'>Local Feature Contribution Rankings</h4>
        <div style='display: flex; flex-direction: column; gap: 10px;'>
    """
    
    for item in sorted_impacts:
        feat = item['feature']
        status = item['status']
        sig = item['sig']
        color = item['color']
        
        explanation_html += f"""
        <div style='background: rgba(255,255,255,0.02); border-left: 4px solid {color}; padding: 8px 12px; border-radius: 4px;'>
            <div style='display: flex; justify-content: space-between; font-size: 0.95rem;'>
                <span style='font-weight: 500; color: #f8fafc;'>{feat}</span>
                <span style='color: {color}; font-weight: 600;'>{status}</span>
            </div>
            <div style='font-size: 0.85rem; color: #94a3b8; margin-top: 2px;'>{sig}</div>
        </div>
        """
        
    explanation_html += """
        </div>
    </div>
    """
    return explanation_html

# -------------------------------------------------------------------------
# Core Prediction Routine
# -------------------------------------------------------------------------
def predict_water_quality(ammonia, bod, do, orthophosphate, ph, temp, nitrogen, nitrate):
    try:
        inputs = [ammonia, bod, do, orthophosphate, ph, temp, nitrogen, nitrate]
        
        # Validation & Null Imputation
        cleaned_inputs = []
        for i, val in enumerate(inputs):
            # Check for empty string, None or NaN
            if val is None or str(val).strip() == "" or (isinstance(val, float) and np.isnan(val)):
                cleaned_inputs.append(MEDIANS[i])  # Impute with median
            else:
                try:
                    cleaned_inputs.append(float(val))
                except ValueError:
                    # Graceful validation error
                    error_html = f"""
                    <div style='text-align: center; padding: 15px; border-radius: 8px; background: rgba(239, 68, 68, 0.1); border: 1px solid rgba(239, 68, 68, 0.3);'>
                        <h4 style='color: #ef4444; margin: 0;'>⚠️ Input Validation Error</h4>
                        <p style='color: #cbd5e1; margin: 5px 0 0 0; font-size: 0.95rem;'>
                            '{val}' is not a valid number for <b>{FEATURE_NAMES[i]}</b>. Please review your inputs.
                        </p>
                    </div>
                    """
                    return error_html, None, "assets/poor.jpg", "<div style='color:#ef4444;'>Validation failed. Please enter numeric values.</div>"

        # Scale parameters using robust scaler
        inputs_df = pd.DataFrame([cleaned_inputs], columns=FEATURE_NAMES)
        scaled_inputs = scaler.transform(inputs_df)
        
        # Run prediction
        probs = model.predict_proba(scaled_inputs)[0]
        pred_class = int(np.argmax(probs))
        confidence = probs[pred_class] * 100
        
        class_name = CLASSES[pred_class]
        image_path = f"assets/{class_name.lower()}.jpg"
        
        # Color-coded badge
        badge_html = f"""
        <div style='text-align: center; padding: 12px; border-radius: 12px; background: rgba(255,255,255,0.03); border: 1px solid var(--glass-border);'>
            <span class='badge badge-{class_name.lower()}'>{class_name}</span>
            <p style='font-size: 1.2rem; margin-top: 12px; color: var(--text-primary); margin-bottom: 0;'>
                Model Confidence: <span style='color: var(--primary-cyan); font-weight: bold;'>{confidence:.2f}%</span>
            </p>
        </div>
        """
        
        # Probability Chart
        prob_chart = plot_probabilities(probs)
        
        # Feature Explanations
        explanation_html = explain_prediction(cleaned_inputs)
        
        return badge_html, prob_chart, image_path, explanation_html

    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        err_html = f"""
        <div style='text-align: center; padding: 15px; border-radius: 8px; background: rgba(239, 68, 68, 0.1); border: 1px solid rgba(239, 68, 68, 0.3);'>
            <h4 style='color: #ef4444; margin: 0;'>⚠️ Prediction Execution Failed</h4>
            <p style='color: #cbd5e1; margin: 5px 0 0 0; font-size: 0.9rem;'>{str(e)}</p>
        </div>
        """
        return err_html, None, "assets/poor.jpg", f"<pre style='color:#ef4444; font-size: 0.85rem;'>{error_details}</pre>"

# Hybrid Switch/Merge function for Sliders & Custom Text overrides
def predict_hybrid(
    s_ammonia, s_bod, s_do, s_orthophosphate, s_ph, s_temp, s_nitrogen, s_nitrate,
    t_ammonia, t_bod, t_do, t_orthophosphate, t_ph, t_temp, t_nitrogen, t_nitrate
):
    def resolve_val(text_val, slider_val):
        if text_val is not None and str(text_val).strip() != "":
            return text_val.strip()
        return slider_val

    ammonia = resolve_val(t_ammonia, s_ammonia)
    bod = resolve_val(t_bod, s_bod)
    do = resolve_val(t_do, s_do)
    orthophosphate = resolve_val(t_orthophosphate, s_orthophosphate)
    ph = resolve_val(t_ph, s_ph)
    temp = resolve_val(t_temp, s_temp)
    nitrogen = resolve_val(t_nitrogen, s_nitrogen)
    nitrate = resolve_val(t_nitrate, s_nitrate)
    
    return predict_water_quality(ammonia, bod, do, orthophosphate, ph, temp, nitrogen, nitrate)

# -------------------------------------------------------------------------
# Build Gradio Interface (Glassmorphic Dashboard Layout)
# -------------------------------------------------------------------------
with open("style.css", "r") as f:
    css_styles = f.read()

with gr.Blocks(title="Water Quality Prediction Dashboard", css=css_styles) as demo:
    
    # ----- ہیڈر (Header) -----
    with gr.Row(elem_classes=["glass-panel"], variant="compact"):
        with gr.Column(scale=1, min_width=100):
            gr.Image("assets/logo.png", show_label=False, container=False, height=90, width=90, elem_classes=["floating-logo"])
        with gr.Column(scale=8):
            gr.HTML("""
            <h1 style='margin: 0; font-size: 2.2rem; background: linear-gradient(90deg, #06b6d4, #3b82f6); -webkit-background-clip: text; -webkit-text-fill-color: transparent;'>
                Water Quality Prediction Dashboard
            </h1>
            <p style='margin: 5px 0 0 0; color: var(--text-secondary); font-size: 1.05rem;'>
                Evaluate ecosystem health, predict contamination risk, and analyze feature distributions using the pre-trained LightGBM classification model.
            </p>
            """)

    # ----- مین گرڈ (Main Grid) -----
    with gr.Row():
        
        # بائیں پینل (Inputs)
        with gr.Column(scale=5, elem_classes=["glass-panel"]):
            gr.HTML("<h3 style='color: var(--primary-cyan); border-bottom: 1px solid var(--glass-border); padding-bottom: 8px; margin-top: 0;'>Water Metrics Configurator</h3>")
            
            with gr.Tabs():
                with gr.TabItem("🎛️ Guided Sliders"):
                    s_ammonia = gr.Slider(minimum=0.0, maximum=10.0, value=0.066, step=0.001, label="Ammonia (mg/l) [Normal Median: 0.066]")
                    s_bod = gr.Slider(minimum=0.0, maximum=30.0, value=2.7, step=0.01, label="Biochemical Oxygen Demand (mg/l) [Normal Median: 2.7]")
                    s_do = gr.Slider(minimum=0.0, maximum=20.0, value=10.2, step=0.01, label="Dissolved Oxygen (mg/l) [Normal Median: 10.2]")
                    s_orthophosphate = gr.Slider(minimum=0.0, maximum=5.0, value=0.144, step=0.001, label="Orthophosphate (mg/l) [Normal Median: 0.144]")
                    s_ph = gr.Slider(minimum=0.0, maximum=14.0, value=7.78, step=0.01, label="pH (ph units) [Normal Median: 7.78]")
                    s_temp = gr.Slider(minimum=-5.0, maximum=45.0, value=11.46, step=0.01, label="Temperature (cel) [Normal Median: 11.46]")
                    s_nitrogen = gr.Slider(minimum=0.0, maximum=25.0, value=4.98, step=0.01, label="Nitrogen (mg/l) [Normal Median: 4.98]")
                    s_nitrate = gr.Slider(minimum=0.0, maximum=20.0, value=4.5, step=0.01, label="Nitrate (mg/l) [Normal Median: 4.5]")

                with gr.TabItem("📝 Precision Inputs (Overrides Sliders)"):
                    gr.HTML("<p style='font-size:0.85rem; color:var(--text-muted); margin-bottom: 12px;'>Enter precise float values below. Any non-empty textbox here overrides its corresponding slider above.</p>")
                    t_ammonia = gr.Textbox(placeholder="E.g., 0.05 (Default: 0.066)", label="Ammonia (mg/l)")
                    t_bod = gr.Textbox(placeholder="E.g., 2.50 (Default: 2.7)", label="Biochemical Oxygen Demand (mg/l)")
                    t_do = gr.Textbox(placeholder="E.g., 10.5 (Default: 10.2)", label="Dissolved Oxygen (mg/l)")
                    t_orthophosphate = gr.Textbox(placeholder="E.g., 0.12 (Default: 0.144)", label="Orthophosphate (mg/l)")
                    t_ph = gr.Textbox(placeholder="E.g., 7.6 (Default: 7.78)", label="pH (ph units)")
                    t_temp = gr.Textbox(placeholder="E.g., 12.0 (Default: 11.46)", label="Temperature (cel)")
                    t_nitrogen = gr.Textbox(placeholder="E.g., 4.5 (Default: 4.98)", label="Nitrogen (mg/l)")
                    t_nitrate = gr.Textbox(placeholder="E.g., 4.0 (Default: 4.5)", label="Nitrate (mg/l)")

        # دائیں پینل (Outputs)
        with gr.Column(scale=5):
            
            with gr.Column(elem_classes=["glass-panel"], variant="panel"):
                gr.HTML("<h3 style='color: var(--primary-cyan); margin-top: 0; border-bottom: 1px solid var(--glass-border); padding-bottom: 8px;'>Prediction Summary Card</h3>")
                
                out_badge = gr.HTML(value="""
                <div style='text-align: center; padding: 12px; border-radius: 12px; background: rgba(255,255,255,0.02); border: 1px solid var(--glass-border);'>
                    <span style='color: var(--text-muted); font-size: 1.1rem;'>Adjust parameters and click Predict</span>
                </div>
                """)
                
                out_image = gr.Image(value="assets/good.jpg", label="Ecosystem Visualisation", height=200, interactive=False)

            with gr.Column(elem_classes=["glass-panel", "top-margin"]):
                with gr.Tabs():
                    with gr.TabItem("📊 Probability Distribution"):
                        out_chart = gr.Plot(label="Confidence Distribution")
                        
                    with gr.TabItem("🔍 Explainable AI (XAI)"):
                        out_explanation = gr.HTML(value="<p style='color:var(--text-muted); font-size: 0.95rem;'>Click predict to run the explainable AI analysis.</p>")
                        
                    with gr.TabItem("📈 Global Feature Importance"):
                        gr.Plot(value=plot_feature_importance(), label="Model Feature Importance")

    # ----- بٹن (Buttons) -----
    with gr.Row(elem_classes=["top-margin"]):
        btn_clear = gr.Button("🔄 Reset to Defaults", elem_classes=["clear-btn"])
        btn_predict = gr.Button("⚡ Predict Water Quality", elem_classes=["action-btn"])

    # ----- ایونٹس (Events) -----
    input_list_sliders = [s_ammonia, s_bod, s_do, s_orthophosphate, s_ph, s_temp, s_nitrogen, s_nitrate]
    input_list_text = [t_ammonia, t_bod, t_do, t_orthophosphate, t_ph, t_temp, t_nitrogen, t_nitrate]
    output_list = [out_badge, out_chart, out_image, out_explanation]

    btn_predict.click(
        fn=predict_hybrid,
        inputs=input_list_sliders + input_list_text,
        outputs=output_list,
        api_name="predict"
    )

    def reset_inputs():
        return [0.066, 2.7, 10.2, 0.144, 7.78, 11.46, 4.98, 4.5] + [""] * 8

    btn_clear.click(
        fn=reset_inputs,
        inputs=[],
        outputs=input_list_sliders + input_list_text
    )

# -------------------------------------------------------------------------
# Launch App
# -------------------------------------------------------------------------
if __name__ == "__main__":
    demo.launch(server_name="127.0.0.1", server_port=7860)