import streamlit as st
from huggingface_hub import hf_hub_download
from terratorch.cli_tools import LightningInferenceModel

# Download the model checkpoint and configuration file
ckpt_path = hf_hub_download(repo_id="ibm-granite/granite-geospatial-biomass", filename="biomass_model.ckpt")
config_path = hf_hub_download(repo_id="ibm-granite/granite-geospatial-biomass", filename="config.yaml")

# Load the model
model = LightningInferenceModel.from_config(config_path, ckpt_path)

# Create a Streamlit app
st.title("Agricultural Yield Prediction App")

# Input field for directory containing input files
input_directory = st.text_input("Enter input directory:")

# Enter Farm Data sidebar
st.sidebar.subheader("Enter Farm Data:")
farm_data = {
    "Soil Type": st.sidebar.selectbox("Soil Type", ["Clay", "Silt", "Sand"]),
    "Weather Conditions": st.sidebar.selectbox("Weather Conditions", ["Sunny", "Rainy", "Cloudy"]),
    "Crop Type": st.sidebar.selectbox("Crop Type", ["Wheat", "Corn", "Soybean"]),
    "Crop Variety": st.sidebar.text_input("Crop Variety"),
    "Soil pH": st.sidebar.number_input("Soil pH", min_value=0.0, max_value=14.0),
    "Fertilizer Application": st.sidebar.text_input("Fertilizer Application"),
    "Irrigation": st.sidebar.selectbox("Irrigation", ["Yes", "No"]),
    "Pest/Disease Management": st.sidebar.text_input("Pest/Disease Management"),
    "Weather Data": {
        "Temperature": st.sidebar.number_input("Temperature (°C)", min_value=-20.0, max_value=40.0),
        "Precipitation": st.sidebar.number_input("Precipitation (mm)", min_value=0.0, max_value=1000.0),
        "Sunshine Hours": st.sidebar.number_input("Sunshine Hours", min_value=0.0, max_value=24.0)
    },
    "Topography": {
        "Slope": st.sidebar.number_input("Slope (%)", min_value=0.0, max_value=100.0),
        "Aspect": st.sidebar.selectbox("Aspect", ["North", "South", "East", "West"])
    },
    "Previous Crop": st.sidebar.selectbox("Previous Crop", ["Wheat", "Corn", "Soybean", "None"])
}

# Button to trigger prediction
if st.button("Predict Yield"):
    if input_directory:
        # Run inference on the input directory
        inference_results, input_file_names = model.inference_on_dir(input_directory)
        
        # Display predicted yields
        st.subheader("Predicted Yields:")
        for file_name, result in zip(input_file_names, inference_results):
            st.write(f"{file_name}: {result}")
    else:
        st.error("Please enter a valid input directory.")