import torch


def Loading_Data(file_path, sub_sampling_steps, resolution, num_train, num_test, batch_size, input_time_slices_num, next_time_slices_num, inverse = False):


    train_data = torch.load(file_path + '/2D_GS.pt')
    x_train = train_data[:num_train,::sub_sampling_steps,::sub_sampling_steps,:input_time_slices_num, :]
    x_train = x_train.reshape(x_train.size(0), x_train.size(1), x_train.size(2), -1)
    y_train = train_data[:num_train,::sub_sampling_steps,::sub_sampling_steps,input_time_slices_num:next_time_slices_num+input_time_slices_num, :]
    # y_train = y_train.reshape(y_train.size(0), y_train.size(1), y_train.size(2), -1)

    test_data = torch.load(file_path + '/2D_GS.pt')
    x_test = test_data[-num_test:,::sub_sampling_steps,::sub_sampling_steps,:input_time_slices_num, :]
    x_test = x_test.reshape(x_test.size(0), x_test.size(1), x_test.size(2), -1)
    y_test = test_data[-num_test:,::sub_sampling_steps,::sub_sampling_steps,input_time_slices_num:next_time_slices_num+input_time_slices_num, :]
    # y_test = y_test.reshape(y_test.size(0), y_test.size(1), y_test.size(2), -1)


    print(y_train.shape)
    print(y_test.shape)

    # x_train = x_train.reshape(num_train,resolution,resolution,input_time_slices_num)
    # x_test = x_test.reshape(num_test,resolution,resolution,input_time_slices_num)

    train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_train, y_train), batch_size=batch_size, shuffle=True)
    test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_test, y_test), batch_size=batch_size, shuffle=False)

    return x_test, y_test, train_loader, test_loader
