import os, sys
cur_dir = os.path.dirname(__file__)
KIE_PATH = os.path.join(os.path.dirname(cur_dir), "sdsvkie")
TD_PATH = os.path.join(os.path.dirname(cur_dir), "sdsvtd")
TR_PATH = os.path.join(os.path.dirname(cur_dir), "sdsvtr")
sys.path.append(KIE_PATH)
sys.path.append(TD_PATH)
sys.path.append(TR_PATH)

from sdsvkie import Predictor
from .AnyKey_Value.anyKeyValue import load_engine, Predictor_KVU
import cv2
import numpy as np
import urllib

model = Predictor(
        cfg = "/models/Kie_invoice_ap/06062023/config.yaml", # TODO: Better be scalable
        device = "cuda:0",
        weights = "/models/Kie_invoice_ap/06062023/best" # TODO: Better be scalable
    )

class_names = ['others', 'title', 'key', 'value', 'header']
save_dir = os.path.join(cur_dir, "AnyKey_Value/visualize/test")

predictor, processor = load_engine(exp_dir="/models/Kie_invoice_fi/key_value_understanding-20230627-164536", 
                                       class_names=class_names, 
                                       mode=3)

def predict_fi(page_numb, image_url):
    """
    module predict function

    Args:
        image_url (str): image url

    Returns:
    example output:
         "data": {
            "document_type": "invoice",
            "fields": [
            {
                "label": "Invoice Number",
                "value": "INV-12345",
                "box": [0, 0, 0, 0],
                "confidence": 0.98
            },
            ...
            ]
        }
    dict: output of model
    """
    req = urllib.request.urlopen(image_url)
    arr = np.asarray(bytearray(req.read()), dtype=np.uint8)
    img = cv2.imdecode(arr, -1)
    # img = cv2.imread(image_url)
    
    # Phan cua LeHoang
    out = model(img)
    output = out["end2end_results"]
    output_kie = {
                field_name: field_item['value'] for field_name, field_item in output.items()
            }
    
    
    kvu_result, _ = Predictor_KVU(image_url, save_dir, predictor, processor)
    # if kvu_result['imei_number'] == None and kvu_result['serial_number'] == None:
    return kvu_result, output_kie

if __name__ == "__main__":
    image_url = "/mnt/hdd2T/dxtan/TannedCung/OCR/workspace/Kie_Invoice_AP/tmp_image/{image_url}.jpg"
    output = predict_fi(0, image_url)