import os
import sys
import time
import glob
import yaml
import random

import numpy as np
import gradio as gr
from PIL import Image, ImageDraw, ImageFilter


def get_template(path):
    # RGB to HSV
    img = Image.open(path).convert('HSV')
    # pick the H channel
    img = img.split()[0]
    # divide by 16
    img = np.array(img) + 8
    img = Image.fromarray(img // 16)
    return img


# resize via height, keep the aspect ratio
def resize_height(img, height, nearest=False):
    width, old_height = img.size
    new_width = int(width * height / old_height)
    return img.resize((new_width, height), Image.NEAREST if nearest else Image.BICUBIC), new_width


# load config yaml file
def load_config(path):
    with open(path, "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    return config


def parse_config(config):
    def random_choice(char):
        if char == "#":
            return random.choice(config["plateNums"] + config["plateChars"])
        elif char == "A":
            return random.choice(config["plateChars"])
        elif char == "0":
            return random.choice(config["plateNums"])
        elif char == "@":
            return random.choice(config["plateCities"])
        return char

    plate_patch = [(random_choice(char) for char in plate) for plate in config["plateFormat"]]

    return config["name"], plate_patch


def get_background(name):
    bgs = os.path.join("config", name, "background")
    bg_path = random.choice(glob.glob(f"{bgs}/*.*"))
    bg = Image.open(bg_path).convert("RGB")
    return bg


# get mask and paste the file from folder to generate the mask
def get_mask(cfg):
    name, plates = parse_config(cfg)
    root_dir = os.path.join("config", name)
    # if none key, return default value
    plateTextColor = cfg.get("plateTextColor", "#000000").lstrip("#")

    img = get_template(os.path.join(root_dir, "template.png"))
    mask = Image.new("L", img.size, 0)
    ord_ = Image.new("L", img.size, 0)
    count = 0

    for idx, plate in enumerate(plates):
        for char in plate:
            count += 1
            m_char_path = random.choice(glob.glob(f"{root_dir}/{idx}/{char}_*.png"))
            mask_rect = np.where(np.array(img) == count)
            h_min, w_min = np.min(mask_rect, axis=1)
            h_max, w_max = np.max(mask_rect, axis=1)
            m_char = Image.open(m_char_path).convert("RGBA").split()[-1]
            m_char, new_width = resize_height(m_char, h_max - h_min)
            x = (w_min + w_max - new_width) // 2

            mask.paste(m_char, (x, h_min), m_char)

            m_char = np.array(np.array(m_char) > 128, dtype=np.uint8) * 255
            m_char = Image.fromarray(m_char)
            ord_.paste(count*16, (x, h_min), m_char)

    mask, _ = resize_height(mask, 512)
    mask_wo_border = mask.copy()
    mask = mask.filter(ImageFilter.MaxFilter(5)).filter(ImageFilter.GaussianBlur(1.5))

    ord_, _ = resize_height(ord_, 512, nearest=True)

    plate_bg = get_background(name).resize(mask.size)
    # make the mask as alpha channel
    bg = Image.new("RGBA", mask.size, (0, 0, 0, 0))
    bg.paste(plate_bg, (0, 0))
    # hex to rgb
    plateTextColor = tuple(int(plateTextColor[i:i + 2], 16) for i in (0, 2, 4))
    bg.paste(plateTextColor, (0, 0), mask_wo_border)

    return mask, ord_, bg


# get all folders in the config folder
config_path = os.path.join(os.path.dirname(__file__), "config")
plates = [f for f in os.listdir(config_path) if os.path.isdir(os.path.join(config_path, f))]


def generate(name):
    print(name)
    cfg = load_config(os.path.join(config_path, name, "config.yaml"))
    mask, ord_, bg = get_mask(cfg)
    return bg, mask, ord_


demo = gr.Interface(
    generate,
    [
        gr.Dropdown(plates, value="cn_truck", label="Plate Style"),
    ],
    [
        gr.Image(label="Background"),
        gr.Image(label="Mask"),
        gr.Image(label="Ord"),
    ]
)


if __name__ == "__main__":
    demo.launch()