import os
import streamlit as st
import logging
import pandas as pd
import numpy as np
import pickle
import numpy as np
from bertopic import BERTopic
from umap import UMAP
from cryptography.fernet import Fernet
from sentence_transformers import SentenceTransformer, util
from pdb import set_trace as trace


def load_db():

    with st.spinner('Loading pre-computed embeddings...'):

        file_path = "./surveybot3000.enc"
        logging.info(f"Loading data from {file_path}!")

        with open(file_path, 'rb') as f:
            encrypted_data = f.read()

        encryption_key = st.session_state['encryption_key']

        if isinstance(encryption_key, str):
            encryption_key = encryption_key.encode('utf-8')

        cipher = Fernet(encryption_key)
        decrypted_df = cipher.decrypt(encrypted_data)
        st.session_state['db'] = pickle.loads(decrypted_df)

        st.session_state['is_authenticated'] = True
        logging.info(f"Loaded {file_path}!")


def load_model():

    with st.spinner('Loading the model...'):


        model_path = os.environ.get('surveybot3000_remote_path') or os.getenv('surveybot3000_path')
        logging.info(f"Loading model from {model_path}!")

        auth_token = os.environ.get('read_models') or True
        logging.info(f"Using auth token: {'***' if auth_token else 'None'}")

        st.session_state['model'] = SentenceTransformer(
            model_name_or_path=model_path,
            token=auth_token,
            trust_remote_code=True,
            tokenizer_kwargs={
                'use_fast': False,
                'legacy': True
            }
        )

        logging.info(f"Loaded model!")


def search():

    with st.spinner('Searching the synthetic net...'):
        query_embeddings = st.session_state['model'].encode(
            sentences=st.session_state['search_query']).mean(axis=0)

        item_embeddings = np.vstack(
            st.session_state['db']['ItemStemEmbeddings'])
        query_scores = util.cos_sim(
            a=np.array(query_embeddings),
            b=item_embeddings
        ).squeeze()

        st.session_state['search_results'] = pd.DataFrame({
            'Match': query_scores,
            'Scale': st.session_state['db']['ScaleName'],
            'Instrument': st.session_state['db']['InstrumentName'],
            'Reference': st.session_state['db']['psyctest_doi'],
        }).sort_values(by='Match', ascending=False)