from sdsvkie import Predictor
import sys, os
cur_dir = os.path.dirname(__file__)
# sys.path.append("cope2n-ai-fi/Kie_Invoice_AP") # Better be relative
from .AnyKey_Value.anyKeyValue import load_engine, Predictor_KVU
import cv2
import numpy as np
import urllib
import random

# model = Predictor(
#         cfg = "/cope2n-ai-fi/models/Kie_invoice_ap/config.yaml",
#         device = "cuda:0",
#         weights = "/cope2n-ai-fi/models/Kie_invoice_ap/best"
#     )

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

predictor, processor = load_engine(exp_dir=os.path.join(cur_dir, "AnyKey_Value/experiments/key_value_understanding-20231003-171748"), 
                                       class_names=class_names, 
                                       mode=3)

def predict(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"]
    
    
    #Phan cua Tuan
    kvu_result = Predictor_KVU(image_url, save_dir, predictor, processor)
    output_dict = {
        "document_type": "invoice",
        "fields": []
    }
    for key in kvu_result.keys():
        field = {
                "label": key,
                "value": kvu_result[key],
                "box": [0, 0, 0, 0],
                "confidence": random.uniform(0.9, 1.0),
                "page": page_numb
            }
        output_dict['fields'].append(field)
    return output_dict
        
    # if kvu_result['imei_number'] == None and kvu_result['serial_number'] == None:
    #     output_dict = {
    #         "document_type": "invoice",
    #         "fields": []
    #     }
    #     for key in output.keys():
    #         field = {
    #             "label": key if key != "id" else "Receipt Number",
    #             "value": output[key]['value'] if output[key]['value'] else "",
    #             "box": output[key]['box'],
    #             "confidence": output[key]['conf'],
    #             "page": page_numb
    #         }
    #         output_dict['fields'].append(field)
    #     table = kvu_result['table']
    #     field_table = {
    #         "label": "table",
    #         "value": table,
    #         "box": [0, 0, 0, 0],
    #         "confidence": 0.98,
    #         "page": page_numb
    #     }
    #     output_dict['fields'].append(field_table)
    #     return output_dict
    
    # else:
    #     output_dict = {
    #         "document_type": "KSU",
    #         "fields": []
    #     }
    #     # for key in output.keys():
    #     #     field = {
    #     #         "label": key if key != "id" else "Receipt Number",
    #     #         "value": output[key]['value'] if output[key]['value'] else "",
    #     #         "box": output[key]['box'],
    #     #         "confidence": output[key]['conf'],
    #     #         "page": page_numb
    #     #     }
    #     #     output_dict['fields'].append(field)

    #     # Serial Number
    #     serial_number = kvu_result['serial_number']
    #     field_serial = {
    #         "label" : "serial_number",
    #         "value": serial_number,
    #         "box": [0, 0, 0, 0],
    #         "confidence": 0.98,
    #         "page": page_numb
    #     }
    #     output_dict['fields'].append(field_serial)

    #     # IMEI Number
    #     imei_number = kvu_result['imei_number']
    #     if imei_number == None:
    #         return output_dict
    #     if imei_number != None:
    #         for i in range(len(imei_number)):
    #             field_imei = {
    #                 "label": "imei_number_{}".format(i+1),
    #                 "value": imei_number[i],
    #                 "box": [0, 0, 0, 0],
    #                 "confidence": 0.98,
    #                 "page": page_numb
    #             }
    #             output_dict['fields'].append(field_imei)
        
    #         return output_dict

if __name__ == "__main__":
    image_url = "/root/thucpd/20230322144639VUzu_16794962527791962785161104697882.jpg"
    output = predict(0, image_url)