# -*- coding: utf-8 -*-
import json
import pandas as pd
import numpy as np
import torch
import streamlit as st
from transformers import AutoTokenizer, AutoModelForTokenClassification

st.set_page_config(
    page_title="NER 기반 민감정보 식별", layout="wide", initial_sidebar_state="expanded"
)

@st.cache
def load_model(model_name):
    model = AutoModelForTokenClassification.from_pretrained(model_name)
    return model


st.title("🔒 NER 기반 민감정보 식별기")
st.write("문장을 입력하시고, CTRL+Enter(CMD+Enter)를 누르세요 🤗")

tokenizer = AutoTokenizer.from_pretrained("klue/roberta-base")
model = load_model("QuoQA-NLP/konec-privacy")

model.eval()


default_value = "영진님, 당뇨 검사한 거 결과 나오셨어요."

src_text = st.text_area(
    "검사하고 싶은 문장을 입력하세요.",
    default_value,
    height=300,
    max_chars=150,
)


def yield_df(default_value):
  tokenized = tokenizer.encode(default_value)
  print(tokenized)

  output = model(input_ids=torch.tensor([tokenized]))
  logits = output.logits
  print(logits.size())

  # get prediction for each tokens for 17 classes
  pred = logits.argmax(-1).squeeze().numpy()
  print(pred)

  class_map = {
  "B-ADD": 0,
  "I-ADD": 1,
  "B-DN": 2,
  "I-DN": 3,
  "B-DT": 4,
  "I-DT": 5,
  "B-LC": 6,
  "I-LC": 7,
  "B-OG": 8,
  "I-OG": 9,
  "B-PS": 10,
  "I-PS": 11,
  "B-QT": 12,
  "I-QT": 13,
  "B-RL": 14,
  "I-RL": 15,
  "O": 16
  }

  class_map_inverted = {v: k for k, v in class_map.items()}

  # decode prediction
  class_decoded = [class_map_inverted[p] for p in pred]
  print(class_decoded)

  label_map = {
  "ADD": "주소 정보",
  "DN": "질환 정보",
  "DT": "날짜 정보",
  "LC": "장소 정보",
  "OG": "기관 정보",
  "PS": "인명/별명 정보",
  "QT": "수량 정보",
  "RL": "관계 정보",
  "O": "비민감 정보"
  }


  # pair tokens with prediction
  tokenized_text = tokenizer.convert_ids_to_tokens(tokenized)
  list_result = []
  for token, pred in zip(tokenized_text, class_decoded):
    splitted_pred = pred.split("-")
    pred_class = splitted_pred[-1]
    label = label_map[pred_class]
    # print with 10 characters with spaces divided with |
    result = {"형태소":token, "예상 라벨":label}
    list_result.append(result)
  
  df = pd.DataFrame(list_result)
  # remove first and last row
  df = df.iloc[1:-1]
  return df

def convert_df(df:pd.DataFrame):
     return df.to_csv(index=False).encode('utf-8')

def convert_json(df:pd.DataFrame):
    result = df.to_json(orient="index")
    parsed = json.loads(result)
    json_string = json.dumps(parsed)
    #st.json(json_string, expanded=True)
    return json_string



filtering_map = {
  "주소 정보": "[주소]",
  "질환 정보": "[질환]",
  "날짜 정보": "[날짜]",
  "장소 정보": "[장소]",
  "기관 정보": "[기관]",
  "인명/별명 정보": "[이름]",
  "수량 정보": "[수량]",
  "관계 정보": "[관계]",
  "비민감 정보": "[비민감]"
  }

if src_text == "":
    st.warning("Please **enter text** for translation")
else:
    df_result = yield_df(src_text)
    st.markdown("### 필터링 된 문장")
    
    display_result = ""
    for index, row in df_result.iterrows():
        token_info = row["형태소"]
        label_info = row["예상 라벨"]
        if label_info != "비민감 정보":
            token_info = filtering_map[label_info]
        
        if "##" in token_info:
            token_info = token_info.replace("##", "")
        else:
            token_info = " " + token_info
        display_result += token_info
    
    st.write(display_result)

    st.markdown("### 분류된 단어들")
    st.header("")
    cs, c1, c2, c3, cLast = st.columns([0.75, 1.5, 1.5, 1.5, 0.75])
    
    st.table(df_result)
    with c1:
        #csvbutton = download_button(results, "results.csv", "📥 Download .csv")
        csvbutton = st.download_button(label="📥 csv로 다운로드", data=convert_df(df_result), file_name= "results.csv", mime='text/csv', key='csv')
    with c2:
        #textbutton = download_button(results, "results.txt", "📥 Download .txt")
        textbutton = st.download_button(label="📥 txt로 다운로드", data=convert_df(df_result), file_name= "results.text", mime='text/plain',  key='text')
    with c3:
        #jsonbutton = download_button(results, "results.json", "📥 Download .json")
        jsonbutton = st.download_button(label="📥 json으로 다운로드", data=convert_json(df_result), file_name= "results.json", mime='application/json',  key='json')



with st.expander("(주) 쿼카에이아이 데모 사사 관련", expanded=True):

  st.write(
      """     
      해당 데모는 2022년도 과학기술정보통신부의 재원으로 정보통신산업진흥원의 지원을 받아 수행된 연구임
      (과제번호: A1504-22-1005)
      """
  )