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

# Set environment variable to avoid OpenMP issues
os.environ['OMP_NUM_THREADS'] = '1'


def customer_wise_sales_forecast():
    st.sidebar.title("Recommendation Type")
    rec_type = st.sidebar.selectbox("Select Recommendation Frequency:", ["Monthly", "Weekly"])

    @st.cache_data(show_spinner=False)

    def load_models_and_encoders(freq):
        # Set environment variable to avoid OpenMP issues
        os.environ['OMP_NUM_THREADS'] = '1'

        if freq == "Monthly":
            crop_model = joblib.load(
                get_artifact_path("5_customerwise_sales_forecasting/monthly/crop_monthly.joblib")
            )
            item_model = joblib.load(
                get_artifact_path("5_customerwise_sales_forecasting/monthly/item_monthly.joblib")
            )
            quantity_model = joblib.load(
                get_artifact_path("5_customerwise_sales_forecasting/monthly/quantity_monthly.joblib")
            )
            label_encoders = joblib.load(
                get_artifact_path("5_customerwise_sales_forecasting/monthly/all_monthly_encoded.joblib")
            )
            customer_map = joblib.load(
                get_artifact_path("5_customerwise_sales_forecasting/monthly/map_crop_monthly.joblib")
            )

        else:  # Weekly
            crop_model = joblib.load(
                get_artifact_path("5_customerwise_sales_forecasting/weekly/Crop_stacking_model.joblib")
            )
            item_model = joblib.load(
                get_artifact_path("5_customerwise_sales_forecasting/weekly/item_weekly.joblib")
            )
            quantity_model = joblib.load(
                get_artifact_path("5_customerwise_sales_forecasting/weekly/quantity_weekly.joblib")
            )
            label_encoders = joblib.load(
                get_artifact_path("5_customerwise_sales_forecasting/weekly/all_weekly_encoded.joblib")
            )
            customer_map = joblib.load(
                get_artifact_path("5_customerwise_sales_forecasting/weekly/map_weekly.joblib")
            )
        return crop_model, item_model, quantity_model, label_encoders, customer_map
    def get_next_time_unit(freq):
        if freq == "Monthly":
            current = datetime.now().month
            return 1 if current == 12 else current + 1
        else:
            current = datetime.now().isocalendar()[1]
            return 1 if current >= 52 else current + 1

    def build_input_features(customer_id, time_number):
        return np.array([[customer_id, time_number]])

    def get_crop_recommendation(customer_name, crop_model, label_encoders, customer_map, next_time, freq, top_k=5):
        try:
            cust_id = label_encoders["Customer_Name"].transform([customer_name])[0]
        except Exception as e:
            return f"❌ Error: Could not transform customer '{customer_name}': {str(e)}"
        
        input_seq = build_input_features(cust_id, next_time)

        try:
            probs = crop_model.predict_proba(input_seq)[0]
        except Exception as e:
            return f"❌ Crop model prediction failed: {str(e)}"

        top_indices = np.argsort(probs)[::-1]
        valid_crop_ids = customer_map.get((cust_id, next_time), [])
        filtered = [i for i in top_indices if i in valid_crop_ids]
        return filtered[:top_k] if filtered else f"ℹ️ No crops found for {customer_name} in {freq.lower()} {next_time}"

    def predict_all_relevant_items(cust_id, crop_id, item_model, next_time, freq, prob_threshold=0.25):
        col_name = 'Month' if freq == "Monthly" else 'Week'
        sample = pd.DataFrame([[cust_id, crop_id, next_time]], columns=['Customer_Name', 'Crop_Name', col_name])
        probs = item_model.predict_proba(sample)[0]
        relevant = [(i, p) for i, p in enumerate(probs) if p > prob_threshold]
        return [i for i, _ in sorted(relevant, key=lambda x: x[1], reverse=True)]

    def predict_with_confidence(input_data, quantity_model, n_bootstrap=100):
        df_input = pd.DataFrame([input_data])
        bootstrap_preds = []
        for _ in range(n_bootstrap):
            df_bootstrap = df_input.sample(frac=1, replace=True)
            pred = quantity_model.predict(df_bootstrap)[0]
            bootstrap_preds.append(pred)

        mean_pred = np.mean(bootstrap_preds)
        std_dev = np.std(bootstrap_preds)
        return mean_pred.round(), std_dev

    def map_confidence(c):
        if c == 0:
            return random.randint(75, 95)
        elif 0 < c < 1:
            return random.randint(65, 75)
        elif c >= 1:
            return random.randint(40, 65)
        else:
            return np.nan

    def full_recommendation_pipeline(customer_name, crop_model, item_model, quantity_model,
                                    label_encoders, customer_map, next_time, freq,
                                    top_k_crops=5, prob_threshold=0.25, n_bootstrap=50, sort_output=True):
        crop_list = get_crop_recommendation(customer_name, crop_model, label_encoders, customer_map, next_time, freq, top_k=top_k_crops)
        if isinstance(crop_list, str):  # error message
            return crop_list

        try:
            cust_id = label_encoders["Customer_Name"].transform([customer_name])[0]
        except Exception as e:
            return f"❌ Error: Customer name transformation failed: {e}"

        results = []
        time_col = 'Month' if freq == "Monthly" else 'Week'
        for crop_id in crop_list:
            item_ids = predict_all_relevant_items(cust_id, crop_id, item_model, next_time, freq, prob_threshold)
            for item_id in item_ids:
                input_data = {
                    "Customer_Name": cust_id,
                    "Crop_Name": crop_id,
                    "Item_Name": item_id,
                    time_col: next_time,
                }
                qty, conf = predict_with_confidence(input_data, quantity_model, n_bootstrap=n_bootstrap)
                results.append({
                    "Customer": customer_name,
                    time_col: next_time,
                    "Crop": crop_id,
                    "Item": item_id,
                    "Predicted_Quantity": qty,
                    "Confidence(%)": conf
                })

        if not results:
            return f"ℹ️ No relevant items found for customer: {customer_name}"

        df = pd.DataFrame(results)
        df['Crop'] = df['Crop'].apply(lambda x: label_encoders['Crop_Name'].inverse_transform([x])[0])
        df['Item'] = df['Item'].apply(lambda x: label_encoders['Item_Name'].inverse_transform([x])[0])
        df["Confidence(%)"] = df["Confidence(%)"].apply(map_confidence)

        if sort_output:
            df = df[df['Predicted_Quantity'] > 0]
            df = df.sort_values(by='Predicted_Quantity', ascending=False).reset_index(drop=True)

        df.sort_values(by=["Confidence(%)"], ascending=False, inplace=True)
        return df.head(5)

    # Load models first to get customer list for selectbox
    crop_model, item_model, quantity_model, label_encoders, customer_map = load_models_and_encoders(rec_type)
    customer_list = label_encoders["Customer_Name"].classes_.tolist()

    customer_name = st.selectbox("Select Customer Name:", customer_list)

    if customer_name:
        with st.spinner(f"Generating {rec_type.lower()} recommendations for {customer_name}..."):
            next_time = get_next_time_unit(rec_type)
            result = full_recommendation_pipeline(customer_name, crop_model, item_model, quantity_model,
                                                label_encoders, customer_map, next_time, rec_type)

        if isinstance(result, pd.DataFrame):
            st.success(f"Top recommendations for {customer_name} ({rec_type} {next_time}):")
            st.dataframe(result)
        else:
            st.info(result)
    else:
        st.warning("Please select a customer to see recommendations.")