import onnx
import onnxruntime as ort
import numpy as np
import cv2
import os
import csv

VIT_MODEL_DSV3_REPO = "SmilingWolf/wd-vit-tagger-v3"

MODEL_DIR='sd/prompt_helper/model'
ORT_SESSION=ort.InferenceSession(f'{MODEL_DIR}/model.onnx', providers=['CPUExecutionProvider'])
#ORT_INPUT_NAME=ort_session.get_inputs()[0].name

IMAGE_SIZE = 448

def img_preporation(img):
    bgr_img = np.array(img)[:, :, ::-1].copy()
    size = max(bgr_img.shape[0:2])
    pad_x = size - bgr_img.shape[1]
    pad_y = size - bgr_img.shape[0]
    pad_l = pad_x // 2
    pad_t = pad_y // 2

    #add paddings to squaring image
    np.pad(bgr_img, ((pad_t, pad_y - pad_t), (pad_l, pad_x - pad_l), (0, 0)), mode="constant", constant_values=255)

    #adaptive resize
    interp = cv2.INTER_AREA if size > IMAGE_SIZE else cv2.INTER_LANCZOS4
    bgr_img = cv2.resize(bgr_img, (IMAGE_SIZE, IMAGE_SIZE), interpolation=interp)

    bgr_img = bgr_img.astype(np.float32)


def get_help(img):
    # = model[0]
    ort_input_name = ORT_SESSION.get_inputs()[0].name

    with open(os.path.join(MODEL_DIR, "selected_tags.csv"), "r", encoding="utf-8") as f:
        reader = csv.reader(f)
        l = [row for row in reader]
        header = l[0]  # tag_id,name,category,count
        rows = l[1:]
    assert header[0] == "tag_id" and header[1] == "name" and header[2] == "category", f"unexpected csv format: {header}"

    general_tags = [row[1] for row in rows[1:] if row[2] == "0"]
    character_tags = [row[1] for row in rows[1:] if row[2] == "4"]

    tag_freq = {}
    undesired_tags = ["transparent background"]

    #img = Image.open(image_path)
    preped_img = img_preporation(img)
    preped_img = np.expand_dims(preped_img, axis=0)

    # Run inference
    prob = ORT_SESSION.run(None, {ort_input_name: preped_img})[0][0]
    # Generate Tags
    combined_tags = []
    general_tag_text = ""
    character_tag_text = ""
    remove_underscore = True
    caption_separator = ", "
    general_threshold = 0.35
    character_threshold = 0.35

    for i, p in enumerate(prob[4:]):
        if i < len(general_tags) and p >= general_threshold:
            tag_name = general_tags[i]
            if remove_underscore and len(tag_name) > 3:  # ignore emoji tags like >_< and ^_^
                tag_name = tag_name.replace("_", " ")

            if tag_name not in undesired_tags:
                tag_freq[tag_name] = tag_freq.get(tag_name, 0) + 1
                general_tag_text += caption_separator + tag_name
                combined_tags.append(tag_name)
        elif i >= len(general_tags) and p >= character_threshold:
            tag_name = character_tags[i - len(general_tags)]
            if remove_underscore and len(tag_name) > 3:
                tag_name = tag_name.replace("_", " ")

            if tag_name not in undesired_tags:
                tag_freq[tag_name] = tag_freq.get(tag_name, 0) + 1
                character_tag_text += caption_separator + tag_name
                combined_tags.append(tag_name)

    # Remove leading comma
    if len(general_tag_text) > 0:
        general_tag_text = general_tag_text[len(caption_separator) :]
    if len(character_tag_text) > 0:
        character_tag_text = character_tag_text[len(caption_separator) :]
    tag_text = caption_separator.join(combined_tags)
    return tag_text