# -*- coding: utf-8 -*-
# @Time    : 2019/8/23 21:52
# @Author  : zhoujun
import copy

import PIL
import numpy as np
import paddle
from paddle.io import DataLoader, DistributedBatchSampler, BatchSampler

from paddle.vision import transforms


def get_dataset(data_path, module_name, transform, dataset_args):
    """
    获取训练dataset
    :param data_path: dataset文件列表，每个文件内以如下格式存储 ‘path/to/Crop_img\tlabel’
    :param module_name: 所使用的自定义dataset名称，目前只支持data_loaders.ImageDataset
    :param transform: 该数据集使用的transforms
    :param dataset_args: module_name的参数
    :return: 如果data_path列表不为空，返回对于的ConcatDataset对象，否则None
    """
    from . import dataset

    s_dataset = getattr(dataset, module_name)(
        transform=transform, data_path=data_path, **dataset_args
    )
    return s_dataset


def get_transforms(transforms_config):
    tr_list = []
    for item in transforms_config:
        if "args" not in item:
            args = {}
        else:
            args = item["args"]
        cls = getattr(transforms, item["type"])(**args)
        tr_list.append(cls)
    tr_list = transforms.Compose(tr_list)
    return tr_list


class ICDARCollectFN:
    def __init__(self, *args, **kwargs):
        pass

    def __call__(self, batch):
        data_dict = {}
        to_tensor_keys = []
        for sample in batch:
            for k, v in sample.items():
                if k not in data_dict:
                    data_dict[k] = []
                if isinstance(v, (np.ndarray, paddle.Tensor, PIL.Image.Image)):
                    if k not in to_tensor_keys:
                        to_tensor_keys.append(k)
                data_dict[k].append(v)
        for k in to_tensor_keys:
            data_dict[k] = paddle.stack(data_dict[k], 0)
        return data_dict


def get_dataloader(module_config, distributed=False):
    if module_config is None:
        return None
    config = copy.deepcopy(module_config)
    dataset_args = config["dataset"]["args"]
    if "transforms" in dataset_args:
        img_transfroms = get_transforms(dataset_args.pop("transforms"))
    else:
        img_transfroms = None
    # 创建数据集
    dataset_name = config["dataset"]["type"]
    data_path = dataset_args.pop("data_path")
    if data_path == None:
        return None

    data_path = [x for x in data_path if x is not None]
    if len(data_path) == 0:
        return None
    if (
        "collate_fn" not in config["loader"]
        or config["loader"]["collate_fn"] is None
        or len(config["loader"]["collate_fn"]) == 0
    ):
        config["loader"]["collate_fn"] = None
    else:
        config["loader"]["collate_fn"] = eval(config["loader"]["collate_fn"])()

    _dataset = get_dataset(
        data_path=data_path,
        module_name=dataset_name,
        transform=img_transfroms,
        dataset_args=dataset_args,
    )
    sampler = None
    if distributed:
        # 3）使用DistributedSampler
        batch_sampler = DistributedBatchSampler(
            dataset=_dataset,
            batch_size=config["loader"].pop("batch_size"),
            shuffle=config["loader"].pop("shuffle"),
        )
    else:
        batch_sampler = BatchSampler(
            dataset=_dataset,
            batch_size=config["loader"].pop("batch_size"),
            shuffle=config["loader"].pop("shuffle"),
        )
    loader = DataLoader(
        dataset=_dataset, batch_sampler=batch_sampler, **config["loader"]
    )
    return loader