import streamlit as st
from streamlit_option_menu import option_menu
from word2vec import *
import pandas as pd

st.set_page_config(page_title="Ancient Greek Word2Vec", layout="centered")

# Horizontal menu
active_tab = option_menu(None, ["Nearest neighbours", "Cosine similarity", "3D graph", 'Dictionary'], 
    menu_icon="cast", default_index=0, orientation="horizontal")

# Nearest neighbours tab
if active_tab == "Nearest neighbours":
    st.write("### TO DO: add description of function")
    col1, col2 = st.columns(2)
    with st.container():
        with col1:
            word = st.text_input("Enter a word", placeholder="ἀνήρ")
            
        with col2:
            time_slice = st.selectbox("Time slice", ["Archaic", "Classical", "Hellenistic", "Early Roman", "Late Roman"])
        
        models = st.multiselect("Select models to search for neighbours", ["Archaic", "Classical", "Hellenistic", "Early Roman", "Late Roman"])
        n = st.slider("Number of neighbours", 1, 50, 15)
        
        nearest_neighbours_button = st.button("Find nearest neighbours")
        
        # If the button to calculate nearest neighbours is clicked
        if nearest_neighbours_button:
            
            # Rewrite timeslices to model names: Archaic -> archaic_cbow
            if time_slice == 'Hellenistic':
                time_slice = 'hellen'
            elif time_slice == 'Early Roman':
                time_slice = 'early_roman'
            elif time_slice == 'Late Roman':
                time_slice = 'late_roman'
            
            time_slice = time_slice.lower() + "_cbow"
            
            
            
            # Check if all fields are filled in
            if validate_nearest_neighbours(word, time_slice, n, models) == False:
                st.error('Please fill in all fields')
            else:
                # Rewrite models to list of all loaded models
                models = load_selected_models(models)
                
                nearest_neighbours = get_nearest_neighbours(word, time_slice, n, models)
                
                df = pd.DataFrame(nearest_neighbours, columns=["Word", "Time slice", "Similarity"])
                st.table(df)
                
                
                
        
    
# Cosine similarity tab
elif active_tab == "Cosine similarity":
    with st.container():
        st.write("Cosine similarity tab")

# 3D graph tab
elif active_tab == "3D graph":
    with st.container():
        st.write("3D graph tab")

# Dictionary tab
elif active_tab == "Dictionary":
    with st.container():
        st.write("Dictionary tab")