import torch
from torch.utils.data import Dataset

from tools import ReadIntArray, OneHotData


class SSQDataset(Dataset):
    def __init__(self, file_path, my_detection, sample_nums):
        self.file_path = file_path
        self.my_detection = my_detection
        self.sample_nums = sample_nums
        self.my_list = ReadIntArray(self.file_path)

    def __getitem__(self, my_index):
        data_list = []
        target_list = []

        start_index = my_index
        for sample_num in range(self.sample_nums):
            data_list.append(self.my_list[start_index + sample_num])

        if self.my_detection in self.my_list[start_index + self.sample_nums]:
            target_list = [1, 0]
        else:
            target_list = [0, 1]

        data_list = OneHotData(data_list)
        data_list = torch.tensor(data_list, dtype=torch.float32)
        target_list = torch.tensor(target_list, dtype=torch.float32)

        return data_list, target_list

    def __len__(self):
        return len(self.my_list) - self.sample_nums