import cv2
import tqdm
from PIL import Image, ImageDraw

from model import CTPN, SVTR
from utils import get_images_from_path, img_read

import IPython
import time

def init_acl(device_id=0):
    acl.init()
    ret = acl.rt.set_device(device_id)
    if ret:
        raise RuntimeError(ret)
    context, ret = acl.rt.create_context(device_id)
    if ret:
        raise RuntimeError(ret)
    print('Init ACL Successfully')
    return context


def deinit_acl(context, device_id=0):
    ret = acl.rt.destroy_context(context)
    if ret:
        raise RuntimeError(ret)

    ret = acl.rt.reset_device(device_id)
    if ret:
        raise RuntimeError(ret)
    print('ACL resources were released successfully.')



def get_images_from_path(img_path):
    img_list = []
    if os.path.isfile(img_path):
        img_list.append(img_path)
    if os.path.isdir(img_path):
        for file in os.listdir(img_path):
            img_list.append(os.path.join(img_path, file))
    return img_list


def img_read(path):
    path = os.path.realpath(path)
    img = cv2.imread(path, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError(f"cannot decode file:{path}")
    return img



CTPN_MODEL_PATH = './ctpn.om'
SVTR_MODEL_PATH = './svtr.om'
SVTR_DICT_PATH = './ppocr_keys_v1.txt'
DEVICE_ID = 0

context = init_acl(DEVICE_ID)
det_model = CTPN(model_path=CTPN_MODEL_PATH, device_id=DEVICE_ID)
rec_model = SVTR(model_path=SVTR_MODEL_PATH, dict_path=SVTR_DICT_PATH,
                     device_id=DEVICE_ID)
IMAGE_PATH = './sample.png'

if not os.path.exists('infer_result'):
    os.makedirs('infer_result')
ans = {}

start = time.time()
img_src = img_read(IMAGE_PATH)
basename = os.path.basename(IMAGE_PATH)
print(f'start infer image: {basename}')
name, ext = os.path.splitext(basename)
image_h, image_w = img_src.shape[:2]

det_input_tensor = det_model.preprocess(img_src)

output = det_model.infer(det_input_tensor)

bboxes = det_model.postprocess(output)
im = Image.open(IMAGE_PATH)
draw = ImageDraw.Draw(im)
ans['image_name'] = basename
ans['result'] = []
for bbox in bboxes:
    bbox_detail = {}
    x1 = int(bbox[0] / det_model.model_width * image_w)
    y1 = int(bbox[1] / det_model.model_height * image_h)
    x2 = int(bbox[2] / det_model.model_width * image_w)
    y2 = int(bbox[3] / det_model.model_height * image_h)
    draw.line([(x1, y1), (x1, y2), (x2, y2), (x2, y1), (x1, y1)], fill='red', width=2)
    bbox_detail['bbox'] = [x1, y1, x1, y2, x2, y2, x2, y1]
    res  = ','.join(map(str,bbox_detail['bbox']))
    print(f'det result: {res}')
    crop_img = img_src[y1:y2, x1:x2]

    rec_input_tensor = rec_model.preprocess(crop_img)
    rec_output = rec_model.infer(rec_input_tensor)
    bbox_detail['text'] = rec_model.postprocess(rec_output)
    print(f'rec result: {bbox_detail["text"]}')
    ans['result'].append(bbox_detail)

im.save(os.path.join('infer_result', name + '_res' + ext))
end = time.time()
print(f"infer time: {str(end - start)}s")

IPython.display.Image(filename=os.path.join('infer_result', name + '_res' + ext))

det_model.deinit()
rec_model.deinit()
deinit_acl(context, DEVICE_ID)