# --------------------------------------------------------------------------------
# 		Import
# --------------------------------------------------------------------------------
import cv2
import json
import os
import torch
from torch.utils.data import Dataset, DataLoader
from utils.data_processing import get_emb
from data_loaders.katakana_dataloader import MyCollate
from PIL import Image
from torchvision import transforms


# --------------------------------------------------------------------------------
#       Classes
# --------------------------------------------------------------------------------
class JPDataset(Dataset):
    def __init__(self, path_data, labels_file, dict_word, max_len):

        self.labels_file = labels_file
        self.max_len = max_len
        self.path_data = path_data
        self.transform_tensor = transforms.ToTensor()

        with open(labels_file, "r", encoding="utf8") as json_file:
            images_labels = json.load(json_file)
            self.path_images = []
            self.labels = []
            for path_image in images_labels:
                self.path_images.append(path_image)
                self.labels.append(images_labels[path_image])

        # Get dictionary
        with open(dict_word, "r", encoding="utf8") as f:
            dictionary = f.readlines()
            self.dictionary = list(map(lambda x: x.strip("\n"), dictionary))

    def __len__(self):
        return len(self.path_images)

    def __getitem__(self, idx):
        label = self.labels[idx]
        # Convert string into index label
        label = get_emb(self.dictionary, label, self.max_len)

        # Get image
        image = Image.open(os.path.join(self.path_data, self.path_images[idx])).convert('L')
        h = image.size[1]
        h_ratio = h / 32
        w_new = round(image.size[0] / h_ratio)
        image = image.resize((w_new, 32), Image.ANTIALIAS)

        image = self.transform_tensor(image)
        label = torch.tensor(label, dtype=torch.float32)

        return image, label


class JPDataLoader(object):
    def __init__(self, batch_size, shuffle, path_data, labels, dictionary, max_len):
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.dataset = JPDataset(path_data, labels, dictionary, max_len)

    def loader(self):
        return DataLoader(self.dataset, batch_size=self.batch_size, shuffle=self.shuffle, collate_fn=collate_fnc)


collate_fnc = MyCollate(False)


# --------------------------------------------------------------------------------
#       Test
# --------------------------------------------------------------------------------
# if __name__ == '__main__':
#     dictionary = "E:/source/Text-Recognition-Using-Seq-To-Seq-Model_1/preprocessing_data/synthesis04/dictionary.txt"
#     data_loader = JPDataLoader(1, True, "E:/data/OCR/synthesis04/labels.json", dictionary, 20)
#     data_loader = data_loader.loader()
#     for batch_idx, (data, label) in enumerate(data_loader):
#         image = data[0]
#         import torchvision.transforms as transforms
#         import matplotlib.pyplot as plt
#         show = transforms.ToPILImage()
#         image = show(image)
#         plt.imshow(image)
#         plt.show()
