import os
import random

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont


def visualize_kie(
    img,
    boxes,
    pred_labels,
    skip_classes=["other"],
    image_name="out_kie.jpg",
    outdir="workdirs/visualize",
    colors=None,
    texts=None, 
):
    """Visualize kie output

    Args:
        image (PIL.Image): _description_
        boxes (list[xyxy]): list of xyxy boxes
        kie_preds (list[str]): _description_
        skip_classes (list[str], optional): list of skip labels. Defaults to ['other'].
        image_path (str, optional): _description_. Defaults to None.
        outdir (_type_, optional): _description_. Defaults to None.
    """
    if not os.path.exists(outdir):
        os.makedirs(outdir, exist_ok=True)

    image = img.copy()
    if isinstance(image, np.ndarray):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    set_labels = sorted(list(set(pred_labels)))
    if colors is None:
        colors = [
            (
                random.randint(0, 255),
                random.randint(0, 255),
                random.randint(0, 255),
            )
            for _ in range(len(set_labels))
        ]
    text_color = (255,0,0)
    text_checkin = [0] * len(set_labels)





    draw = ImageDraw.Draw(image)
    # font = ImageFont.load_default()
    font = ImageFont.truetype("/usr/share/fonts/truetype/msttcorefonts/Arial.ttf", 16)
    for i, pred_label in enumerate(pred_labels):
        if pred_label in skip_classes:
            continue
        pred_label_idx = set_labels.index(pred_label)
        box = boxes[i]
        color = colors[pred_label_idx] if isinstance(colors, list) else colors[pred_label]
        draw.rectangle(box, fill=None, outline=color)

        # skip visualize text label for other boxes
        if texts is not None and text_checkin[set_labels.index(pred_label)] == 0:
            draw.text(
                (int(box[0]) + 10, int(box[1]) - 10),
                fill=text_color if isinstance(colors, list) else colors[pred_label],
                # text=str(pred_label) + " - " + str(texts[i]),
                text=str(texts[i]),
                font=font,
            )
            text_checkin[pred_label_idx] = 1
    image.save(os.path.join(outdir, image_name))


def visualize_ocr(
    img, boxes, texts, image_name="out_ocr.jpg", outdir="workdirs/visualize"
):
    """Visualize kie output

    Args:
        image (PIL.Image): _description_
        boxes (list[xyxy]): list of xyxy boxes
        kie_preds (list[str]): _description_
        skip_classes (list[str], optional): list of skip labels. Defaults to ['other'].
        image_path (str, optional): _description_. Defaults to None.
        outdir (_type_, optional): _description_. Defaults to None.
    """
    if not os.path.exists(outdir):
        os.makedirs(outdir, exist_ok=True)

    image = img.copy()
    if isinstance(image, np.ndarray):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    color = (
        random.randint(0, 255),
        random.randint(0, 255),
        random.randint(0, 255),
    )

    draw = ImageDraw.Draw(image)
    font = ImageFont.load_default()
    for i, text in enumerate(texts):
        box = boxes[i]
        draw.rectangle(box, fill=None, outline=color)
        # draw.text((int(box[0])+10, int(box[1])-10),fill=color, text=text, font = font)
    image.save(os.path.join(outdir, image_name))
