import streamlit as st
from sklearn.linear_model import LinearRegression
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

st.header("Prediction: Regression")

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

    with st.form("regression_form"):
        st.subheader("Linear Regression Parameters")
        data_name = st.multiselect("Features", data.select_dtypes(include="number").columns)
        target_name = st.selectbox("Target", data.select_dtypes(include="number").columns)
        st.form_submit_button('Train and Predict')

    if data_name and target_name:
        X = data[data_name]
        y = data[target_name]

        model = LinearRegression()
        model.fit(X, y)
        
        st.subheader("Enter values for prediction")
        pred_values = [st.number_input(f"Value for {feature}", value=0.0) for feature in data_name]
        prediction = model.predict(pd.DataFrame([pred_values], columns=data_name))

        st.write("Prediction:", prediction[0])

        fig = plt.figure()
        dataframe_sorted = pd.concat([X, y], axis=1).sort_values(by=data_name)

        if len(data_name) == 1:
            X = dataframe_sorted[data_name[0]]
            y = dataframe_sorted[target_name]

            prediction_array_y = [
                model.predict(pd.DataFrame([[dataframe_sorted[data_name[0]].iloc[i]]], columns=data_name))[0]
                for i in range(dataframe_sorted.shape[0])
            ]

            plt.scatter(dataframe_sorted[data_name[0]], dataframe_sorted[target_name], color='b')
            plt.scatter(dataframe_sorted[data_name[0]], prediction_array_y, color='r')
        else:
            ax = fig.add_subplot(111, projection='3d')

            prediction_array_y = [
                model.predict(pd.DataFrame([[dataframe_sorted[data_name[0]].iloc[i], dataframe_sorted[data_name[1]].iloc[i]]], columns=data_name))[0]
                for i in range(dataframe_sorted.shape[0])
            ]

            ax.scatter(dataframe_sorted[data_name[0]], dataframe_sorted[data_name[1]], dataframe_sorted[target_name], color='b')
            ax.scatter(dataframe_sorted[data_name[0]], dataframe_sorted[data_name[1]], prediction_array_y, color='r')

        st.pyplot(fig)


else:
    st.error("File not loaded")