# ==============================
# 🚖 Uber Driver Recommendation System
# # ==============================

import numpy as np
import pandas as pd
import gradio as gr
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor

# ------------------------------
# 1. Data Generation
# ------------------------------
def generate_data(n=1000):
    np.random.seed(42)
    return pd.DataFrame({
        "pickup_distance": np.random.uniform(0.5, 10, n),
        "trip_distance": np.random.uniform(1, 20, n),
        "fare": np.random.uniform(50, 500, n),
        "surge": np.random.choice([1, 1.5, 2], n),
        "rating": np.random.uniform(3, 5, n)
    })

# ------------------------------
# 2. Feature Engineering
# ------------------------------
def feature_engineering(df):
    df = df.copy()
    df["earning_per_km"] = df["fare"] / (df["trip_distance"] + 1)
    df["efficiency"] = (df["fare"] * df["surge"]) / (
        df["pickup_distance"] + df["trip_distance"]
    )
    return df

# ------------------------------
# 3. Train Model
# ------------------------------
def train_model():
    data = generate_data()
    data = feature_engineering(data)

    data["reward"] = data["efficiency"]

    X = data.drop("reward", axis=1)
    y = data["reward"]

    X_train, _, y_train, _ = train_test_split(X, y, test_size=0.2, random_state=42)

    model = RandomForestRegressor(n_estimators=50, random_state=42)
    model.fit(X_train, y_train)

    return model, X.columns.tolist()

model, feature_columns = train_model()

# ------------------------------
# 4. Generate Ride Options (FIXED)
# ------------------------------
def generate_rides(pickup, trip, fare, surge):
    rides = []
    for _ in range(5):
        rides.append({
            "pickup_distance": max(0.5, pickup + np.random.uniform(-1, 1)),
            "trip_distance": max(1, trip + np.random.uniform(-2, 2)),
            "fare": max(50, fare + np.random.uniform(-50, 50)),
            "surge": min(2, max(1, surge + np.random.choice([0, 0.5]))),
            "rating": np.random.uniform(3, 5)  # ✅ FIX
        })
    return pd.DataFrame(rides)

# ------------------------------
# 5. Explanation Logic
# ------------------------------
def explain(row):
    reasons = []

    if row["fare"] > 300:
        reasons.append("High Fare")
    if row["pickup_distance"] < 3:
        reasons.append("Close Pickup")
    if row["surge"] > 1:
        reasons.append("Surge Benefit")
    if row["trip_distance"] > 10:
        reasons.append("Long Trip")

    return ", ".join(reasons) if reasons else "Balanced Ride"

# ------------------------------
# 6. Recommendation Engine (FIXED)
# ------------------------------
def recommend(pickup, trip, fare, surge):
    rides = generate_rides(pickup, trip, fare, surge)
    rides = feature_engineering(rides)

    # ✅ Ensure feature consistency
    rides = rides[feature_columns]

    scores = model.predict(rides)
    rides["score"] = scores

    rides = rides.sort_values(by="score", ascending=False).head(3)

    # ✅ Clean UI Output
    output = ""
    for idx, row in rides.iterrows():
        output += (
            f"🚖 Ride Option\n"
            f"Score: {round(row['score'], 2)}\n"
            f"Fare: ₹{round(row['fare'], 2)}\n"
            f"Pickup: {round(row['pickup_distance'], 2)} km\n"
            f"Trip: {round(row['trip_distance'], 2)} km\n"
            f"Surge: {row['surge']}\n"
            f"Why: {explain(row)}\n"
            f"-----------------------------\n"
        )

    return output

# ------------------------------
# 7. Gradio UI (STABLE)
# ------------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🚖 Uber Driver Recommendation System")
    gr.Markdown("AI-based smart ride selection")

    with gr.Row():
        pickup = gr.Slider(0.5, 10, value=2, label="Pickup Distance (km)")
        trip = gr.Slider(1, 20, value=5, label="Trip Distance (km)")

    with gr.Row():
        fare = gr.Slider(50, 500, value=200, label="Fare (₹)")
        surge = gr.Slider(1, 2, value=1, step=0.5, label="Surge")

    btn = gr.Button("Get Recommendation")

    output = gr.Textbox(
        label="Top Ride Recommendations",
        lines=15
    )

    btn.click(
        fn=recommend,
        inputs=[pickup, trip, fare, surge],
        outputs=output
    )

# ------------------------------
# 8. Launch
# ------------------------------
if __name__ == "__main__":
    demo.launch()