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)

# List of available routes
routes = [
    "254U", "13U", "16D", "116BRU", "716U", "403D", "17AEU", "18U", "14U", "116BRD", "706D", "14D", "213D", "207ED", "716D", "706U", "15CU", "403U", "118D", "16U", "137JU", "402U", "20D", "11U", "103SD", "658U", "12U", "02U", "205GU", "103VD", "137JD", "305U", "117U", "108RD", "105D", "136U", "112D", "305D", "108RU", "12D", "22D", "20U", "19U", "204D", "136D", "658D", "112U", "504D", "17AD", "21D", "504U", "107JU", "106RU", "216KD", "11D", "117D", "104U", "22U", "104D", "506U", "106RD", "21U", "105U", "17AU", "216BD", "116RD", "207EU", "116RU", "402D", "18D", "118U", "903U", "01U", "206D", "206U", "903D"
]

# 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.Dropdown(label="Route ID", choices=routes),
        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()