from data_provider.data_loader import Dataset_LPI4AI
from torch.utils.data import DataLoader

data_dict = {
    'LPI4AI': Dataset_LPI4AI
}


def data_provider(args, flag):
    Data = data_dict[args.data]

    if flag == 'test':
        shuffle_flag = False
        drop_last = True
        batch_size = args.batch_size
        freq = args.freq
    else:
        shuffle_flag = True
        drop_last = True
        batch_size = args.batch_size
        freq = args.freq

    if args.data == 'LPI4AI':
        data_set = Data(
            root_path=args.root_path,
            data_path=args.data_path,
            flag=flag,
        )
    else:
        print('ERROR: Data Not Supported!')
        exit()
    
    data_loader = DataLoader(
        data_set,
        batch_size=batch_size,
        shuffle=shuffle_flag,
        num_workers=args.num_workers,
        drop_last=drop_last)
    return data_set, data_loader
