import gradio as gr
import pickle
import pandas as pd

# Load trained XGBoost model
with open("xgboost_trip_delay_model.pkl", "rb") as f:
    model = pickle.load(f)

# Define the prediction function for Gradio
def predict_bus_delay(route_id, trip_direction, speed, trip_delay_y, hour, minute, last_stop_arrival_seconds):
    data = pd.DataFrame({
        "route_id": [route_id],  
        "trip_direction": [1 if trip_direction == "UP" else 0],
        "speed": [float(speed)],
        "trip_delay_y": [float(trip_delay_y)],
        "hour": [int(hour)],
        "minute": [int(minute)],
        "last_stop_arrival_seconds": [int(last_stop_arrival_seconds)]
    })

    data["route_id"] = data["route_id"].astype("category")

    # Make prediction
    prediction = model.predict(data)[0]
    return f"Predicted Bus Delay Between Stops: {prediction:.2f} seconds"

# Create Gradio interface
interface = gr.Interface(
    fn=predict_bus_delay,
    inputs=[
        gr.Textbox(label="Route ID"),
        gr.Dropdown(label="Direction", choices=["UP", "DN"]),
        gr.Number(label="Speed"),
        gr.Number(label="Previous Stop Delay"),
        gr.Number(label="Hour"),
        gr.Number(label="Minute"),
        gr.Number(label="Last Stop Arrival (s)")
    ],
    outputs="text"
)

# Launch the Gradio interface
interface.launch()