import pickle

import numpy as np
import torch
import torch.nn.functional as F
import torch.nn.utils.rnn as rnn_utils
from omegaconf import OmegaConf
from sklearn.model_selection import KFold, StratifiedKFold
from torch import nn
from torch.autograd import Variable
from torch.utils import data
from torch.utils.data import (
    ConcatDataset,
    DataLoader,
    Dataset,
    Subset,
    SubsetRandomSampler,
    TensorDataset,
    random_split,
)


class Dataset(data.Dataset):
    def __init__(self, x, y, x_lab_length):
        self.x = x
        self.y = y
        self.x_lab_length = x_lab_length

    def __getitem__(self, index):  # 返回的是tensor
        return self.x[index], self.y[index], self.x_lab_length[index]

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


def get_dataset(x, y, x_lab_length):
    return Dataset(x, y, x_lab_length)