import streamlit as st
import torch
import pickle
import time
import pandas as pd
from iteration_utilities import unique_everseen
from sentence_transformers import util
from loader import bi_encoder, cross_encoder, df, job_corpus_ecoded, job_corpus


def jobsearch(query,df, top_k=100):
  #print("Answer by NinjaBot : ")
  ans = []
  question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
  hits = util.semantic_search(question_embedding, job_corpus_ecoded, top_k=top_k)
  hits = hits[0]
  cross_inp = [[query, job_corpus[hit['corpus_id']]] for hit in hits]
  cross_scores = cross_encoder.predict(cross_inp)
  for idx in range(len(cross_scores)):
      hits[idx]['cross-score'] = cross_scores[idx]

  hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
  #indexes = []
  search_result = []
  for idx, hit in enumerate(hits[0:10]):
      obj = {}
      ans.append(job_corpus[hit['corpus_id']])
      #indexes.append(job_corpus.index(job_corpus[hit['corpus_id']]))
      obj['title'] = df.at[job_corpus.index(job_corpus[hit['corpus_id']]),'title']
      obj['link'] = df.at[job_corpus.index(job_corpus[hit['corpus_id']]),'url']
      search_result.append(obj)
  final_search_result = list(unique_everseen(search_result))
  return final_search_result
  #return df.at[indexes[0],'title'],df.at[indexes[1],'title'],df.at[indexes[2],'title'],df.at[indexes[3],'title'],df.at[indexes[4],'title']
  #return ans[0],ans[1],ans[2],ans[3],ans[4]

def main():
    if 'submitted' not in st.session_state:
        st.session_state.submitted = False

    def callback():
      st.session_state.submitted = True

    st.title('Job Search Engine 💼')
    st.text("")
    st.text("")
    query = st.text_input('Enter your job query here ! ')
    if (st.button("Search", on_click=callback) and query) :
        with st.spinner('Fetching the best jobs for you!...'):
            #time.sleep(10)
            result = jobsearch(query, df)
        #result = jobsearch(query, df)
        st.success('NinjaBot : Here are a few suggestions')
        #st.write(f"This is the query : {query}")  
        st.write(result)
main()