import streamlit as st
import pandas as pd
import joblib

MODEL_PATH = 'src/water_quality.joblib'

FEATURES = [
    'ph', 'Hardness', 'Solids', 'Chloramines', 'Sulfate', 
    'Conductivity', 'Organic_carbon', 'Trihalomethanes', 'Turbidity'
]

@st.cache_resource
def load_pipeline():
    try:
        pipeline = joblib.load(MODEL_PATH)
        return pipeline
    except Exception as e:
        st.error(f"Error loading model. Ensure '{MODEL_PATH}' exists. Error: {e}")
        return None

def make_prediction(pipeline, input_data):
    df_input = pd.DataFrame([input_data])
    predictions = pipeline.predict(df_input)
    probability = pipeline.predict_proba(df_input)[:, 1][0]  # probability of potable
    label = int(predictions[0])
    return label, probability

# --- Streamlit Interface ---
st.set_page_config(page_title="Water Quality Predictor", layout="centered")
st.title("💧 Water Potability Prediction")

pipeline = load_pipeline()

if pipeline:
    st.sidebar.header("Water Quality Parameters")
    ph = st.sidebar.number_input("pH Level:", 0.0, 14.0, 7.0, 0.1)
    hardness = st.sidebar.number_input("Hardness:", 50.0, 400.0, 200.0)
    solids = st.sidebar.number_input("Solids (ppm):", 300.0, 60000.0, 20000.0)
    chloramines = st.sidebar.number_input("Chloramines (ppm):", 0.0, 15.0, 7.0, 0.1)
    sulfate = st.sidebar.number_input("Sulfate (ppm):", 100.0, 500.0, 300.0)
    conductivity = st.sidebar.number_input("Conductivity (μS/cm):", 200.0, 800.0, 400.0)
    organic_carbon = st.sidebar.number_input("Organic Carbon (ppm):", 5.0, 30.0, 15.0)
    trihalomethanes = st.sidebar.number_input("Trihalomethanes (ppm):", 10.0, 150.0, 70.0)
    turbidity = st.sidebar.number_input("Turbidity (NTU):", 1.0, 7.0, 4.0)

    input_data = {
        'ph': ph,
        'Hardness': hardness,
        'Solids': solids,
        'Chloramines': chloramines,
        'Sulfate': sulfate,
        'Conductivity': conductivity,
        'Organic_carbon': organic_carbon,
        'Trihalomethanes': trihalomethanes,
        'Turbidity': turbidity
    }

    if st.button("Predict Potability"):
        label, prob = make_prediction(pipeline, input_data)
        status = "POTABLE ✅" if label == 1 else "NOT POTABLE ❌"
        st.success(f"Prediction: {status}")
        st.info(f"Probability of being potable: {prob*100:.2f}%")