#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   dataset_lmdb.py
@Time    :   2023/11/06 11:14:32
@Author  :   moyan 
@Contact :   ice_moyan@163.com
'''
import os
import cv2
import six
import lmdb
import string
import pickle
import numpy as np
from PIL import Image
from torch.utils.data import Dataset

class LMDBDataSet(Dataset):
    def __init__(self, data_dir, do_shuffle=False):
        super(LMDBDataSet, self).__init__()
        self.do_shuffle = do_shuffle
        
        self.lmdb_sets = self.load_hierarchical_lmdb_dataset(data_dir)
        print("Initialize indexs of datasets:%s" % data_dir)
        self.data_idx_order_list = self.dataset_traversal()

        if self.do_shuffle:
            np.random.shuffle(self.data_idx_order_list)

    def load_hierarchical_lmdb_dataset(self, data_dir):
        lmdb_sets = {}
        dataset_idx = 0
        for dirpath, dirnames, filenames in os.walk(data_dir + '/'):
            if not dirnames:
                # # moyan debug
                # if os.path.basename(dirpath) == "gen_text_char9k_600w_moyan_230816_train":
                #     continue
                env = lmdb.open(
                    dirpath,
                    max_readers=32,
                    readonly=True,
                    lock=False,
                    readahead=False,
                    meminit=False)
                txn = env.begin(write=False)
                num_samples = int(txn.get('num-samples'.encode()))
                lmdb_sets[dataset_idx] = {"dirpath":dirpath, "env":env, \
                    "txn":txn, "num_samples":num_samples}
                dataset_idx += 1
        return lmdb_sets

    def dataset_traversal(self):
        lmdb_num = len(self.lmdb_sets)
        total_sample_num = 0
        for lno in range(lmdb_num):
            total_sample_num += self.lmdb_sets[lno]['num_samples']
        data_idx_order_list = np.zeros((total_sample_num, 2))
        beg_idx = 0
        for lno in range(lmdb_num):
            tmp_sample_num = self.lmdb_sets[lno]['num_samples']
            end_idx = beg_idx + tmp_sample_num
            data_idx_order_list[beg_idx:end_idx, 0] = lno
            data_idx_order_list[beg_idx:end_idx, 1] \
                = list(range(tmp_sample_num))
            data_idx_order_list[beg_idx:end_idx, 1] += 1
            beg_idx = beg_idx + tmp_sample_num
        return data_idx_order_list

    def get_img_data(self, value):
        """get_img_data"""
        if not value:
            return None
        imgdata = np.frombuffer(value, dtype='uint8')
        if imgdata is None:
            return None
        imgori = cv2.imdecode(imgdata, 1)
        if imgori is None:
            return None
        return imgori

    def get_lmdb_sample_info(self, txn, index):
        label_key = 'label-%09d'.encode() % index
        label = txn.get(label_key)
        if label is None:
            return None
        label = label.decode('utf-8')
        img_key = 'image-%09d'.encode() % index
        imgbuf = txn.get(img_key)
        return imgbuf, label

    def __getitem__(self, idx):
        lmdb_idx, file_idx = self.data_idx_order_list[idx]
        lmdb_idx = int(lmdb_idx)
        file_idx = int(file_idx)
        sample_info = self.get_lmdb_sample_info(self.lmdb_sets[lmdb_idx]['txn'], file_idx)
        if sample_info is None:
            return self.__getitem__(np.random.randint(self.__len__()))
        img, label = sample_info
        data = {'image': self.get_img_data(img), 'label': label}
        return data

    def __len__(self):
        return self.data_idx_order_list.shape[0]


def main():
    data_dir = "/mnt/e/DataSet/OCR/HW_Chinese/lmdb/HW_Chinese_Test/"
    dataset = LMDBDataSet(data_dir, do_shuffle=False)

    for data in dataset:
        print(data['image'].size, data['label'])


if __name__=='__main__':
    main()
