import pandas as pd
import gradio as gr
import matplotlib.pyplot as plt
from statsmodels.tsa.arima.model import ARIMA
import os

# =========================
# Load dataset safely
# =========================
DATA_PATH = "train.csv"

if not os.path.exists(DATA_PATH):
    raise FileNotFoundError(
        "train.csv not found. Please upload train.csv to the Hugging Face Space root directory."
    )

df = pd.read_csv(DATA_PATH)
df["date"] = pd.to_datetime(df["date"])

stores = sorted(df["store"].unique())
items = sorted(df["item"].unique())

# =========================
# Forecast single store-item
# =========================
def forecast_single(store_id, item_id, days):
    data = df[(df["store"] == store_id) & (df["item"] == item_id)]
    data = data[["date", "sales"]].sort_values("date")

    if len(data) < 50:
        return None, None, "❌ Not enough data for forecasting"

    ts = data.set_index("date")["sales"]

    model = ARIMA(ts, order=(1, 1, 1))
    model_fit = model.fit()
    forecast = model_fit.forecast(steps=days)

    forecast_df = forecast.reset_index()
    forecast_df.columns = ["Date", "Forecasted_Sales"]

    fig, ax = plt.subplots()
    ax.plot(ts, label="Historical Sales")
    ax.plot(forecast, label="Forecast")
    ax.set_title(f"Store {store_id} – Item {item_id}")
    ax.legend()

    return forecast_df, fig, forecast_df

# =========================
# Forecast ALL stores for one item
# =========================
def forecast_all_stores(item_id, days):
    results = []

    for store_id in stores:
        data = df[(df["store"] == store_id) & (df["item"] == item_id)]
        data = data[["date", "sales"]].sort_values("date")

        if len(data) < 50:
            continue

        ts = data.set_index("date")["sales"]

        model = ARIMA(ts, order=(1, 1, 1))
        model_fit = model.fit()
        forecast = model_fit.forecast(steps=days)

        forecast_df = forecast.reset_index()
        forecast_df.columns = ["Date", "Forecasted_Sales"]
        forecast_df["Store"] = store_id

        results.append(forecast_df)

    final_df = pd.concat(results)
    return final_df, final_df

# =========================
# Gradio UI
# =========================
with gr.Blocks() as demo:
    gr.Markdown("# 📊 Demand Forecasting Tool")
    gr.Markdown(
        "Forecast demand for any **Store–Item combination** using historical sales data."
    )

    gr.Markdown("## 🔹 Forecast for Single Store & Item")

    with gr.Row():
        store = gr.Dropdown(stores, label="Select Store", value=stores[0])
        item = gr.Dropdown(items, label="Select Item", value=items[0])
        days = gr.Slider(7, 90, value=30, label="Forecast Horizon (Days)")

    run_single = gr.Button("Generate Forecast")

    table_single = gr.Dataframe(label="Forecast Table")
    plot_single = gr.Plot(label="Forecast Chart")
    download_single = gr.File(label="Download CSV")

    run_single.click(
        forecast_single,
        inputs=[store, item, days],
        outputs=[table_single, plot_single, download_single]
    )

    gr.Markdown("---")
    gr.Markdown("## 🔁 Forecast One Item Across ALL Stores")

    item_all = gr.Dropdown(items, label="Select Item")
    days_all = gr.Slider(7, 90, value=30, label="Forecast Horizon (Days)")

    run_all = gr.Button("Forecast All Stores")

    table_all = gr.Dataframe(label="All Stores Forecast")
    download_all = gr.File(label="Download Combined CSV")

    run_all.click(
        forecast_all_stores,
        inputs=[item_all, days_all],
        outputs=[table_all, download_all]
    )

demo.launch()