import streamlit as st
import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../backend')))
from regression_strategy import perform_regression, make_prediction

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, key="regression_features")
        target_name = st.selectbox("Target", data.select_dtypes(include="number").columns, key="regression_target")
        submitted = st.form_submit_button('Train and Predict')

    if submitted and data_name and target_name:
        try:
            model = perform_regression(data, data_name, target_name)
            st.session_state.regression_model = model
            st.session_state.regression_features_selected = data_name
            st.session_state.regression_target_selected = target_name
        except ValueError as e:
            st.error(e)

    if "regression_model" in st.session_state:
        st.subheader("Enter values for prediction")
        input_values = [st.number_input(f"Value for {feature}", value=0.0, key=f"regression_input_{feature}") for feature in st.session_state.regression_features_selected]
        prediction = make_prediction(st.session_state.regression_model, st.session_state.regression_features_selected, input_values)

        st.write("Prediction:", prediction)
else:
    st.error("File not loaded")