from azure.core.credentials import AzureKeyCredential
from azure.ai.formrecognizer import FormRecognizerClient, DocumentAnalysisClient
import pandas as pd
import numpy as np
import gradio as gr
from PIL import Image
import os


from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
model_name = "deepset/bert-base-uncased-squad2"

# a) Get predictions
nlp = pipeline('question-answering', model=model_name, tokenizer=model_name)


endpoint = "https://invoice-extract.cognitiveservices.azure.com/"
key = "af03126391e141d482255941fdb16b7c"
# form_recognizer_client = FormRecognizerClient(endpoint=endpoint, credential=AzureKeyCredential(key))
document_analysis_client = DocumentAnalysisClient(
        endpoint=endpoint, credential=AzureKeyCredential(key))


def extract_check_info(img):
    f1 = Image.open(img)
    f1.save('upload.jpg',optimize=True, quality=95)
    os.remove(img)
    with open("upload.jpg", "rb") as f:
        poller = document_analysis_client.begin_analyze_document(
            "prebuilt-document", document=f)
        document_data = poller.result()
    os.remove("upload.jpg")
    content = document_data.content.strip()
    QA_input = {
    'question': 'What is the IFSC code?',
    'context': content}
    r1= nlp(QA_input)
    ifsc_code = r1['answer']
    
    QA_input = {
    'question': 'What is the A/C Number?',
    'context': content}
    r2 = nlp(QA_input)
    ac_number = r2['answer']
    
    QA_input = {
    'question': 'What is the person name?',
    'context': content}
    r3 = nlp(QA_input)
    person_name = r3['answer']
    
    return ifsc_code, ac_number, person_name




if __name__ == "__main__":

    iface = gr.Interface(fn=extract_check_info,
                    inputs=[gr.inputs.Image(type='filepath', label='Input')],
                    outputs = [gr.outputs.Textbox(type='text', label='IFSC code'),
                              gr.outputs.Textbox(type='text', label='Account Number'),
                              gr.outputs.Textbox(type='text', label='Signatory Name')],
                    title="Check information extraction demo")
    iface.launch()