import streamlit as st
import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
import joblib
from huggingface_hub import hf_hub_download, upload_file

# Load or initialize the model
def load_model():
    try:
        # Attempt to download model from Hugging Face
        model_path = hf_hub_download(repo_id="your-huggingface-repo", filename="crop_yield_model.pkl")
        model = joblib.load(model_path)
        return model
    except Exception as e:
        st.warning("Model not found on Hugging Face. Initializing a new one.")
        # Initialize with dummy training data
        model = LinearRegression()
        dummy_features = np.array([[100, 50, 25]])  # Example realistic values for rainfall, fertilizer, temperature
        dummy_target = np.array([2])  # Example realistic yield value
        model.fit(dummy_features, dummy_target)
        return model

def save_model(model):
    joblib.dump(model, "crop_yield_model.pkl")
    upload_file(
        path_or_fileobj="crop_yield_model.pkl",
        path_in_repo="crop_yield_model.pkl",
        repo_id="your-huggingface-repo",
        repo_type="model",
    )

# Streamlit app
st.title("Crop Yield Prediction")

# Input features
st.sidebar.header("Input Features")
rainfall = st.sidebar.number_input("Rainfall (mm)", min_value=0.0, max_value=5000.0, step=0.1)
fertilizer = st.sidebar.number_input("Fertilizer Used (kg/ha)", min_value=0.0, max_value=1000.0, step=0.1)
temperature = st.sidebar.number_input("Temperature (°C)", min_value=-10.0, max_value=50.0, step=0.1)

# Train new model option
train_new = st.sidebar.checkbox("Train New Model")

# Load data for training if needed
if train_new:
    st.sidebar.header("Training Data")
    uploaded_file = st.sidebar.file_uploader("Upload a CSV file", type="csv")
    if uploaded_file:
        data = pd.read_csv(uploaded_file)
        st.write("### Training Data Preview", data.head())
        features = data[["Rainfall", "Fertilizer", "Temperature"]]
        target = data["Yield"]

        # Train model
        model = LinearRegression()
        model.fit(features, target)
        save_model(model)
        st.success("Model trained and saved successfully!")
    else:
        st.warning("Please upload a CSV file to train a new model.")
else:
    model = load_model()

# Predict crop yield
if st.button("Predict Crop Yield"):
    if model is not None:
        input_data = np.array([[rainfall, fertilizer, temperature]])
        try:
            prediction = model.predict(input_data)
            st.write(f"### Predicted Crop Yield: {prediction[0]:.2f} tons/ha")
        except Exception as e:
            st.error("An error occurred during prediction. Please ensure the model is properly trained.")
    else:
        st.warning("Model not found or not trained. Please train a new model.")

# Notes for Hugging Face integration
st.sidebar.markdown("---")
st.sidebar.write("Hugging Face Integration:")
st.sidebar.write("- Replace `your-huggingface-repo` with your repository name.")
st.sidebar.write("- Ensure the repository permissions allow model uploads and downloads.")