import os
import pickle
import pandas as pd
import streamlit as st

import model

from utils import check_columns


# Function to validate the trained model with a new uploaded CSV file
def main():

    st.title("Model Validation")

    # Display the session ID
    # st.write(f"Session ID: {st.session_state.key}")
    session_id = st.session_state.key

    if not os.path.isdir(f"models/{session_id}"):
        st.write("Model is not available")
        st.stop()

    model_options = [model_name for model_name in os.listdir(f"models/{session_id}")]

    models = {
        model_name: os.path.abspath(os.path.join(f"models/{session_id}", model_name))
        for model_name in model_options
    }

    model_name = st.selectbox("Select a model", options=model_options)

    # Create file uploader for validation CSV file
    validation_file = st.file_uploader(
        "Choose a CSV file to validate the model", type="csv"
    )

    # Check if validation file was uploaded
    if validation_file is not None:
        # Read CSV file into pandas DataFrame
        validation_df = pd.read_csv(validation_file)

        # Check if DataFrame has expected columns
        if check_columns(validation_df):
            # Display DataFrame as a table
            st.write(validation_df)

            # Create pipeline for text classification using the trained model
            classifier = model.create_classifier(models[model_name])

            with open(f"{models[model_name]}/label.pkl", "rb") as f:
                label_map = pickle.load(f)

            results = classifier(validation_df["text"].tolist())

            # Predict labels for validation DataFrame
            validation_df["predicted_label"] = [
                label_map[result.item()] for result in results
            ]

            # Display validation DataFrame with predicted labels
            st.write("Validation results:")
            st.write(validation_df)
        else:
            st.error("Validation file must have 'text' and 'label' columns.")