import gradio as gr
import pandas as pd
import pickle
import numpy as np

# Load the saved artifacts
with open('random_forest_model.pkl', 'rb') as f:
    model = pickle.load(f)
with open('scaler.pkl', 'rb') as f:
    scaler = pickle.load(f)
with open('kmeans.pkl', 'rb') as f:
    kmeans = pickle.load(f)
with open('df_2024_preprocessed.pkl', 'rb') as f:
    df_2024 = pickle.load(f)

# Define feature columns (same as in training)
feature_cols = [
    'ndvi_Flowering', 'ndvi_FruitDev', 'savi_FruitDev',
    'era5_precipitation_FruitDev', 'era5_temperature_Flowering', 'era5_humidity_Harvest',
    'smap_FruitDev', 'ndvi_FruitDev_lag1', 'precip_FruitDev_lag1', 'ndvi_precip_FruitDev',
    'woreda_cluster'
]

# Get unique woredas for the dropdown
woredas = sorted(df_2024['woreda'].unique())

def predict_yield(woreda, farm_size):
    """
    Predict coffee yield for a given woreda and farm size.
    
    Args:
        woreda (str): Selected woreda.
        farm_size (float): Farm size in hectares.
    
    Returns:
        str: Predicted yield in quintals.
    """
    if farm_size <= 0:
        return "Error: Farm size must be greater than 0 hectares."
    
    # Filter df_2024 for the selected woreda
    df_woreda = df_2024[df_2024['woreda'] == woreda].copy()
    
    if df_woreda.empty:
        return f"Error: Woreda '{woreda}' not found in the dataset."
    
    # Prepare features for prediction
    X_woreda = df_woreda[feature_cols].copy()
    
    # Ensure no NaNs (should already be handled, but adding as a safeguard)
    if X_woreda.isna().any().any():
        return "Error: Missing data for the selected woreda. Please contact support."
    
    # Transform features using the saved scaler
    X_woreda_scaled = scaler.transform(X_woreda)
    
    # Predict yield per hectare
    yield_per_ha = model.predict(X_woreda_scaled)[0]
    yield_per_ha = max(0, yield_per_ha)  # Ensure non-negative prediction
    
    # Calculate total yield based on farm size
    total_yield = yield_per_ha * farm_size
    
    return (f"Predicted Coffee Yield for {woreda} (2025 Harvest):\n"
            f"- Yield per hectare: {yield_per_ha:.2f} quintals/ha\n"
            f"- Total yield for {farm_size:.2f} hectares: {total_yield:.2f} quintals")

# Create Gradio interface
interface = gr.Interface(
    fn=predict_yield,
    inputs=[
        gr.Dropdown(choices=woredas, label="Select Your Woreda"),
        gr.Number(label="Enter Farm Size (hectares)", value=1.0)
    ],
    outputs=gr.Textbox(label="Prediction Result"),
    title="Coffee Yield Prediction for Sidama Region (2025 Harvest)",
    description=("Select your woreda and enter your farm size to predict your coffee yield for the 2025 harvest season (Oct–Dec 2025). "
                 "Note: Predictions are based on historical data (2019–2022) and may have limitations. See the paper for details.")
)

# Launch the interface (this won't work locally; deploy on Hugging Face Spaces)
interface.launch()