import pandas as pd
import joblib
import gradio as gr

# Load model artifact once
artifact = joblib.load("pue_artifact.joblib")

def predict_pue(
    zone_id, hour, day_of_week, is_weekend,
    weather_factor, event_factor, traffic_index,
    distance_km, duration_min, base_fare,
    avg_distance, avg_duration, avg_fare,
    discount_usage_rate, total_rides
):

    input_data = {
        "zone_id": int(zone_id),
        "hour": int(hour),
        "day_of_week": int(day_of_week),
        "is_weekend": int(is_weekend),
        "weather_factor": float(weather_factor),
        "event_factor": float(event_factor),
        "traffic_index": float(traffic_index),
        "distance_km": float(distance_km),
        "duration_min": float(duration_min),
        "base_fare": float(base_fare),
        "avg_distance": float(avg_distance),
        "avg_duration": float(avg_duration),
        "avg_fare": float(avg_fare),
        "discount_usage_rate": float(discount_usage_rate),
        "total_rides": int(total_rides)
    }

    X = pd.DataFrame([input_data])
    X = X.reindex(columns=artifact["features"], fill_value=0)

    ride = artifact["ride_encoder"].inverse_transform(
        artifact["ride_model"].predict(X)
    )[0]

    discount_prob = artifact["discount_model"].predict_proba(X)[0][1]

    route = artifact["route_encoder"].inverse_transform(
        artifact["route_model"].predict(X)
    )[0]

    return {
        "Recommended Ride Type": ride,
        "Discount Probability": round(float(discount_prob), 2),
        "Preferred Route": route
    }

# Gradio UI
app = gr.Interface(
    fn=predict_pue,
    inputs=[
        gr.Number(label="Zone ID"),
        gr.Slider(0,23,step=1,label="Hour"),
        gr.Slider(0,6,step=1,label="Day of Week"),
        gr.Radio([0,1],label="Weekend"),
        gr.Slider(1,1.5,label="Weather Factor"),
        gr.Slider(1,1.6,label="Event Factor"),
        gr.Slider(0.5,2,label="Traffic Index"),
        gr.Number(label="Distance (km)"),
        gr.Number(label="Duration (min)"),
        gr.Number(label="Base Fare"),
        gr.Number(label="Avg Distance"),
        gr.Number(label="Avg Duration"),
        gr.Number(label="Avg Fare"),
        gr.Slider(0,1,label="Discount Usage Rate"),
        gr.Number(label="Total Rides")
    ],
    outputs="json",
    title="Personalized User Experience – Real-Time ML",
    description="Real-time personalization for ride-hailing apps"
)

app.launch()