import streamlit as st
import pandas as pd
import joblib

# Load model and data
@st.cache_resource
def load_model_and_data():
    model = joblib.load("deployment_files/Quarterly_Sales_Revenue_model_v1_0.joblib")
    #model = joblib.load("Quarterly_Sales_Revenue_model_v1_0.joblib")

    data = pd.read_csv("deployment_files/SuperKart_local.csv")

    return model, data

# Prediction logic
def predict_quarterly_sales_by_store_and_product(store_ids, df, model):
    store_results = []
    product_results = []
    total_sales = 0.0

    feature_cols = [
        'Product_Id', 'Store_Id', 'Product_Type', 'Product_Sugar_Content',
        'Store_Type', 'Store_Location_City_Type', 'Store_Size',
        'Product_Weight', 'Product_Allocated_Area', 'Product_MRP',
        'Store_Establishment_Year', 'Total_Historical_Sales',
        'Num_Past_Transactions', 'Average_Historical_Sales'
    ]

    for store_id in store_ids:
        store_df = df[df['Store_Id'] == store_id].copy()

        if store_df.empty:
            store_results.append((store_id, 0.0))
            continue

        X_input = store_df[feature_cols]
        y_pred = model.predict(X_input)

        store_total_sales = sum(y_pred)
        store_results.append((store_id, round(store_total_sales, 2)))
        total_sales += store_total_sales

        for product_id, sales in zip(store_df['Product_Id'], y_pred):
            product_results.append((store_id, product_id, round(sales, 2)))

    return store_results, product_results, round(total_sales, 2)

# Load model and data
model, df = load_model_and_data()

# UI
st.title("📈 Predict Quarterly Sales by Store and Product")

# Dropdowns
store_ids = st.multiselect("🏬 Select Store ID(s)", sorted(df["Store_Id"].unique()))
product_ids = st.multiselect("📦 Select Product ID(s) (optional)", sorted(df["Product_Id"].unique()))

# Predict button
if st.button("Predict Quarterly Sales"):
    if not store_ids:
        st.warning("Please select at least one Store ID.")
    else:
        store_results, product_results, total_sales = predict_quarterly_sales_by_store_and_product(store_ids, df, model)

        # Display Store-wise sales
        st.subheader("🏪 Store-wise Total Quarterly Sales")
        for store, sales in store_results:
            st.write(f"🛒 Store {store}: ₹{sales:,.2f}")

        # Display Product-wise sales (optional filter)
        st.subheader("📦 Product-wise Sales (Filtered by Selected Stores)")
        filtered_products = product_results
        if product_ids:
            filtered_products = [r for r in product_results if r[1] in product_ids]

        if filtered_products:
            product_df = pd.DataFrame(filtered_products, columns=["Store_Id", "Product_Id", "Predicted_Sales"])
            st.dataframe(product_df)
        else:
            st.info("No product sales to display for selected filters.")

        # Display total sales
        st.subheader("📊 Total Quarterly Sales Across All Selected Stores")
        st.success(f"₹{total_sales:,.2f}")