import pickle
import pandas as pd
import streamlit as st

from datasets import Dataset

import model
from utils import check_columns, count_labels, get_download_link

# Main function to run the Streamlit app
def main():
    # Set app title
    st.title("Few Shot Learning Demo using SetFit")

    st.write("Prepare CSV file with text and label header, here is the sample file")

    df = pd.read_csv("data/sample.csv")

    # Display a link to download the file
    st.markdown(get_download_link(df), unsafe_allow_html=True)

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

    # Create file uploader
    uploaded_file = st.file_uploader("Choose a CSV file to upload", type="csv")

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

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

            # Calculate the number of instances of each label class
            label_counts = count_labels(df)
            st.write(f"Number of instances of each label class: {label_counts}")

            labels = set(df["label"].tolist())
            label_map = {label: idx for idx, label in enumerate(labels)}

            df["label"] = df["label"].map(label_map)

            dataset = Dataset.from_pandas(df)

            model_name = st.text_input("Input the model name")

            pretrained_model_options = ["all-MiniLM-L6-v2", "paraphrase-MiniLM-L3-v2"]

            pretrained_model = st.selectbox(
                "Select a pretrained model", options=pretrained_model_options
            )

            # Add Train button
            if st.button("Train"):
                # Train the model
                with st.spinner("Training model..."):
                    model_path = model.run_setfit_training(
                        session_id,
                        pretrained_model,
                        model_name,
                        dataset,
                        1,
                        10,
                    )

                st.write(f"Model checkpoint saved {model_path.split('/')[-1]}")

                label_map = {v: k for k, v in label_map.items()}
                with open(f"{model_path}/label.pkl", "wb") as f:
                    pickle.dump(label_map, f)

                st.write("Training Finished")
                st.write("Go to Validation Page")

        else:
            st.error("File must have 'text' and 'label' columns.")