import streamlit as st
from utils.data_utils import load_dataset
from utils.model_utils import load_model, train_model
from utils.ui_utils import display_model_selection, display_training_method_selection, display_dataset_upload, display_model_details, display_training_progress

# Set the page title
st.title("FallnAI Trainer")

# Display the model selection dropdown
model = display_model_selection()

# Display the training method selection dropdown
training_method = display_training_method_selection()

# Display the dataset upload
uploaded_file = display_dataset_upload()

# Display the model details
display_model_details(model)

# Add a button to start training
if st.button("Train Model"):
    # Load the model
    model = load_model(model)

    # Load the dataset
    data = load_dataset(uploaded_file)

    # Train the model
    trained_model = train_model(model, data, training_method)

    # Display the trained model
    st.write("Trained Model:", trained_model)

    # Display the training progress
    display_training_progress(trained_model)