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

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

cat_col = ['Country_Name', 'Region_Name', 'State_Name', 'Crop_Name']

def stress_prediction():
    # Load model and label encoders inside the function
    best_model = joblib.load(get_artifact_path("7_stress_prediction/pest_disease_model.joblib"))
    label_enc = joblib.load(get_artifact_path("7_stress_prediction/label_pest_disease.joblib"))
    # Prediction function
    def predict(data):
        input_data = pd.DataFrame([data])

        # Encode categorical columns
        for col in cat_col:
            try:
                input_data[col] = label_enc[col].transform(input_data[col])
            except ValueError as e:
                return f"[Encoding Error] Column '{col}': {e}"

        # Predict probabilities
        proba = best_model.predict_proba(input_data)[0]
        class_labels = label_enc["Pest_Disease"].inverse_transform(range(len(proba)))

        label_percentage_list = [
            (str(label), float(round(prob * 100, 2))) 
            for label, prob in zip(class_labels, proba)
        ]

        # Get Top 3 predictions
        top_3 = sorted(label_percentage_list, key=lambda x: x[1], reverse=True)[:3]

        return top_3


    # -----------------------------
    # Streamlit Interface
    # -----------------------------
    st.title("🌱 Pest & Disease Prediction System")
    st.write("Provide environmental and crop details to predict possible pests/diseases.")

    # Sidebar Inputs
    st.sidebar.header("Input Features")

    country = st.sidebar.text_input("Country Name", "India")
    region = st.sidebar.text_input("Region Name", "Asia")
    state = st.sidebar.text_input("State Name", "Maharashtra")
    crop = st.sidebar.text_input("Crop Name", "Pomegranate")

    avg_temp = st.sidebar.number_input("Average Temperature (°C)", 0, 60, 26)
    avg_humidity = st.sidebar.number_input("Average Humidity (%)", 0, 100, 65)

    # Submit button
    if st.sidebar.button("Predict"):
        user_input = {
            "Country_Name": country,
            "Region_Name": region,
            "State_Name": state,
            "Crop_Name": crop,
            "avg_temp": avg_temp,
            "avg_humidity": avg_humidity
        }

        result = predict(user_input)

        st.subheader("🔍 Top 3 Predicted Pest/Diseases")

        if isinstance(result, str):  
            st.error(result)
        else:
            for label, score in result:
                st.write(f"**{label}** — {score}%")

            st.success("Prediction completed successfully.")