import numpy as np
import cv2, os, argparse
from glob import glob
from tqdm import tqdm

from plate_number import random_select, generate_plate_number_white, generate_plate_number_army, generate_plate_number_yellow_xue, generate_plate_number_yellow
from plate_number import generate_plate_number_black_gangao
from plate_number import generate_plate_number_black_shi_old, generate_plate_number_black_shi_new
from plate_number import generate_plate_number_black_ling_old, generate_plate_number_black_ling_new
from plate_number import generate_plate_number_blue
from plate_number import letters, digits


def get_location_data(length=7, split_id=1, height=140):
    """
    获取车牌号码在底牌中的位置
    length: 车牌字符数，7或者8，7为普通车牌、8为新能源车牌
    split_id: 分割空隙
    height: 车牌高度，对应单层和双层车牌
    """
    # 字符位置
    location_xy = np.zeros((length, 4), dtype=np.int32)

    # 单层车牌高度
    if height == 140:
        # 单层车牌，y轴坐标固定
        location_xy[:, 1] = 25
        location_xy[:, 3] = 115
        # 螺栓间隔
        step_split = 34 if length == 7 else 49
        # 字符间隔
        step_font = 12 if length == 7 else 9

        # 字符宽度
        width_font = 45
        for i in range(length):
            if i == 0:
                location_xy[i, 0] = 15
            elif i == split_id:
                location_xy[i, 0] = location_xy[i - 1, 2] + step_split
            else:
                location_xy[i, 0] = location_xy[i - 1, 2] + step_font
            # 新能源车牌
            if length == 8 and i > 0:
                width_font = 43
            location_xy[i, 2] = location_xy[i, 0] + width_font
    else:
        # 双层车牌第一层
        location_xy[0, :] = [110, 15, 190, 75]
        location_xy[1, :] = [250, 15, 330, 75]

        # 第二层
        width_font = 65
        step_font = 15
        for i in range(2, length):
            location_xy[i, 1] = 90
            location_xy[i, 3] = 200
            if i == 2:
                location_xy[i, 0] = 27
            else:
                location_xy[i, 0] = location_xy[i - 1, 2] + step_font
            location_xy[i, 2] = location_xy[i, 0] + width_font

    return location_xy


# 字符贴上底板
def copy_to_image_multi(img, font_img, bbox, bg_color, is_red):
    x1, y1, x2, y2 = bbox
    font_img = cv2.resize(font_img, (x2 - x1, y2 - y1))
    img_crop = img[y1: y2, x1: x2, :]

    if is_red:
        img_crop[font_img < 200, :] = [0, 0, 255]
    elif 'blue' in bg_color or 'black' in bg_color:
        img_crop[font_img < 200, :] = [255, 255, 255]
    else:
        img_crop[font_img < 200, :] = [0, 0, 0]
    return img


class MultiPlateGenerator:
    def __init__(self, adr_plate_model, adr_font):
        # 车牌底板路径
        self.adr_plate_model = adr_plate_model
        # 车牌字符路径
        self.adr_font = adr_font

        # 车牌字符图片，预存处理
        self.font_imgs = {}
        font_filenames = glob(os.path.join(adr_font, '*jpg'))
        for font_filename in font_filenames:
            font_img = cv2.imdecode(np.fromfile(font_filename, dtype=np.uint8), 0)

            if '140' in font_filename:
                font_img = cv2.resize(font_img, (45, 90))
            elif '220' in font_filename:
                font_img = cv2.resize(font_img, (65, 110))
            elif font_filename.split('_')[-1].split('.')[0] in letters + digits:
                font_img = cv2.resize(font_img, (43, 90))
            # 仅对"1"膨胀加粗（字为暗色，需反转→膨胀→反转回）
            key = os.path.basename(font_filename).split('.')[0]
            if key.endswith('_1') or key == '1':
                font_img = 255 - font_img
                font_img = cv2.dilate(font_img, np.ones((3, 3), np.uint8), iterations=2)
                font_img = 255 - font_img
            self.font_imgs[key] = font_img

        # 字符位置
        self.location_xys = {}
        for i in [7, 8]:
            for j in [1, 2, 3, 4]:
                for k in [140, 220]:
                    self.location_xys['{}_{}_{}'.format(i, j, k)] = \
                        get_location_data(length=i, split_id=j, height=k)

        # 底板背景缓存（避免每次 generate_plate_special 读盘）
        self._plate_bg_cache = {}

    # 获取字符位置
    def get_location_multi(self, plate_number, height=140):
        length = len(plate_number)
        if '警' in plate_number:
            split_id = 1
        elif '使' in plate_number:
            # 旧版：使在首位 → 间隙在第4位前；新版：使在末位 → 间隙在第3位前
            split_id = 4 if plate_number[0] == '使' else 3
        else:
            split_id = 2
        return self.location_xys['{}_{}_{}'.format(length, split_id, height)]

    # 随机生成车牌号码（黄/白/黑 + 大型新能源）
    def generate_plate_number(self):
        # 白牌警车/军车、黄牌教练车/大型车、港澳、使领馆新旧版、大型新能源
        generate_plate_number_funcs = [
            generate_plate_number_white,
            generate_plate_number_army,
            generate_plate_number_yellow_xue,
            generate_plate_number_yellow,
            generate_plate_number_black_gangao,
            generate_plate_number_black_shi_old,
            generate_plate_number_black_shi_new,
            generate_plate_number_black_ling_old,
            generate_plate_number_black_ling_new,
            lambda: generate_plate_number_blue(length=8),  # 大型新能源
        ]
        plate_number = random_select(generate_plate_number_funcs)()

        # 按GA36规范确定底板颜色
        if '港' in plate_number or '澳' in plate_number:
            bg_color = 'black'
        elif '领' in plate_number:
            bg_color = 'black'
        elif '使' in plate_number:
            bg_color = 'black_shi'
        elif '警' in plate_number:
            bg_color = 'white'
        elif plate_number[0] in letters:
            bg_color = 'white_army'
        elif '学' in plate_number:
            bg_color = 'yellow'
        else:
            bg_color = 'yellow'  # 大型车

        # 大型新能源（8位）
        if len(plate_number) == 8:
            bg_color = 'green_truck'

        return plate_number, bg_color, False

    # 随机生成车牌图片（全部单层，已移除双层挂车）
    def generate_plate(self, enhance=False):
        plate_number, bg_color, _ = self.generate_plate_number()
        height = 140

        # 获取底板图片
        number_xy = self.get_location_multi(plate_number, height)
        img_plate_model = cv2.imread(os.path.join(self.adr_plate_model, '{}_{}.PNG'.format(bg_color, height)))
        img_plate_model = cv2.resize(img_plate_model, (440 if len(plate_number) == 7 else 480, height))

        for i in range(len(plate_number)):
            if len(plate_number) == 8:
                # 新能源（green_car / green_truck）
                font_img = self.font_imgs['green_{}'.format(plate_number[i])]
            else:
                if '{}_{}'.format(height, plate_number[i]) in self.font_imgs:
                    font_img = self.font_imgs['{}_{}'.format(height, plate_number[i])]
                else:
                    # 兜底：220字体缩放到140
                    font_img = self.font_imgs.get('220_{}'.format(plate_number[i]),
                                self.font_imgs.get('220_up_{}'.format(plate_number[i]),
                                    self.font_imgs['220_down_{}'.format(plate_number[i])]))

            # 字符是否红色
            # 军车首字母、警、旧版使馆(使在首位)、旧版领馆(省A·XXXX领)为红色
            is_red = False
            if (i == 0 and plate_number[0] in letters) or plate_number[i] == '警':
                is_red = True
            elif plate_number[i] == '使' and i == 0:
                is_red = True  # 旧版使馆：使在首位 → 红字
            elif plate_number[i] == '领' and len(plate_number) >= 2 and plate_number[1] in letters:
                is_red = True  # 旧版领馆：省A·XXXX领，第二位是城市字母 → 红字
            elif i == 1 and plate_number[0] in letters and np.random.random(1) > 0.5:
                is_red = True  # 军车第二字符随机红色

            if enhance:
                k = np.random.randint(1, 6)
                kernel = np.ones((k, k), np.uint8)
                if np.random.random(1) > 0.5:
                    font_img = np.copy(cv2.erode(font_img, kernel, iterations=1))
                else:
                    font_img = np.copy(cv2.dilate(font_img, kernel, iterations=1))

            # 贴上底板
            img_plate_model = copy_to_image_multi(img_plate_model, font_img,
                                                  number_xy[i, :], bg_color, is_red)

        img_plate_model = cv2.blur(img_plate_model, (3, 3))

        return img_plate_model, number_xy, plate_number, bg_color, False

    def generate_plate_special(self, plate_number, bg_color, is_double=False, enhance=False):
        """
        生成特定号码、颜色车牌
        :param plate_number: 车牌号码
        :param bg_color: 背景颜色，blue/yellow/white/white_army/black/black_shi/green_car/green_truck
        :param is_double: 是否双层（默认False，已移除双层挂车生成）
        :param enhance: 图像增强
        :return: 车牌图
        """
        height = 220 if is_double else 140

        number_xy = self.get_location_multi(plate_number, height)

        # 底板背景（缓存读取，避免每张图读盘）
        cache_key = '{}_{}'.format(bg_color, height)
        if cache_key not in self._plate_bg_cache:
            bg_path = os.path.join(self.adr_plate_model, '{}.PNG'.format(cache_key))
            self._plate_bg_cache[cache_key] = cv2.imdecode(
                np.fromfile(bg_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        img_plate_model = self._plate_bg_cache[cache_key].copy()
        img_plate_model = cv2.resize(img_plate_model, (440 if len(plate_number) == 7 else 480, height))

        for i in range(len(plate_number)):
            if len(plate_number) == 8:
                font_img = self.font_imgs['green_{}'.format(plate_number[i])]
            else:
                if '{}_{}'.format(height, plate_number[i]) in self.font_imgs:
                    font_img = self.font_imgs['{}_{}'.format(height, plate_number[i])]
                else:
                    if i < 2:
                        font_img = self.font_imgs['220_up_{}'.format(plate_number[i])]
                    else:
                        font_img = self.font_imgs['220_down_{}'.format(plate_number[i])]

            # 字符是否红色（同 generate_plate 逻辑）
            is_red = False
            if (i == 0 and plate_number[0] in letters) or plate_number[i] == '警':
                is_red = True
            elif plate_number[i] == '使' and i == 0:
                is_red = True  # 旧版使馆：使在首位 → 红字
            elif plate_number[i] == '领' and len(plate_number) >= 2 and plate_number[1] in letters:
                is_red = True  # 旧版领馆：省A·XXXX领 → 红字
            elif i == 1 and plate_number[0] in letters and np.random.random(1) > 0.5:
                is_red = True

            if enhance:
                k = np.random.randint(1, 6)
                kernel = np.ones((k, k), np.uint8)
                if np.random.random(1) > 0.5:
                    font_img = np.copy(cv2.erode(font_img, kernel, iterations=1))
                else:
                    font_img = np.copy(cv2.dilate(font_img, kernel, iterations=1))

            img_plate_model = copy_to_image_multi(img_plate_model, font_img,
                                                  number_xy[i, :], bg_color, is_red)

        # is_double = 'double' if is_double else 'single'
        img_plate_model = cv2.blur(img_plate_model, (3, 3))

        return img_plate_model


def parse_args():
    parser = argparse.ArgumentParser(description='中国车牌生成器')
    parser.add_argument('--number', default=10, type=int, help='生成车牌数量')
    parser.add_argument('--save-adr', default='multi_val', help='车牌保存路径')
    args = parser.parse_args()
    return args


def mkdir(path):
    try:
        os.makedirs(path)
    except:
        pass


if __name__ == '__main__':
    args = parse_args()
    print(args)
    # 随机生成车牌
    print('save in {}'.format(args.save_adr))

    mkdir(args.save_adr)
    generator = MultiPlateGenerator('plate_model', 'font_model')

    for i in tqdm(range(args.number)):
        img, number_xy, gt_plate_number, bg_color, _ = generator.generate_plate()
        cv2.imwrite(os.path.join(args.save_adr, '{}_{}.jpg'.format(gt_plate_number, bg_color)), img)