import pandas as pd
from transformers import pipeline

import PIL
from PIL import Image
from PIL import ImageDraw
import gradio as gr
import torch
import easyocr
import omegaconf
import cv2
import json

from vietocr.vietocr.tool.predictor import Predictor
from vietocr.vietocr.tool.config import Cfg

# Configure of VietOCR
config = Cfg.load_config_from_name('vgg_transformer')
# config = Cfg.load_config_from_file('vietocr/config.yml')
# config['weights'] = '/Users/bmd1905/Desktop/pretrain_ocr/vi00_vi01_transformer.pth'

config['cnn']['pretrained'] = True
config['predictor']['beamsearch'] = True
config['device'] = 'cpu' # mps

recognitor = Predictor(config)
model_name = "microsoft/xdoc-base-squad2.0"
nlp = pipeline('question-answering', model=model_name)

def query(doc, labels):
    questions = labels.split(", ")
    result={}
    for question in questions:
        QA_input = {
                    'question': question,
                    'context': doc
                }
        print(QA_input)
        res= nlp(QA_input)
        print(res)
        value = res['answer']
        result[question]=value
    return result
    
def draw_boxes(image, bounds, color='yellow', width=2):
    draw = ImageDraw.Draw(image)
    for bound in bounds:
        p0, p1, p2, p3 = bound[0]
        draw.line([*p0, *p1, *p2, *p3, *p0], fill=color, width=width)
    return image

def inference(filepath, lang, labels):
    img = cv2.imread(filepath)
    width, height, _ = img.shape
    reader = easyocr.Reader(lang)
    bounds = reader.readtext(filepath)
    new_bounds=[]
    text=''
    for (bbox, text, prob) in bounds:
        (tl, tr, br, bl) = bbox
        tl = (int(tl[0]), int(tl[1]))
        tr = (int(tr[0]), int(tr[1]))
        br = (int(br[0]), int(br[1]))
        bl = (int(bl[0]), int(bl[1]))

        min_x = min(tl[0], tr[0], br[0], bl[0])
        min_x = max(0, min_x)
        max_x = max(tl[0], tr[0], br[0], bl[0])
        max_x = min(width-1, max_x)
        min_y = min(tl[1], tr[1], br[1], bl[1])
        min_y = max(0, min_y)
        max_y = max(tl[1], tr[1], br[1], bl[1])
        max_y = min(height-1, max_y)
        # crop the region of interest (ROI)
        
        cropped_image = img[min_y:max_y,min_x:max_x] # crop the image
        cropped_image = Image.fromarray(cropped_image)
        out = recognitor.predict(cropped_image)
        text = text + '\t' + out
    result = query(text, labels)
    jsonText = json.dumps(result)
    im = PIL.Image.open(filepath)
    draw_boxes(im, bounds)
    im.save('result.jpg')
    return ['result.jpg', jsonText]

title = 'EasyOCR'
description = 'Gradio demo for EasyOCR. EasyOCR demo supports 80+ languages.To use it, simply upload your image and choose a language from the dropdown menu, or click one of the examples to load them. Read more at the links below.'
article = "<p style='text-align: center'><a href='https://www.jaided.ai/easyocr/'>Ready-to-use OCR with 80+ supported languages and all popular writing scripts including Latin, Chinese, Arabic, Devanagari, Cyrillic and etc.</a> | <a href='https://github.com/JaidedAI/EasyOCR'>Github Repo</a></p>"
css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"
choices = [
    "vi"
]
gr.Interface(
    inference,
    [gr.inputs.Image(type='filepath', label='Input'),gr.inputs.CheckboxGroup(choices, type="value", default=['vi'], label='language'), gr.inputs.Textbox(label='Labels',default='Người nộp, Tiếp nhận hồ sơ của')],
    [gr.outputs.Image(type='pil', label='Output'), gr.outputs.Textbox(label='Json')],
    title=title,
    description=description,
    article=article,
    css=css,
    examples=['giaytiepnhan.jpg','giaytiepnhan2.jpg']
    enable_queue=True
    ).launch(debug=True)