import os
import glob
import time
import json
import random

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


def get_template(path):
    # RGB to HSV
    img = Image.open(path).convert('HSV')
    copy = img.convert('L').copy()
    # where is white
    copy = np.where(np.array(copy) == 255, 255, 0).astype(np.uint8)
    # convert numpy array to image
    copy = Image.fromarray(copy, "L")
    # pick the H channel
    img = img.split()[0]
    # divide by 16
    img = np.array(img) + 8
    img = Image.fromarray(img // 16)
    return img, copy


# 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 json file
def load_config(path):
    with open(path, "r") as f:
        config = json.load(f)
    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, dot = 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
    plateText = ""

    for idx, plate in enumerate(plates):
        for char in plate:
            count += 1
            plateText += char
            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)

    # add the dot img on the mask
    mask = Image.composite(dot, mask, dot)
    mask, _ = resize_height(mask, 512)
    mask_wo_border = mask.copy()
    mask_wo_border = mask_wo_border.filter(ImageFilter.MaxFilter(3))
    mask = mask.filter(ImageFilter.MaxFilter(5)).filter(ImageFilter.GaussianBlur(1.5))

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

    plate_bg = get_background(name)
    mask_wo_border = mask_wo_border.resize(plate_bg.size)

    # make the mask as alpha channel
    bg = Image.new("RGBA", plate_bg.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 plateText, bg, mask, ord_


# 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):
    cfg = load_config(os.path.join(config_path, name, "config.json"))
    return get_mask(cfg) # text, bg, mask, ord_


demo = gr.Interface(
    generate,
    [
        gr.Dropdown(plates, value="cn_truck", label="Plate Style"),
    ],
    [
        gr.Text(label="Plate Text"),
        gr.Image(label="Background", show_label=False, format="png"),
        gr.Image(label="Mask", show_label=False, format="png"),
        gr.Image(label="Ord", show_label=False, format="png"),
    ],
    allow_flagging=False,
)


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