import streamlit as st
from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LogisticRegression
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from sklearn import datasets
from sklearn.impute import SimpleImputer  # Add this line
import pandas as pd
import numpy as np
st.header("Prediction")

if "data" in st.session_state:
    data = st.session_state.data

    with st.form("my_form"):
        header = st.columns([2,1,2])
        header[0].subheader("Model")
        header[1].subheader("Data Name")

        row1 = st.columns([2,1,2])
        model = row1[0].selectbox("", ["Random Forest Classifier", "Random Forest Regressor", "Logistic Regression", "Linear Regression"])
        data_name = row1[1].selectbox("", data.columns)

        st.form_submit_button('launch')
            
    if model == "Random Forest Classifier":
        model = RandomForestClassifier()
    elif model == "Random Forest Regressor":
        model = RandomForestRegressor()
    elif model == "Logistic Regression":
        model = LogisticRegression()
    elif model == "Linear Regression":
        model = LinearRegression()
    
    x = data.drop(data_name, axis=1)
    y = data[data_name]
    
    # Convert categorical data to numerical values
    x = pd.get_dummies(x)
    
    # Handle missing values
    imputer = SimpleImputer()
    x = imputer.fit_transform(x)
    x = pd.get_dummies(x)
    
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42)
    model.fit(x_train, y_train)
    y_pred = model.predict(x_test)
    if model == "Random Forest Classifier":
        st.write("Accuracy: ", accuracy_score(y_test, y_pred))
    elif model == "Random Forest Regressor" or model == "Logistic Regression" or model == "Linear Regression":
        st.write("Mean Squared Error: ", np.mean((y_pred - y_test) ** 2))
else:
    st.error("file not loaded")