# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import numpy as np
import os
import lmdb
import random
import signal
import paddle
from paddle.io import Dataset, DataLoader, DistributedBatchSampler, BatchSampler

from .imaug import transform, create_operators
from ppocr.utils.logging import get_logger


def term_mp(sig_num, frame):
    """ kill all child processes
    """
    pid = os.getpid()
    pgid = os.getpgid(os.getpid())
    print("main proc {} exit, kill process group " "{}".format(pid, pgid))
    os.killpg(pgid, signal.SIGKILL)


signal.signal(signal.SIGINT, term_mp)
signal.signal(signal.SIGTERM, term_mp)


class ModeException(Exception):
    """
    ModeException
    """

    def __init__(self, message='', mode=''):
        message += "\nOnly the following 3 modes are supported: " \
                   "train, valid, test. Given mode is {}".format(mode)
        super(ModeException, self).__init__(message)


class SampleNumException(Exception):
    """
    SampleNumException
    """

    def __init__(self, message='', sample_num=0, batch_size=1):
        message += "\nError: The number of the whole data ({}) " \
                   "is smaller than the batch_size ({}), and drop_last " \
                   "is turnning on, so nothing  will feed in program, " \
                   "Terminated now. Please reset batch_size to a smaller " \
                   "number or feed more data!".format(sample_num, batch_size)
        super(SampleNumException, self).__init__(message)


def get_file_list(file_list, data_dir, delimiter='\t'):
    """
    read label list from file and shuffle the list

    Args:
        params(dict):
    """
    if isinstance(file_list, str):
        file_list = [file_list]
    data_source_list = []
    for file in file_list:
        with open(file) as f:
            full_lines = [line.strip() for line in f]
            for line in full_lines:
                try:
                    img_path, label = line.split(delimiter)
                except:
                    logger = get_logger()
                    logger.warning('label error in {}'.format(line))
                img_path = os.path.join(data_dir, img_path)
                data = {'img_path': img_path, 'label': label}
                data_source_list.append(data)
    return data_source_list


class LMDBDateSet(Dataset):
    def __init__(self, config, global_config):
        super(LMDBDateSet, self).__init__()
        self.data_list = self.load_lmdb_dataset(
            config['file_list'], global_config['max_text_length'])
        random.shuffle(self.data_list)

        self.ops = create_operators(config['transforms'], global_config)

        # for rec
        character = ''
        for op in self.ops:
            if hasattr(op, 'character'):
                character = getattr(op, 'character')

        self.info_dict = {'character': character}

    def load_lmdb_dataset(self, data_dir, max_text_length):
        self.env = lmdb.open(
            data_dir,
            max_readers=32,
            readonly=True,
            lock=False,
            readahead=False,
            meminit=False)
        if not self.env:
            print('cannot create lmdb from %s' % (data_dir))
            exit(0)

        filtered_index_list = []
        with self.env.begin(write=False) as txn:
            nSamples = int(txn.get('num-samples'.encode()))
            self.nSamples = nSamples
            for index in range(self.nSamples):
                index += 1  # lmdb starts with 1
                label_key = 'label-%09d'.encode() % index
                label = txn.get(label_key).decode('utf-8')
                if len(label) > max_text_length:
                    # print(f'The length of the label is longer than max_length: length
                    # {len(label)}, {label} in dataset {self.root}')
                    continue

                # By default, images containing characters which are not in opt.character are filtered.
                # You can add [UNK] token to `opt.character` in utils.py instead of this filtering.
                filtered_index_list.append(index)
        return filtered_index_list

    def print_lmdb_sets_info(self, lmdb_sets):
        lmdb_info_strs = []
        for dataset_idx in range(len(lmdb_sets)):
            tmp_str = " %s:%d," % (lmdb_sets[dataset_idx]['dirpath'],
                                   lmdb_sets[dataset_idx]['num_samples'])
            lmdb_info_strs.append(tmp_str)
        lmdb_info_strs = ''.join(lmdb_info_strs)
        logger = get_logger()
        logger.info("DataSummary:" + lmdb_info_strs)
        return

    def __getitem__(self, idx):
        idx = self.data_list[idx]
        with self.env.begin(write=False) as txn:
            label_key = 'label-%09d'.encode() % idx
            label = txn.get(label_key)
            if label is not None:
                label = label.decode('utf-8')
                img_key = 'image-%09d'.encode() % idx
                imgbuf = txn.get(img_key)
                data = {'image': imgbuf, 'label': label}
                outs = transform(data, self.ops)
            else:
                outs = None
            if outs is None:
                return self.__getitem__(np.random.randint(self.__len__()))
            return outs

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


class SimpleDataSet(Dataset):
    def __init__(self, config, global_config):
        super(SimpleDataSet, self).__init__()
        delimiter = config.get('delimiter', '\t')
        self.data_list = get_file_list(config['file_list'], config['data_dir'],
                                       delimiter)
        random.shuffle(self.data_list)

        self.ops = create_operators(config['transforms'], global_config)

        # for rec
        character = ''
        for op in self.ops:
            if hasattr(op, 'character'):
                character = getattr(op, 'character')

        self.info_dict = {'character': character}

    def __getitem__(self, idx):
        data = copy.deepcopy(self.data_list[idx])
        with open(data['img_path'], 'rb') as f:
            img = f.read()
            data['image'] = img
        outs = transform(data, self.ops)
        if outs is None:
            return self.__getitem__(np.random.randint(self.__len__()))
        return outs

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


class BatchBalancedDataLoader(object):
    def __init__(self,
                 dataset_list: list,
                 ratio_list: list,
                 distributed,
                 device,
                 loader_args: dict):
        """
        对datasetlist里的dataset按照ratio_list里对应的比例组合，似的每个batch里的数据按按照比例采样的
        :param dataset_list: 数据集列表
        :param ratio_list: 比例列表
        :param loader_args: dataloader的配置
        """
        assert sum(ratio_list) == 1 and len(dataset_list) == len(ratio_list)

        self.dataset_len = 0
        self.data_loader_list = []
        self.dataloader_iter_list = []
        all_batch_size = loader_args.pop('batch_size')
        batch_size_list = list(
            map(int, [max(1.0, all_batch_size * x) for x in ratio_list]))
        remain_num = all_batch_size - sum(batch_size_list)
        batch_size_list[np.argmax(ratio_list)] += remain_num

        for _dataset, _batch_size in zip(dataset_list, batch_size_list):
            if distributed:
                batch_sampler_class = DistributedBatchSampler
            else:
                batch_sampler_class = BatchSampler
            batch_sampler = batch_sampler_class(
                dataset=_dataset,
                batch_size=_batch_size,
                shuffle=loader_args['shuffle'],
                drop_last=loader_args['drop_last'], )
            _data_loader = DataLoader(
                dataset=_dataset,
                batch_sampler=batch_sampler,
                places=device,
                num_workers=loader_args['num_workers'],
                return_list=True, )
            self.data_loader_list.append(_data_loader)
            self.dataloader_iter_list.append(iter(_data_loader))
            self.dataset_len += len(_dataset)

    def __iter__(self):
        return self

    def __len__(self):
        return min([len(x) for x in self.data_loader_list])

    def __next__(self):
        batch = []
        for i, data_loader_iter in enumerate(self.dataloader_iter_list):
            try:
                _batch_i = next(data_loader_iter)
                batch.append(_batch_i)
            except StopIteration:
                self.dataloader_iter_list[i] = iter(self.data_loader_list[i])
                _batch_i = next(self.dataloader_iter_list[i])
                batch.append(_batch_i)
            except ValueError:
                pass
        if len(batch) > 0:
            batch_list = []
            batch_item_size = len(batch[0])
            for i in range(batch_item_size):
                cur_item_list = [batch_i[i] for batch_i in batch]
                batch_list.append(paddle.concat(cur_item_list, axis=0))
        else:
            batch_list = batch[0]
        return batch_list


def fill_batch(batch):
    """
    2020.09.08： The current paddle version only supports returning data with the same length.
                Therefore, fill in the batches with inconsistent lengths.
                this method is currently only useful for text detection
    """
    keys = list(range(len(batch[0])))
    v_max_len_dict = {}
    for k in keys:
        v_max_len_dict[k] = max([len(item[k]) for item in batch])
    for item in batch:
        length = []
        for k in keys:
            v = item[k]
            length.append(len(v))
            assert isinstance(v, np.ndarray)
            if len(v) == v_max_len_dict[k]:
                continue
            try:
                tmp_shape = [v_max_len_dict[k] - len(v)] + list(v[0].shape)
            except:
                a = 1
            tmp_array = np.zeros(tmp_shape, dtype=v[0].dtype)
            new_array = np.concatenate([v, tmp_array])
            item[k] = new_array
        item.append(length)
    return batch