# Copyright (c) OpenMMLab. All rights reserved.
import math
import os
import shutil
import urllib
import warnings

import cv2
import mmcv
import numpy as np
import torch
from matplotlib import pyplot as plt
from PIL import Image, ImageDraw, ImageFont

import mmocr.utils as utils


def overlay_mask_img(img, mask):
    """Draw mask boundaries on image for visualization.

    Args:
        img (ndarray): The input image.
        mask (ndarray): The instance mask.

    Returns:
        img (ndarray): The output image with instance boundaries on it.
    """
    assert isinstance(img, np.ndarray)
    assert isinstance(mask, np.ndarray)

    contours, _ = cv2.findContours(
        mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    cv2.drawContours(img, contours, -1, (0, 255, 0), 1)

    return img


def show_feature(features, names, to_uint8, out_file=None):
    """Visualize a list of feature maps.

    Args:
        features (list(ndarray)): The feature map list.
        names (list(str)): The visualized title list.
        to_uint8 (list(1|0)): The list indicating whether to convent
            feature maps to uint8.
        out_file (str): The output file name. If set to None,
            the output image will be shown without saving.
    """
    assert utils.is_type_list(features, np.ndarray)
    assert utils.is_type_list(names, str)
    assert utils.is_type_list(to_uint8, int)
    assert utils.is_none_or_type(out_file, str)
    assert utils.equal_len(features, names, to_uint8)

    num = len(features)
    row = col = math.ceil(math.sqrt(num))

    for i, (f, n) in enumerate(zip(features, names)):
        plt.subplot(row, col, i + 1)
        plt.title(n)
        if to_uint8[i]:
            f = f.astype(np.uint8)
        plt.imshow(f)
    if out_file is None:
        plt.show()
    else:
        plt.savefig(out_file)


def show_img_boundary(img, boundary):
    """Show image and instance boundaires.

    Args:
        img (ndarray): The input image.
        boundary (list[float or int]): The input boundary.
    """
    assert isinstance(img, np.ndarray)
    assert utils.is_type_list(boundary, (int, float))

    cv2.polylines(
        img, [np.array(boundary).astype(np.int32).reshape(-1, 1, 2)],
        True,
        color=(0, 255, 0),
        thickness=1)
    plt.imshow(img)
    plt.show()


def show_pred_gt(preds,
                 gts,
                 show=False,
                 win_name='',
                 wait_time=0,
                 out_file=None):
    """Show detection and ground truth for one image.

    Args:
        preds (list[list[float]]): The detection boundary list.
        gts (list[list[float]]): The ground truth boundary list.
        show (bool): Whether to show the image.
        win_name (str): The window name.
        wait_time (int): The value of waitKey param.
        out_file (str): The filename of the output.
    """
    assert utils.is_2dlist(preds)
    assert utils.is_2dlist(gts)
    assert isinstance(show, bool)
    assert isinstance(win_name, str)
    assert isinstance(wait_time, int)
    assert utils.is_none_or_type(out_file, str)

    p_xy = [p for boundary in preds for p in boundary]
    gt_xy = [g for gt in gts for g in gt]

    max_xy = np.max(np.array(p_xy + gt_xy).reshape(-1, 2), axis=0)

    width = int(max_xy[0]) + 100
    height = int(max_xy[1]) + 100

    img = np.ones((height, width, 3), np.int8) * 255
    pred_color = mmcv.color_val('red')
    gt_color = mmcv.color_val('blue')
    thickness = 1

    for boundary in preds:
        cv2.polylines(
            img, [np.array(boundary).astype(np.int32).reshape(-1, 1, 2)],
            True,
            color=pred_color,
            thickness=thickness)
    for gt in gts:
        cv2.polylines(
            img, [np.array(gt).astype(np.int32).reshape(-1, 1, 2)],
            True,
            color=gt_color,
            thickness=thickness)
    if show:
        mmcv.imshow(img, win_name, wait_time)
    if out_file is not None:
        mmcv.imwrite(img, out_file)

    return img


def imshow_pred_boundary(img,
                         boundaries_with_scores,
                         labels,
                         score_thr=0,
                         boundary_color='blue',
                         text_color='blue',
                         thickness=1,
                         font_scale=0.5,
                         show=True,
                         win_name='',
                         wait_time=0,
                         out_file=None,
                         show_score=False):
    """Draw boundaries and class labels (with scores) on an image.

    Args:
        img (str or ndarray): The image to be displayed.
        boundaries_with_scores (list[list[float]]): Boundaries with scores.
        labels (list[int]): Labels of boundaries.
        score_thr (float): Minimum score of boundaries to be shown.
        boundary_color (str or tuple or :obj:`Color`): Color of boundaries.
        text_color (str or tuple or :obj:`Color`): Color of texts.
        thickness (int): Thickness of lines.
        font_scale (float): Font scales of texts.
        show (bool): Whether to show the image.
        win_name (str): The window name.
        wait_time (int): Value of waitKey param.
        out_file (str or None): The filename of the output.
        show_score (bool): Whether to show text instance score.
    """
    assert isinstance(img, (str, np.ndarray))
    assert utils.is_2dlist(boundaries_with_scores)
    assert utils.is_type_list(labels, int)
    assert utils.equal_len(boundaries_with_scores, labels)
    if len(boundaries_with_scores) == 0:
        warnings.warn('0 text found in ' + out_file)
        return None

    utils.valid_boundary(boundaries_with_scores[0])
    img = mmcv.imread(img)

    scores = np.array([b[-1] for b in boundaries_with_scores])
    inds = scores > score_thr
    boundaries = [boundaries_with_scores[i][:-1] for i in np.where(inds)[0]]
    scores = [scores[i] for i in np.where(inds)[0]]
    labels = [labels[i] for i in np.where(inds)[0]]

    boundary_color = mmcv.color_val(boundary_color)
    text_color = mmcv.color_val(text_color)
    font_scale = 0.5

    for boundary, score in zip(boundaries, scores):
        boundary_int = np.array(boundary).astype(np.int32)

        cv2.polylines(
            img, [boundary_int.reshape(-1, 1, 2)],
            True,
            color=boundary_color,
            thickness=thickness)

        if show_score:
            label_text = f'{score:.02f}'
            cv2.putText(img, label_text,
                        (boundary_int[0], boundary_int[1] - 2),
                        cv2.FONT_HERSHEY_COMPLEX, font_scale, text_color)
    if show:
        mmcv.imshow(img, win_name, wait_time)
    if out_file is not None:
        mmcv.imwrite(img, out_file)

    return img


def imshow_text_char_boundary(img,
                              text_quads,
                              boundaries,
                              char_quads,
                              chars,
                              show=False,
                              thickness=1,
                              font_scale=0.5,
                              win_name='',
                              wait_time=-1,
                              out_file=None):
    """Draw text boxes and char boxes on img.

    Args:
        img (str or ndarray): The img to be displayed.
        text_quads (list[list[int|float]]): The text boxes.
        boundaries (list[list[int|float]]): The boundary list.
        char_quads (list[list[list[int|float]]]): A 2d list of char boxes.
            char_quads[i] is for the ith text, and char_quads[i][j] is the jth
            char of the ith text.
        chars (list[list[char]]). The string for each text box.
        thickness (int): Thickness of lines.
        font_scale (float): Font scales of texts.
        show (bool): Whether to show the image.
        win_name (str): The window name.
        wait_time (int): Value of waitKey param.
        out_file (str or None): The filename of the output.
    """
    assert isinstance(img, (np.ndarray, str))
    assert utils.is_2dlist(text_quads)
    assert utils.is_2dlist(boundaries)
    assert utils.is_3dlist(char_quads)
    assert utils.is_2dlist(chars)
    assert utils.equal_len(text_quads, char_quads, boundaries)

    img = mmcv.imread(img)
    char_color = [mmcv.color_val('blue'), mmcv.color_val('green')]
    text_color = mmcv.color_val('red')
    text_inx = 0
    for text_box, boundary, char_box, txt in zip(text_quads, boundaries,
                                                 char_quads, chars):
        text_box = np.array(text_box)
        boundary = np.array(boundary)

        text_box = text_box.reshape(-1, 2).astype(np.int32)
        cv2.polylines(
            img, [text_box.reshape(-1, 1, 2)],
            True,
            color=text_color,
            thickness=thickness)
        if boundary.shape[0] > 0:
            cv2.polylines(
                img, [boundary.reshape(-1, 1, 2)],
                True,
                color=text_color,
                thickness=thickness)

        for b in char_box:
            b = np.array(b)
            c = char_color[text_inx % 2]
            b = b.astype(np.int32)
            cv2.polylines(
                img, [b.reshape(-1, 1, 2)], True, color=c, thickness=thickness)

        label_text = ''.join(txt)
        cv2.putText(img, label_text, (text_box[0, 0], text_box[0, 1] - 2),
                    cv2.FONT_HERSHEY_COMPLEX, font_scale, text_color)
        text_inx = text_inx + 1

    if show:
        mmcv.imshow(img, win_name, wait_time)
    if out_file is not None:
        mmcv.imwrite(img, out_file)

    return img


def tile_image(images):
    """Combined multiple images to one vertically.

    Args:
        images (list[np.ndarray]): Images to be combined.
    """
    assert isinstance(images, list)
    assert len(images) > 0

    for i, _ in enumerate(images):
        if len(images[i].shape) == 2:
            images[i] = cv2.cvtColor(images[i], cv2.COLOR_GRAY2BGR)

    widths = [img.shape[1] for img in images]
    heights = [img.shape[0] for img in images]
    h, w = sum(heights), max(widths)
    vis_img = np.zeros((h, w, 3), dtype=np.uint8)

    offset_y = 0
    for image in images:
        img_h, img_w = image.shape[:2]
        vis_img[offset_y:(offset_y + img_h), 0:img_w, :] = image
        offset_y += img_h

    return vis_img


def imshow_text_label(img,
                      pred_label,
                      gt_label,
                      show=False,
                      win_name='',
                      wait_time=-1,
                      out_file=None):
    """Draw predicted texts and ground truth texts on images.

    Args:
        img (str or np.ndarray): Image filename or loaded image.
        pred_label (str): Predicted texts.
        gt_label (str): Ground truth texts.
        show (bool): Whether to show the image.
        win_name (str): The window name.
        wait_time (int): Value of waitKey param.
        out_file (str): The filename of the output.
    """
    assert isinstance(img, (np.ndarray, str))
    assert isinstance(pred_label, str)
    assert isinstance(gt_label, str)
    assert isinstance(show, bool)
    assert isinstance(win_name, str)
    assert isinstance(wait_time, int)

    img = mmcv.imread(img)

    src_h, src_w = img.shape[:2]
    resize_height = 64
    resize_width = int(1.0 * src_w / src_h * resize_height)
    img = cv2.resize(img, (resize_width, resize_height))
    h, w = img.shape[:2]

    if is_contain_chinese(pred_label):
        pred_img = draw_texts_by_pil(img, [pred_label], None)
    else:
        pred_img = np.ones((h, w, 3), dtype=np.uint8) * 255
        cv2.putText(pred_img, pred_label, (5, 40), cv2.FONT_HERSHEY_SIMPLEX,
                    0.9, (0, 0, 255), 2)
    images = [pred_img, img]

    if gt_label != '':
        if is_contain_chinese(gt_label):
            gt_img = draw_texts_by_pil(img, [gt_label], None)
        else:
            gt_img = np.ones((h, w, 3), dtype=np.uint8) * 255
            cv2.putText(gt_img, gt_label, (5, 40), cv2.FONT_HERSHEY_SIMPLEX,
                        0.9, (255, 0, 0), 2)
        images.append(gt_img)

    img = tile_image(images)

    if show:
        mmcv.imshow(img, win_name, wait_time)
    if out_file is not None:
        mmcv.imwrite(img, out_file)

    return img


def imshow_node(img,
                result,
                boxes,
                idx_to_cls={},
                show=False,
                win_name='',
                wait_time=-1,
                out_file=None):

    img = mmcv.imread(img)
    h, w = img.shape[:2]

    max_value, max_idx = torch.max(result['nodes'].detach().cpu(), -1)
    node_pred_label = max_idx.numpy().tolist()
    node_pred_score = max_value.numpy().tolist()

    texts, text_boxes = [], []
    for i, box in enumerate(boxes):
        new_box = [[box[0], box[1]], [box[2], box[1]], [box[2], box[3]],
                   [box[0], box[3]]]
        Pts = np.array([new_box], np.int32)
        cv2.polylines(
            img, [Pts.reshape((-1, 1, 2))],
            True,
            color=(255, 255, 0),
            thickness=1)
        x_min = int(min([point[0] for point in new_box]))
        y_min = int(min([point[1] for point in new_box]))

        # text
        pred_label = str(node_pred_label[i])
        if pred_label in idx_to_cls:
            pred_label = idx_to_cls[pred_label]
        pred_score = '{:.2f}'.format(node_pred_score[i])
        text = pred_label + '(' + pred_score + ')'
        texts.append(text)

        # text box
        font_size = int(
            min(
                abs(new_box[3][1] - new_box[0][1]),
                abs(new_box[1][0] - new_box[0][0])))
        char_num = len(text)
        text_box = [
            x_min * 2, y_min, x_min * 2 + font_size * char_num, y_min,
            x_min * 2 + font_size * char_num, y_min + font_size, x_min * 2,
            y_min + font_size
        ]
        text_boxes.append(text_box)

    pred_img = np.ones((h, w * 2, 3), dtype=np.uint8) * 255
    pred_img = draw_texts_by_pil(
        pred_img, texts, text_boxes, draw_box=False, on_ori_img=True)

    vis_img = np.ones((h, w * 3, 3), dtype=np.uint8) * 255
    vis_img[:, :w] = img
    vis_img[:, w:] = pred_img

    if show:
        mmcv.imshow(vis_img, win_name, wait_time)
    if out_file is not None:
        mmcv.imwrite(vis_img, out_file)

    return vis_img


def gen_color():
    """Generate BGR color schemes."""
    color_list = [(101, 67, 254), (154, 157, 252), (173, 205, 249),
                  (123, 151, 138), (187, 200, 178), (148, 137, 69),
                  (169, 200, 200), (155, 175, 131), (154, 194, 182),
                  (178, 190, 137), (140, 211, 222), (83, 156, 222)]
    return color_list


def draw_polygons(img, polys):
    """Draw polygons on image.

    Args:
        img (np.ndarray): The original image.
        polys (list[list[float]]): Detected polygons.
    Return:
        out_img (np.ndarray): Visualized image.
    """
    dst_img = img.copy()
    color_list = gen_color()
    out_img = dst_img
    for idx, poly in enumerate(polys):
        poly = np.array(poly).reshape((-1, 1, 2)).astype(np.int32)
        cv2.drawContours(
            img,
            np.array([poly]),
            -1,
            color_list[idx % len(color_list)],
            thickness=cv2.FILLED)
        out_img = cv2.addWeighted(dst_img, 0.5, img, 0.5, 0)
    return out_img


def get_optimal_font_scale(text, width):
    """Get optimal font scale for cv2.putText.

    Args:
        text (str): Text in one box.
        width (int): The box width.
    """
    for scale in reversed(range(0, 60, 1)):
        textSize = cv2.getTextSize(
            text,
            fontFace=cv2.FONT_HERSHEY_SIMPLEX,
            fontScale=scale / 10,
            thickness=1)
        new_width = textSize[0][0]
        if new_width <= width:
            return scale / 10
    return 1


def draw_texts(img, texts, boxes=None, draw_box=True, on_ori_img=False):
    """Draw boxes and texts on empty img.

    Args:
        img (np.ndarray): The original image.
        texts (list[str]): Recognized texts.
        boxes (list[list[float]]): Detected bounding boxes.
        draw_box (bool): Whether draw box or not. If False, draw text only.
        on_ori_img (bool): If True, draw box and text on input image,
            else, on a new empty image.
    Return:
        out_img (np.ndarray): Visualized image.
    """
    color_list = gen_color()
    h, w = img.shape[:2]
    if boxes is None:
        boxes = [[0, 0, w, 0, w, h, 0, h]]
    assert len(texts) == len(boxes)

    if on_ori_img:
        out_img = img
    else:
        out_img = np.ones((h, w, 3), dtype=np.uint8) * 255
    for idx, (box, text) in enumerate(zip(boxes, texts)):
        if draw_box:
            new_box = [[x, y] for x, y in zip(box[0::2], box[1::2])]
            Pts = np.array([new_box], np.int32)
            cv2.polylines(
                out_img, [Pts.reshape((-1, 1, 2))],
                True,
                color=color_list[idx % len(color_list)],
                thickness=1)
        min_x = int(min(box[0::2]))
        max_y = int(
            np.mean(np.array(box[1::2])) + 0.2 *
            (max(box[1::2]) - min(box[1::2])))
        font_scale = get_optimal_font_scale(
            text, int(max(box[0::2]) - min(box[0::2])))
        cv2.putText(out_img, text, (min_x, max_y), cv2.FONT_HERSHEY_SIMPLEX,
                    font_scale, (0, 0, 0), 1)

    return out_img


def draw_texts_by_pil(img,
                      texts,
                      boxes=None,
                      draw_box=True,
                      on_ori_img=False,
                      font_size=None,
                      fill_color=None,
                      draw_pos=None,
                      return_text_size=False):
    """Draw boxes and texts on empty image, especially for Chinese.

    Args:
        img (np.ndarray): The original image.
        texts (list[str]): Recognized texts.
        boxes (list[list[float]]): Detected bounding boxes.
        draw_box (bool): Whether draw box or not. If False, draw text only.
        on_ori_img (bool): If True, draw box and text on input image,
            else on a new empty image.
        font_size (int, optional): Size to create a font object for a font.
        fill_color (tuple(int), optional): Fill color for text.
        draw_pos (list[tuple(int)], optional): Start point to draw each text.
        return_text_size (bool): If True, return the list of text size.

    Returns:
        (np.ndarray, list[tuple]) or np.ndarray: Return a tuple
        ``(out_img, text_sizes)``, where ``out_img`` is the output image
        with texts drawn on it and ``text_sizes`` are the size of drawing
        texts. If ``return_text_size`` is False, only the output image will be
        returned.
    """

    color_list = gen_color()
    h, w = img.shape[:2]
    if boxes is None:
        boxes = [[0, 0, w, 0, w, h, 0, h]]
    if draw_pos is None:
        draw_pos = [None for _ in texts]
    assert len(boxes) == len(texts) == len(draw_pos)

    if fill_color is None:
        fill_color = (0, 0, 0)

    if on_ori_img:
        out_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    else:
        out_img = Image.new('RGB', (w, h), color=(255, 255, 255))
    out_draw = ImageDraw.Draw(out_img)

    text_sizes = []
    for idx, (box, text, ori_point) in enumerate(zip(boxes, texts, draw_pos)):
        if len(text) == 0:
            continue
        min_x, max_x = min(box[0::2]), max(box[0::2])
        min_y, max_y = min(box[1::2]), max(box[1::2])
        color = tuple(list(color_list[idx % len(color_list)])[::-1])
        if draw_box:
            out_draw.line(box, fill=color, width=1)
        dirname, _ = os.path.split(os.path.abspath(__file__))
        font_path = os.path.join(dirname, 'font.TTF')
        if not os.path.exists(font_path):
            url = ('https://download.openmmlab.com/mmocr/data/font.TTF')
            print(f'Downloading {url} ...')
            local_filename, _ = urllib.request.urlretrieve(url)
            shutil.move(local_filename, font_path)
        tmp_font_size = font_size
        if tmp_font_size is None:
            box_width = max(max_x - min_x, max_y - min_y)
            tmp_font_size = int(0.9 * box_width / len(text))
        fnt = ImageFont.truetype(font_path, tmp_font_size)
        if ori_point is None:
            ori_point = (min_x + 1, min_y + 1)
        out_draw.text(ori_point, text, font=fnt, fill=fill_color)
        text_sizes.append(fnt.getsize(text))

    del out_draw

    out_img = cv2.cvtColor(np.asarray(out_img), cv2.COLOR_RGB2BGR)

    if return_text_size:
        return out_img, text_sizes

    return out_img


def is_contain_chinese(check_str):
    """Check whether string contains Chinese or not.

    Args:
        check_str (str): String to be checked.

    Return True if contains Chinese, else False.
    """
    for ch in check_str:
        if u'\u4e00' <= ch <= u'\u9fff':
            return True
    return False


def det_recog_show_result(img, end2end_res, out_file=None):
    """Draw `result`(boxes and texts) on `img`.

    Args:
        img (str or np.ndarray): The image to be displayed.
        end2end_res (dict): Text detect and recognize results.
        out_file (str): Image path where the visualized image should be saved.
    Return:
        out_img (np.ndarray): Visualized image.
    """
    img = mmcv.imread(img)
    boxes, texts = [], []
    for res in end2end_res['result']:
        boxes.append(res['box'])
        texts.append(res['text'])
    box_vis_img = draw_polygons(img, boxes)

    if is_contain_chinese(''.join(texts)):
        text_vis_img = draw_texts_by_pil(img, texts, boxes)
    else:
        text_vis_img = draw_texts(img, texts, boxes)

    h, w = img.shape[:2]
    out_img = np.ones((h, w * 2, 3), dtype=np.uint8)
    out_img[:, :w, :] = box_vis_img
    out_img[:, w:, :] = text_vis_img

    if out_file:
        mmcv.imwrite(out_img, out_file)

    return out_img


def draw_edge_result(img, result, edge_thresh=0.5, keynode_thresh=0.5):
    """Draw text and their relationship on empty images.

    Args:
        img (np.ndarray): The original image.
        result (dict): The result of model forward_test, including:
            - img_metas (list[dict]): List of meta information dictionary.
            - nodes (Tensor): Node prediction with size:
                number_node * node_classes.
            - edges (Tensor): Edge prediction with size: number_edge * 2.
        edge_thresh (float): Score threshold for edge classification.
        keynode_thresh (float): Score threshold for node
            (``key``) classification.

    Returns:
        np.ndarray: The image with key, value and relation drawn on it.
    """

    h, w = img.shape[:2]

    vis_area_width = w // 3 * 2
    vis_area_height = h
    dist_key_to_value = vis_area_width // 2
    dist_pair_to_pair = 30

    bbox_x1 = dist_pair_to_pair
    bbox_y1 = 0

    new_w = vis_area_width
    new_h = vis_area_height
    pred_edge_img = np.ones((new_h, new_w, 3), dtype=np.uint8) * 255

    nodes = result['nodes'].detach().cpu()
    texts = result['img_metas'][0]['ori_texts']
    num_nodes = result['nodes'].size(0)
    edges = result['edges'].detach().cpu()[:, -1].view(num_nodes, num_nodes)

    # (i, j) will be a valid pair
    # either edge_score(node_i->node_j) > edge_thresh
    # or edge_score(node_j->node_i) > edge_thresh
    pairs = (torch.max(edges, edges.T) > edge_thresh).nonzero(as_tuple=True)
    pairs = (pairs[0].numpy().tolist(), pairs[1].numpy().tolist())

    # 1. "for n1, n2 in zip(*pairs) if n1 < n2":
    #     Only (n1, n2) will be included if n1 < n2 but not (n2, n1), to
    #     avoid duplication.
    # 2. "(n1, n2) if nodes[n1, 1] > nodes[n1, 2]":
    #     nodes[n1, 1] is the score that this node is predicted as key,
    #     nodes[n1, 2] is the score that this node is predicted as value.
    #     If nodes[n1, 1] > nodes[n1, 2], n1 will be the index of key,
    #     so that n2 will be the index of value.
    result_pairs = [(n1, n2) if nodes[n1, 1] > nodes[n1, 2] else (n2, n1)
                    for n1, n2 in zip(*pairs) if n1 < n2]

    result_pairs.sort()
    result_pairs_score = [
        torch.max(edges[n1, n2], edges[n2, n1]) for n1, n2 in result_pairs
    ]

    key_current_idx = -1
    pos_current = (-1, -1)
    newline_flag = False

    key_font_size = 15
    value_font_size = 15
    key_font_color = (0, 0, 0)
    value_font_color = (0, 0, 255)
    arrow_color = (0, 0, 255)
    score_color = (0, 255, 0)
    for pair, pair_score in zip(result_pairs, result_pairs_score):
        key_idx = pair[0]
        if nodes[key_idx, 1] < keynode_thresh:
            continue
        if key_idx != key_current_idx:
            # move y-coords down for a new key
            bbox_y1 += 10
            # enlarge blank area to show key-value info
            if newline_flag:
                bbox_x1 += vis_area_width
                tmp_img = np.ones(
                    (new_h, new_w + vis_area_width, 3), dtype=np.uint8) * 255
                tmp_img[:new_h, :new_w] = pred_edge_img
                pred_edge_img = tmp_img
                new_w += vis_area_width
                newline_flag = False
                bbox_y1 = 10
        key_text = texts[key_idx]
        key_pos = (bbox_x1, bbox_y1)
        value_idx = pair[1]
        value_text = texts[value_idx]
        value_pos = (bbox_x1 + dist_key_to_value, bbox_y1)
        if key_idx != key_current_idx:
            # draw text for a new key
            key_current_idx = key_idx
            pred_edge_img, text_sizes = draw_texts_by_pil(
                pred_edge_img, [key_text],
                draw_box=False,
                on_ori_img=True,
                font_size=key_font_size,
                fill_color=key_font_color,
                draw_pos=[key_pos],
                return_text_size=True)
            pos_right_bottom = (key_pos[0] + text_sizes[0][0],
                                key_pos[1] + text_sizes[0][1])
            pos_current = (pos_right_bottom[0] + 5, bbox_y1 + 10)
            pred_edge_img = cv2.arrowedLine(
                pred_edge_img, (pos_right_bottom[0] + 5, bbox_y1 + 10),
                (bbox_x1 + dist_key_to_value - 5, bbox_y1 + 10), arrow_color,
                1)
            score_pos_x = int(
                (pos_right_bottom[0] + bbox_x1 + dist_key_to_value) / 2.)
            score_pos_y = bbox_y1 + 10 - int(key_font_size * 0.3)
        else:
            # draw arrow from key to value
            if newline_flag:
                tmp_img = np.ones((new_h + dist_pair_to_pair, new_w, 3),
                                  dtype=np.uint8) * 255
                tmp_img[:new_h, :new_w] = pred_edge_img
                pred_edge_img = tmp_img
                new_h += dist_pair_to_pair
            pred_edge_img = cv2.arrowedLine(pred_edge_img, pos_current,
                                            (bbox_x1 + dist_key_to_value - 5,
                                             bbox_y1 + 10), arrow_color, 1)
            score_pos_x = int(
                (pos_current[0] + bbox_x1 + dist_key_to_value - 5) / 2.)
            score_pos_y = int((pos_current[1] + bbox_y1 + 10) / 2.)
        # draw edge score
        cv2.putText(pred_edge_img, '{:.2f}'.format(pair_score),
                    (score_pos_x, score_pos_y), cv2.FONT_HERSHEY_COMPLEX, 0.4,
                    score_color)
        # draw text for value
        pred_edge_img = draw_texts_by_pil(
            pred_edge_img, [value_text],
            draw_box=False,
            on_ori_img=True,
            font_size=value_font_size,
            fill_color=value_font_color,
            draw_pos=[value_pos],
            return_text_size=False)
        bbox_y1 += dist_pair_to_pair
        if bbox_y1 + dist_pair_to_pair >= new_h:
            newline_flag = True

    return pred_edge_img


def imshow_edge(img,
                result,
                boxes,
                show=False,
                win_name='',
                wait_time=-1,
                out_file=None):
    """Display the prediction results of the nodes and edges of the KIE model.

    Args:
        img (np.ndarray): The original image.
        result (dict): The result of model forward_test, including:
            - img_metas (list[dict]): List of meta information dictionary.
            - nodes (Tensor): Node prediction with size: \
                number_node * node_classes.
            - edges (Tensor): Edge prediction with size: number_edge * 2.
        boxes (list): The text boxes corresponding to the nodes.
        show (bool): Whether to show the image. Default: False.
        win_name (str): The window name. Default: ''
        wait_time (float): Value of waitKey param. Default: 0.
        out_file (str or None): The filename to write the image.
            Default: None.

    Returns:
        np.ndarray: The image with key, value and relation drawn on it.
    """
    img = mmcv.imread(img)
    h, w = img.shape[:2]
    color_list = gen_color()

    for i, box in enumerate(boxes):
        new_box = [[box[0], box[1]], [box[2], box[1]], [box[2], box[3]],
                   [box[0], box[3]]]
        Pts = np.array([new_box], np.int32)
        cv2.polylines(
            img, [Pts.reshape((-1, 1, 2))],
            True,
            color=color_list[i % len(color_list)],
            thickness=1)

    pred_img_h = h
    pred_img_w = w

    pred_edge_img = draw_edge_result(img, result)
    pred_img_h = max(pred_img_h, pred_edge_img.shape[0])
    pred_img_w += pred_edge_img.shape[1]

    vis_img = np.zeros((pred_img_h, pred_img_w, 3), dtype=np.uint8)
    vis_img[:h, :w] = img
    vis_img[:, w:] = 255

    height_t, width_t = pred_edge_img.shape[:2]
    vis_img[:height_t, w:(w + width_t)] = pred_edge_img

    if show:
        mmcv.imshow(vis_img, win_name, wait_time)
    if out_file is not None:
        mmcv.imwrite(vis_img, out_file)
        res_dic = {
            'boxes': boxes,
            'nodes': result['nodes'].detach().cpu(),
            'edges': result['edges'].detach().cpu(),
            'metas': result['img_metas'][0]
        }
        mmcv.dump(res_dic, f'{out_file}_res.pkl')

    return vis_img