from data_provider.data_loader import Dataset_public,Dataset_synthetic,Dataset_MMD
from torch.utils.data import DataLoader

data_dict = {
    'synthetic': Dataset_synthetic,
    'synthetic_MMD': Dataset_MMD,
    'Etth1': Dataset_public,
    'Etth1_MMD': Dataset_MMD,
    'Etth2': Dataset_public,
    'Etth2_MMD': Dataset_MMD,
    'Ettm1': Dataset_public,
    'Ettm1_MMD': Dataset_MMD,
    'Ettm2': Dataset_public,
    'Ettm2_MMD': Dataset_MMD,
    'Electricity': Dataset_public,
    'Electricity_MMD': Dataset_MMD,
}


def data_provider(args, flag):
    Data = data_dict[args.data]
    timeenc = 0 if args.embed != 'timeF' else 1
    train_only = args.train_only

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

    data_set = Data(
        args=args,
        root_path=args.root_path,
        traindata_path=args.traindata_path,
        valdata_path=args.valdata_path,
        predict_path=args.predict_path,
        flag=flag,
        size=[args.seq_len, args.label_len, args.pred_len],
        features=args.features,
        target=args.target,
        timeenc=timeenc,
        freq=freq,
        train_only=train_only
    )
    print(flag, len(data_set))
    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
