import streamlit as st
import pickle
import numpy as np
import os, glob, json, sys
import pickle
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer

from src import data, utils
from src.embeddings import EmbeddingsRegressor


# load the models
with open('models/2d_ridge_roberta-suicide-regchain-pca-final.pkl', 'rb') as f:
    regressor = pickle.load(f)

model_name = 'hackathon-somos-nlp-2023/roberta-base-bne-finetuned-suicide-es'
tokenizer = SentenceTransformer(model_name)
model = EmbeddingsRegressor(tokenizer, regressor, normalize_output=True)
predict = utils.make_predict(model.predict)

# model_selector = st.sidebar.selectbox(
#     'Select model:', 
#     ['roberta', 'roberta_seq_multi', 'roberta_seq_multi_2']
# )

text_input = st.text_input('Enter your text here:')
if text_input:
    prediction = predict([text_input]).tolist()
    prediction = np.array(prediction).reshape(-1,4)
    prediction = utils.normalize(prediction)
    preds_df = data.make_task_labels_from_d(prediction, include_d=True).rename(
        columns={c:'d_'+c.replace('+','_').replace('|','_') for c in data.task_d_cols}
    )
    preds_df['b_label'] = np.clip(preds_df['b_label'], 0, 1)
    # show the dataframe 
    table = st.table(preds_df)