import pandas as pd
import joblib
import gradio as gr
import numpy as np

# -----------------------------------
# Load trained PUE artifact
# -----------------------------------
artifact = joblib.load("pue_artifact_v1.joblib")

FEATURES = artifact["features"]          # must match list below
ride_model = artifact["ride_model"]
discount_model = artifact["discount_model"]
route_model = artifact["route_model"]
ride_encoder = artifact["ride_encoder"]
route_encoder = artifact["route_encoder"]

# -----------------------------------
# Prediction Logic (CLEAN PUE)
# -----------------------------------
def predict_pue(
    zone_id,
    hour,
    is_weekend,
    weather_factor,
    traffic_index,
    avg_fare,
    avg_distance,
    avg_duration,
    discount_usage_rate,
    total_rides
):
    input_data = {
        "zone_id": zone_id,
        "hour": hour,
        "is_weekend": is_weekend,
        "weather_factor": weather_factor,
        "traffic_index": traffic_index,
        "avg_fare": avg_fare,
        "avg_distance": avg_distance,
        "avg_duration": avg_duration,
        "discount_usage_rate": discount_usage_rate,
        "total_rides": total_rides
    }

    # Enforce strict feature order (ONNX / Java safe)
    row = {f: float(input_data.get(f, 0.0)) for f in FEATURES}
    X = pd.DataFrame([[row[f] for f in FEATURES]], columns=FEATURES)

    # Predictions
    ride_pred = ride_model.predict(X)
    ride_type = ride_encoder.inverse_transform(ride_pred)[0]

    discount_prob = float(discount_model.predict_proba(X)[0][1])
    discount_prob = round(np.clip(discount_prob, 0, 1), 2)

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

    return {
        "Recommended Ride Type": ride_type,
        "Discount Probability": discount_prob,
        "Preferred Route": route
    }

# -----------------------------------
# Gradio UI (Simplified & Correct)
# -----------------------------------
app = gr.Interface(
    fn=predict_pue,
    inputs=[
        gr.Number(label="Zone ID", value=50),
        gr.Slider(0, 23, step=1, label="Hour", value=15),
        gr.Radio([0, 1], label="Weekend", value=0),
        gr.Slider(1.0, 1.5, label="Weather Factor", value=1.1),
        gr.Slider(0.5, 2.0, label="Traffic Index", value=1.3),
        gr.Number(label="Avg Fare", value=120),
        gr.Number(label="Avg Distance (km)", value=5),
        gr.Number(label="Avg Duration (min)", value=10),
        gr.Slider(0, 1, label="Discount Usage Rate", value=0.3),
        gr.Number(label="Total Rides", value=15),
    ],
    outputs=gr.JSON(label="Prediction"),
    title="Personalized User Experience",
    description=(
        "Personalized ride recommendations based purely on user behavior "
        "and real-time context. Pricing and demand signals are intentionally excluded."
    )
)

app.launch()