import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import joblib

# Load the pre-trained model and scaler
model = joblib.load("water_quality_model.pkl")  # Ensure the model file exists
scaler = joblib.load("scaler.pkl")  # Ensure the scaler file exists

# Define water quality categories
def classify_quality(score):
    if score >= 0.8:
        return "Excellent"
    elif score >= 0.6:
        return "Good"
    elif score >= 0.4:
        return "Fair"
    else:
        return "Poor"

# Define the prediction function
def predict_water_quality(pH, turbidity, dissolved_oxygen, nitrate_levels):
    # Prepare the input data
    data = np.array([[pH, turbidity, dissolved_oxygen, nitrate_levels]])
    scaled_data = scaler.transform(data)

    # Predict the quality score
    score = model.predict(scaled_data)[0]
    quality = classify_quality(score)

    # Generate insights (graphical output)
    plt.figure(figsize=(6, 4))
    categories = ["Excellent", "Good", "Fair", "Poor"]
    scores = [0.8, 0.6, 0.4, 0.2]
    colors = ["green", "blue", "orange", "red"]

    plt.bar(categories, scores, color=colors, alpha=0.7)
    plt.axhline(y=score, color="black", linestyle="--", label=f"Predicted Score: {score:.2f} ({quality})")
    plt.legend()
    plt.title("Water Quality Prediction")
    plt.xlabel("Categories")
    plt.ylabel("Score")
    plt.tight_layout()

    # Save the graph to a file
    graph_path = "water_quality_graph.png"
    plt.savefig(graph_path)
    plt.close()

    return quality, graph_path

# Define the Gradio interface
interface = gr.Interface(
    fn=predict_water_quality,
    inputs=[
        gr.inputs.Number(label="pH"),
        gr.inputs.Number(label="Turbidity (NTU)"),
        gr.inputs.Number(label="Dissolved Oxygen (mg/L)"),
        gr.inputs.Number(label="Nitrate Levels (mg/L)"),
    ],
    outputs=[
        gr.outputs.Textbox(label="Water Quality"),
        gr.outputs.Image(label="Water Quality Graph")
    ],
    title="Water Quality Prediction",
    description="Enter water parameters to predict the quality of water. Results include the classification and a visual representation of quality."
)

# Launch the application
if __name__ == "__main__":
    interface.launch()