import streamlit as st
from tools import *
import pandas as pd
import time
from streamlit_extras.stateful_button import button

_, file_tree_ph = sidebar(globals())

st.title("Finalize model")

exp_id = st.session_state.get('exp_id')
if exp_id is None:
    st.write("Please select an experiment first.")
    st.stop()

use_synthetic_data = st.session_state.get('use_synth_data')
df = pd.read_csv(
    PATH_TO_GEN_DATASET if use_synthetic_data else PATH_TO_TRAIN_DATASET)

with st.container(border=True):
	chosen_target = st.selectbox(
		'Choose the Target Column', df.columns, key="preserve_chosen_target", index=len(df.columns)-1)
	chosen_ignore = st.multiselect('Choose the Ignore Columns',
                                df.columns, key="preserve_chosen_ignore")
	st.selectbox('Choose the Problem Type', [
		'classification', 'regression'], key="preserve_chosen_problem", index=identify_problem_type(df) == 'regression')

	exp = create_experiment(
		st.session_state.get('preserve_chosen_problem'), df, chosen_target, chosen_ignore, exp_id=exp_id)

	models = exp.models()
	chosen_model = st.selectbox(
		'Choose the Model', options=models.index, format_func=lambda x: models.loc[x, 'Name'], key="preserve_chosen_model")

if st.button('Finalize model', disabled=chosen_model is None or chosen_target is None or chosen_ignore is None):
	df = pd.read_csv(PATH_TO_TRAIN_DATASET)
	final_model = exp.create_model(chosen_model, verbose=False)
	exp.finalize_model(final_model)
	exp.save_model(
		final_model, f'{PATH_TO_MODELS}/cl_{chosen_model}'+("_gen" if use_synthetic_data else ""))
	st.success("Model finalized successfully")
	update_file_tree(file_tree_ph, globals())