from transformers import BertTokenizerFast
from gradio_pdf import PDF
from src.bert import *
from src.legalNER import *
import gradio as gr
from pathlib import Path

dir_ = Path(__file__).parent

ids_to_labels = {0: 'B_ADVO', 1: 'B_ARTV', 2: 'B_CRIA', 3: 'B_DEFN', 4: 'B_JUDG', 5: 'B_JUDP', 6: 'B_PENA', 7: 'B_PROS', 8: 'B_PUNI', 9: 'B_REGI', 10: 'B_TIMV', 11: 'B_VERN', 12: 'I_ADVO', 13: 'I_ARTV', 14: 'I_CRIA', 15: 'I_DEFN', 16: 'I_JUDG', 17: 'I_JUDP', 18: 'I_PENA', 19: 'I_PROS', 20: 'I_PUNI', 21: 'I_REGI', 22: 'I_TIMV', 23: 'I_VERN', 24: 'O'}
indolem = 'indolem/indobert-base-uncased'
indonlu = 'indobenchmark/indobert-base-p2'
model_indolem = BertModel(indolem, len(ids_to_labels))
model_indonlu = BertModel(indonlu, len(ids_to_labels))
tokenizer_indolem = BertTokenizerFast.from_pretrained(indolem)
tokenizer_indonlu = BertTokenizerFast.from_pretrained(indonlu)

def text_extraction(text, model, progress=gr.Progress()):
  if model == 'IndoBERT (IndoLEM)':
    use_model = model_indolem
    use_tokenizer = tokenizer_indolem

  else:
    use_model = model_indonlu
    use_tokenizer = tokenizer_indonlu

  legalner =  LegalNER(use_model, use_tokenizer, ids_to_labels, model)
  entitas = legalner.predict(text)
  new_text = legalner.tokenizer_decode

  return {"text": new_text, "entities": entitas}

def pdf_extraction(doc, model, progress=gr.Progress()):
  if model == 'IndoBERT (IndoLEM)':
    use_model = model_indolem
    use_tokenizer = tokenizer_indolem

  else:
    use_model = model_indonlu
    use_tokenizer = tokenizer_indonlu

  legalner =  LegalNER(use_model, use_tokenizer, ids_to_labels, model)

  return legalner.predict(doc)


with gr.Blocks() as ner:
  gr.Markdown("#Sistem Ekstraksi Informasi Dokumen Putusan Hukum")
  gr.Markdown("## Uji Coba Model dengan Potongan Kalimat")
  # Input Text
  with gr.Row():
    with gr.Column(scale=2):
      text = gr.Textbox(label="Text")
      model_text = gr.Dropdown(['IndoBERT (IndoLEM)', 'IndoBERT (IndoNLU)'], label='Model', value='IndoBERT (IndoLEM)', info='Pilih Model yang ingin digunakan *Default : IndoBERT (IndoLEM)')
      button_text = gr.Button(value="Predict", variant='primary')
      gr.ClearButton(text, value='Reset')
    with gr.Column(scale=3):
      output_text = gr.HighlightedText(label="Output Text")
    
    button_text.click(fn=text_extraction, inputs=[text, model_text], outputs=output_text, api_name="text")
  
  gr.Markdown("## Contoh Inputan Potongan Kalimat")
  gr.Examples(
    examples=[   
        ["PUTUSAN . NOMOR : 187 / Pid . Sus / 2014 / PN . JKT . TIM . DEMI KEADILAN BERDASARKAN KETUHANAN YANG MAHA ESA . MENUNTUT : 1 Menyatakan terdakwa AGNES TRI AHADI Als AGNES telah terbukti secara sah dan meyakinkan bersalah melakukan tindak pidana Narkotika memiliki , menyimpan , menguasai , atau menyediakan Narkotika golongan I bukan tanaman sebagaimana didakwakan dalam dakwaan kedua yaitu melanggar ketentuan unsure pasal 112 ayat ( 1 ) UURI No . 35 tahun 2009 tentang Narkotika ;", "IndoBERT (IndoLEM)"],
        ["MENUNTUT : 1 Menyatakan terdakwa AGNES TRI AHADI Als AGNES telah terbukti secara sah dan meyakinkan bersalah melakukan tindak pidana Narkotika memiliki , menyimpan , menguasai , atau menyediakan Narkotika golongan I bukan tanaman sebagaimana didakwakan dalam dakwaan kedua yaitu melanggar ketentuan unsure pasal 112 ayat ( 1 ) UURI No . 35 tahun 2009 tentang Narkotika ;", "IndoBERT (IndoNLU)"],
        ["PUTUSAN Nomor 77/Pid.B/2023/PN Jkt.Pst DEMI KEADILAN BERDASARKAN KETUHANAN YANG MAHA ESA Pengadilan Negeri Jakarta Pusat yang mengadili perkara pidana dengan acara pemeriksaan biasa dalam tingkat pertama menjatuhkan putusan sebagai berikut dalam perkara Terdakwa : 1.	Nama lengkap	: Arif Bin Santung", "IndoBERT (IndoLEM)"],
    ],
    inputs=[text, model_text],
    outputs=output_text,
    fn=text_extraction,
    )

  gr.Markdown("## Ekstrak Entitas pada Dokumen Putusan Hukum")
  # Input PDF
  with gr.Row():
    with gr.Column(scale=2):
      doc = PDF(label="Document")
      model_pdf = gr.Dropdown(['IndoBERT (IndoLEM)', 'IndoBERT (IndoNLU)'], label='Model',value='IndoBERT (IndoLEM)', info='Pilih Model yang ingin digunakan *Default : IndoBERT (IndoLEM)')
      button_pdf = gr.Button(value="Extract", variant='primary')
      gr.ClearButton(doc, value="Reset")
    
    with gr.Column(scale=3):
      output_pdf = gr.Textbox(label="Output PDF")
  
  button_pdf.click(fn=pdf_extraction, inputs=[doc, model_pdf], outputs=output_pdf, api_name="pdf")
  
  gr.Examples(
    examples=[[str(dir_ /"data/165_Pdt.P_2023_PN_Bkl.pdf")],
              [str(dir_ /'data/162_Pid.Sus_2023_PN_Bkl.pdf')],
              [str(dir_ /'data/164_Pid.Sus_2023_PN_Bkl.pdf')],
              [str(dir_ /'data/167_Pid.Sus_2023_PN_Bkl.pdf')],
              [str(dir_ /'data/168_Pid.Sus_2023_PN_Bkl.pdf')],
              [str(dir_ /'data/169_Pid.Sus_2023_PN_Bkl.pdf')],
    ],
    inputs=[doc],
    outputs=output_pdf,
    fn=pdf_extraction,
    )

if __name__ == "__main__":
  ner.launch()