# app.py

import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
import gradio as gr
import matplotlib.pyplot as plt

# -------------------------------------------
# 1️⃣ Create sample dataset
# -------------------------------------------

np.random.seed(42)
num_samples = 200

distance = np.random.uniform(1, 30, num_samples)
order_size = np.random.randint(1, 10, num_samples)
hour_of_day = np.random.randint(8, 23, num_samples)

# delivery_time = base + 1.2*distance + 2*order_size + 0.5*hour + noise
noise = np.random.normal(0, 5, num_samples)
delivery_time = 5 + 1.2*distance + 2*order_size + 0.5*hour_of_day + noise

df = pd.DataFrame({
    "distance": distance,
    "order_size": order_size,
    "hour_of_day": hour_of_day,
    "delivery_time": delivery_time
})

# -------------------------------------------
# 2️⃣ Train linear regression model
# -------------------------------------------

X = df[["distance", "order_size", "hour_of_day"]]
y = df["delivery_time"]

model = LinearRegression()
model.fit(X, y)

# -------------------------------------------
# 3️⃣ Define prediction + graph function
# -------------------------------------------

def predict_and_plot(distance, order_size, hour_of_day):
    # Predict delivery time
    features = np.array([[distance, order_size, hour_of_day]])
    prediction = model.predict(features)[0]

    # Create graph: vary distance from 1 to 30, keep other inputs fixed
    distances = np.linspace(1, 30, 100)
    inputs = np.column_stack((distances, np.full(100, order_size), np.full(100, hour_of_day)))
    predicted_times = model.predict(inputs)

    # Plot
    plt.figure(figsize=(6, 4))
    plt.plot(distances, predicted_times, label='Predicted delivery time', color='blue')
    plt.scatter([distance], [prediction], color='red', label='Your input', zorder=5)
    plt.xlabel('Distance (km)')
    plt.ylabel('Estimated delivery time (minutes)')
    plt.title('Delivery Time vs Distance')
    plt.legend()
    plt.tight_layout()

    # Save plot to file
    plot_path = "delivery_plot.png"
    plt.savefig(plot_path)
    plt.close()

    return f"⏱️ Estimated delivery time: {prediction:.2f} minutes", plot_path

# -------------------------------------------
# 4️⃣ Gradio interface
# -------------------------------------------

iface = gr.Interface(
    fn=predict_and_plot,
    inputs=[
        gr.Number(label="Distance (km)", value=5),
        gr.Number(label="Order Size (number of items)", value=3),
        gr.Number(label="Hour of Day (24h, e.g., 14 for 2 PM)", value=12)
    ],
    outputs=[
        gr.Text(label="Prediction"),
        gr.Image(label="Delivery Time vs Distance Graph")
    ],
    title="🚚 Delivery Time Estimator with Graph",
    description="Predicts delivery time and shows how it changes with distance (using linear regression)."
)

# -------------------------------------------
# 5️⃣ Launch app
# -------------------------------------------

if __name__ == "__main__":
    iface.launch()