import streamlit as st
import numpy as np
import faiss
import pickle
from sentence_transformers import SentenceTransformer
import pandas as pd

# 제목 표시
st.title('직장 괴롭힘 참고 사례 검색과 판단')

# SentenceTransformer 모델 로드
model = SentenceTransformer('bongsoo/kpf-sbert-v1')   # 768차원

# 참고 사례 데이터프레임 로드
df = pd.read_excel('직장괴롭힘_071424.xlsx')

# 문장 임베딩
#embeddings = model.encode(df['사례'].tolist())

# 저장된 임베딩을 pickle 파일에서 로드
with open('embeddingsNorm_직장괴롭힘071424_kpfsbert.pkl', 'rb') as f:
    embeddings = pickle.load(f)
    
# 임베딩을 코사인 유사도 계산을 위해 정규화
embeddings_norm = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)

# FAISS 인덱스 생성 (내적 검색용)    
d = embeddings_norm.shape[1] # 임베딩의 차원 수를 저장:  embeddings_norm은 (문장 수, 차원 수) 형태의 2차원 배열이므로, shape[1]은 임베딩의 차원 수를 나타냄   
index = faiss.IndexFlatIP(d) # 내적(IP, Inner Product) 검색을 위한 FAISS 인덱스 생성. IndexFlatIP는 내적 기반의 유사도 검색을 위한 인덱스를 의미    
index.add(embeddings_norm) # 정규화된 임베딩을 인덱스에 추가. 검색을 위해 미리 계산된 모든 문장 임베딩을 인덱스에 추가하여 검색이 가능하도록 함

# 찾고자 하는 문장을 입력받음
query_sentence = st.text_area('검색할 문장을 입력하세요', '''예시: 회사 과장님이 회의에서 "조용히 있으라"며 야단치고, 동료들도 회식 등 모임에서 따돌림합니다''')

if st.button('분석'):
    # 입력된 문장의 임베딩
    query_embedding = model.encode([query_sentence])
    # 입력된 문장의 임베딩을 정규화
    query_embedding_norm = query_embedding / np.linalg.norm(query_embedding, axis=1, keepdims=True)
    # 유사도 검색 (상위 5개 문장)
    k = 5
    distances, indices = index.search(query_embedding_norm, k)
    # 내적 값을 코사인 유사도로 변환
    cosine_similarities = distances[0]
    # 인덱스와 유사도를 저장하는 딕셔너리 생성
    #indices_distances = {'indices': indices[0], 'cosine_similarities': cosine_similarities}
    # 결과를 저장할 데이터프레임 생성
    df_result = df.iloc[indices[0]].copy()
    df_result['유사도'] = cosine_similarities
    # 유사도 기준 내림차순으로 정렬
    df_result = df_result.sort_values(by='유사도', ascending=False)
    # 결과 표시
    st.write('분석 결과:')
    st.write(df_result)