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


def expense_forecasting():
    st.set_page_config(page_title="🌾 Crop Expense Predictor", layout="centered")

    # ------------------------
    # Load models and encoders
    # ------------------------
    @st.cache_resource
    def load_models():
        model = joblib.load(get_artifact_path("8_expense_forecasting/combined_model.joblib"))
        return model

    model = load_models()
    label_enc = model["label_encoder"]
    activity_model = model["activity"]
    activity_count_model = model["activity_count"]
    expense_model = model["expense"]

    # Get original category names from LabelEncoders
    crop_names = label_enc["Crop_Name"].classes_.tolist()
    variety_names = label_enc["Variety_Name"].classes_.tolist()
    season_names = label_enc["Season_Name"].classes_.tolist()

    # ------------------------
    # Full Prediction Pipeline
    # ------------------------
    def full_pipeline(input_dict):
        df = pd.DataFrame([input_dict])

        # Encode categorical inputs
        cat_cols = ["Crop_Name", "Variety_Name", "Season_Name"]
        for col in cat_cols:
            df[col] = label_enc[col].transform(df[col].astype(str).str.strip().str.title())

        # Predict Expense_Activity
        probas = activity_model.predict_proba(df)[0]
        class_names = activity_model.classes_

        # Filter based on probability threshold
        filtered_activities = [class_names[i] for i in range(len(probas)) if probas[i] >= 0.30]

        results = []
        for encoded_activity in filtered_activities:
            row = df.copy()
            row["Activity"] = encoded_activity

            # Predict Count
            mean_count = activity_count_model.predict(row)[0].round().astype("int64")
            row["Count"] = mean_count

            # Predict Expense
            expense = expense_model.predict(row)[0].round().astype("int64")

            # Decode activity
            decoded_activity = label_enc["Activity"].inverse_transform([encoded_activity])[0]

            results.append({
                "Predicted_Activity": decoded_activity,
                "Predicted_Count": int(mean_count),
                "Predicted_Expense": int(expense)
            })

        return results

    # ------------------------
    # Streamlit UI
    # ------------------------
    st.title("🌱 Crop Expense & Activity Predictor")

    st.markdown("Enter crop details to predict expense activities, counts, and total expense.")

    # Input Form
    with st.form("prediction_form"):
        crop_name = st.selectbox("Crop Name", sorted(crop_names), index=0)
        variety_name = st.selectbox("Variety Name", sorted(variety_names), index=0)
        season_name = st.selectbox("Season Name", sorted(season_names), index=0)
        submitted = st.form_submit_button("Predict")

    # Prediction logic
    if submitted:
        input_data = {
            "Crop_Name": crop_name,
            "Variety_Name": variety_name,
            "Season_Name": season_name
        }

        st.subheader("🔍 Input Data")
        st.json(input_data)

        try:
            results = full_pipeline(input_data)

            if results:
                st.subheader("📊 Prediction Results")
                for res in results:
                    st.markdown(f"""
                    **Activity**: `{res['Predicted_Activity']}`  
                    **Estimated Count**: **{res['Predicted_Count']}**  
                    **Estimated Expense**: **₹{res['Predicted_Expense']}**  
                    ---  
                    """)
            else:
                st.warning("⚠️ No activity met the 30% probability threshold. Try different inputs.")

        except Exception as e:
            st.error(f"🚫 Error: {e}")