from PIL import Image
import numpy as np
import random
import string
from tqdm import tqdm
from config import *


class LabelHandler:

    labels = string.digits + string.ascii_uppercase
    classes = len(labels)

    def __init__(self, m):
        self.word_nums = m

    def encode(self, words):
        res = np.zeros(shape=(self.word_nums, self.classes), dtype=np.float32)
        for i, word in enumerate(words):
            res[i][self.labels.index(word)] = 1.0
        return res

    def decode(self, matrix):
        return "".join(self.labels[np.argmax(row)] for row in matrix)


class LabelHandlerCTC:

    labels = string.digits + string.ascii_uppercase + " "
    classes = len(labels)

    def __init__(self, m):
        self.word_nums = m

    def encode(self, words):
        res = np.zeros(shape=(self.word_nums, ), dtype=np.float32)
        for i, word in enumerate(words):
            res[i] = self.labels.index(word)
        return res

    def decode(self, matrix):
        # print(matrix)
        return "".join(self.labels[int(idx)] for idx in matrix)


def img_pretreat(file):
    # print(file)
    row_img = Image.open(file)  # 打开图片
    gray_img = row_img.convert('L')  # 转灰度图
    gray_img = gray_img.resize((Width, Height))
    points = np.expand_dims(np.asarray(gray_img, dtype=np.float32), axis=-1)  # 图片转数组增加一个维度
    points = (points - 127.5) / 255

    return points


def process(ratio=0.0):
    """
    构造数据集
    :param ratio: 测试集比例
    :return:
    """
    inputs = []
    labels = []

    test_inputs = []
    test_labels = []

    # lh = LabelHandler(NUM)
    lh = LabelHandlerCTC(NUM)
    with open(label_path, 'r') as f:
        lines = f.readlines()
    random.shuffle(lines)
    n = len(lines)
    for i, line in tqdm(enumerate(lines)):
        filename, label = line.rstrip("\n").split(":")
        a_input = img_pretreat(os.path.join(pic_path, filename+".png"))
        # print(label)
        a_label = lh.encode(label)
        if i < n*(1-ratio):
            inputs.append(a_input)
            labels.append(a_label)
        else:
            test_inputs.append(a_input)
            test_labels.append(a_label)

    inputs = np.asarray(inputs)
    labels = np.asarray(labels)

    np.save(os.path.join(train_path, 'inputs.npy'), inputs)
    np.save(os.path.join(train_path, 'labels_ctc.npy'), labels)

    np.save(os.path.join(test_path, 'inputs.npy'), test_inputs)
    np.save(os.path.join(test_path, 'labels_ctc.npy'), test_labels)


if __name__ == "__main__":
    # save_file()
    # img_pretreat("../resources/dataset/pics/11254-0.png")
    process(0.02)