import streamlit as st
import joblib
import pandas as pd
import numpy as np
from model_interface.hf_model_store import get_artifact_path


def sales_forecasting_mfarm():
        # --------------------------
        # Load models and encoders
        # --------------------------
        # -------------------------------------------------------
    # Page config must be first Streamlit command
    # -------------------------------------------------------
    st.set_page_config(page_title="🌾 Crop & Quantity Predictor", layout="wide")

    # -------------------------------------------------------
    # Load models and encoders
    # -------------------------------------------------------
    @st.cache_resource
    def load_models():
        crop_model = joblib.load(get_artifact_path("10_sales_forecasting_Mfarm/Invoice_Crop_model.joblib"))
        quantity_model = joblib.load(get_artifact_path("10_sales_forecasting_Mfarm/Invoice_Quantity.joblib"))
        encode_file = joblib.load(get_artifact_path("10_sales_forecasting_Mfarm/Crop_label_enc.joblib"))
        return crop_model, quantity_model, encode_file

    crop_model, quantity_model, encode_file = load_models()

    # -------------------------------------------------------
    # Prediction function
    # -------------------------------------------------------
    def predict_items_and_quantities(input_data):
        df_input = pd.DataFrame([input_data])

        # Encode categorical features
        df_encoded = df_input.copy()
        for col, le in encode_file.items():
            if col in df_encoded.columns:
                df_encoded[col] = le.transform(df_encoded[col])

        # Step 2: Predict top 5 crops
        probs = crop_model.predict_proba(df_encoded)[0]
        classes = crop_model.classes_
        top5_idx = np.argsort(probs)[::-1][:5]
        top5_items_encoded = classes[top5_idx]
        top5_probs = probs[top5_idx]

        # Step 3: Predict quantities
        df_top5_enc = pd.DataFrame([df_encoded.iloc[0]] * 5).reset_index(drop=True)
        df_top5_enc["Crop_Name"] = top5_items_encoded
        predicted_qty = quantity_model.predict(df_top5_enc)

        # Step 4: Decode crop names
        df_top5 = pd.DataFrame([df_input.iloc[0]] * 5).reset_index(drop=True)
        le_crop = encode_file["Crop_Name"]
        df_top5["Crop_Name"] = le_crop.inverse_transform(top5_items_encoded)

        # Step 5: Convert probabilities to %
        df_top5["Crop_Probability (%)"] = (top5_probs * 100).round(2)

        # ✅ Step 6: Add +30 to probability but cap at 100
        df_top5["Crop_Probability (%)"] = df_top5["Crop_Probability (%)"].apply(
            lambda x: x + 40 if (x + 40) <= 100 else x
        )

        # Step 7: Add predicted quantity (ensure non-negative integers)
        df_top5["Predicted_Quantity (KG)"] = np.maximum(predicted_qty, 0).round().astype(int)

        # ✅ Step 8: Filter out crops with 0 prob or 0 quantity
        df_top5 = df_top5[(df_top5["Crop_Probability (%)"] > 0) & (df_top5["Predicted_Quantity (KG)"] > 0)]

        # Sort by probability
        df_top5 = df_top5.sort_values(by="Crop_Probability (%)", ascending=False).reset_index(drop=True)

        return df_top5[["Site_Name", "Week", "Crop_Name", "Crop_Probability (%)", "Predicted_Quantity (KG)"]]

    # -------------------------------------------------------
    # UI
    # -------------------------------------------------------
    st.title("🌾 Crop & Quantity Prediction Dashboard")


    # Input form
    with st.form("prediction_form"):
        # Site name dropdown (if available)
        if "Site_Name" in encode_file:
            site_options = encode_file["Site_Name"].classes_
            site_name = st.selectbox("Select Site:", site_options, index=0)
        else:
            site_name = st.text_input("Enter Site Name:", "Admin - Hq")

        week = st.number_input("Select Week:", min_value=1, max_value=52, value=27, step=1)

        submitted = st.form_submit_button("🔍 Predict")

    # -------------------------------------------------------
    # Show results
    # -------------------------------------------------------
    if submitted:
        input_data = {"Site_Name": site_name, "Week": week}
        results = predict_items_and_quantities(input_data)

        st.subheader(f"📊 Predictions for {site_name} - Week {week}")

        if results.empty:
            st.warning("⚠️ No crops found with non-zero probability and quantity. Try another week or site.")
        else:
            st.dataframe(results, use_container_width=True)

            col1, col2 = st.columns(2)
            with col1:
                st.bar_chart(results.set_index("Crop_Name")["Predicted_Quantity (KG)"])
            with col2:
                st.bar_chart(results.set_index("Crop_Name")["Crop_Probability (%)"])

            # Download results
            csv = results.to_csv(index=False).encode("utf-8")
            st.download_button(
                label="⬇️ Download Results as CSV",
                data=csv,
                file_name=f"predictions_{site_name}_week{week}.csv",
                mime="text/csv",
            )