import streamlit as st
import pandas as pd
import joblib
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from xgboost import XGBRegressor
from lightgbm import LGBMRegressor
from catboost import CatBoostRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error

# Title of the application
st.title("اداره زراعت و کنترل محصول")

# Model selection
model_choice = st.sidebar.selectbox("انتخاب مدل", ["Linear Regression", "Random Forest", "XGBoost", "LightGBM", "CatBoost"])

# Data upload
uploaded_file = st.sidebar.file_uploader("بارگذاری فایل داده‌ها", type=["csv", "xlsx"])
if uploaded_file is not None:
    if uploaded_file.name.endswith('.csv'):
        data = pd.read_csv(uploaded_file)
    else:
        data = pd.read_excel(uploaded_file)
    
    st.write("پیش‌نمایش داده‌ها:")
    st.dataframe(data.head())

    # Feature and label selection
    features = st.sidebar.multiselect("انتخاب ویژگی‌ها", data.columns)
    target = st.sidebar.selectbox("انتخاب برچسب", data.columns)
    
    if st.sidebar.button("آموزش مدل"):
        if features and target:
            X = data[features]
            y = data[target]

            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

            if model_choice == "Linear Regression":
                model = LinearRegression()
            elif model_choice == "Random Forest":
                model = RandomForestRegressor()
            elif model_choice == "XGBoost":
                model = XGBRegressor()
            elif model_choice == "LightGBM":
                model = LGBMRegressor()
            else:
                model = CatBoostRegressor(silent=True)

            model.fit(X_train, y_train)
            predictions = model.predict(X_test)
            error = mean_absolute_error(y_test, predictions)
            
            st.write(f"خطای مدل: {error}")
            
            # Save model
            joblib.dump(model, 'model.pkl')
            st.success("مدل با موفقیت آموزش دیده و ذخیره شد.")
        else:
            st.warning("لطفاً ویژگی‌ها و برچسب را انتخاب کنید.")
else:
    st.warning("لطفاً یک فایل داده بارگذاری کنید.")

# Prediction with saved model
st.header("پیش‌بینی بازده عملکرد")

uploaded_model = st.file_uploader("بارگذاری مدل ذخیره شده", type=["pkl"])
if uploaded_model is not None:
    model = joblib.load(uploaded_model)
    
    uploaded_data = st.file_uploader("بارگذاری فایل داده‌های جدید", type=["csv", "xlsx"])
    if uploaded_data is not None:
        if uploaded_data.name.endswith('.csv'):
            new_data = pd.read_csv(uploaded_data)
        else:
            new_data = pd.read_excel(uploaded_data)
        
        st.write("پیش‌نمایش داده‌های جدید:")
        st.dataframe(new_data.head())

        predictions = model.predict(new_data)
        st.write("پیش‌بینی‌ها:")
        st.write(predictions)

        # Generate a CSV for download
        predictions_df = pd.DataFrame(predictions, columns=['Predictions'])
        csv = predictions_df.to_csv(index=False).encode('utf-8')

        st.download_button(
            label="دانلود پیش‌بینی‌ها",
            data=csv,
            file_name='predictions.csv',
            mime='text/csv',
        )
else:
    st.warning("لطفاً یک مدل ذخیره شده بارگذاری کنید.")