import torch
import os
import random
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from PIL import Image
import numpy as np
import collections
import numbers
import math
import pandas as pd
from sklearn.preprocessing import StandardScaler, MinMaxScaler
import pickle


def data_name2nc(data_name: str):
    # data_name = data_name.upper()
    nc = 0
    if data_name == 'MSL':
        nc = 55
    elif data_name == 'SMAP':
        nc = 25
    elif data_name == 'SMD':
        nc = 38
    elif data_name == 'PSM':
        nc = 25
    elif data_name == 'SWAT':
        nc = 51 - 1
    elif data_name == 'WADI':
        nc = 123 - 1
    elif data_name == 'UCR' or data_name == 'UCR_AUG':
        nc = 1
    elif data_name == 'NIPS_TS_Swan':
        nc = 38
    elif data_name == 'NIPS_TS_Water':
        nc = 9
    elif data_name.startswith('SP'):
        nc = 38
    elif data_name == 'MSL1':
        nc = 1

    else:
        raise ValueError(f'no this dataset {data_name}')
    return nc

class GetDataset(object):
    def load_data(self):
        assert hasattr(self,'train') and self.train is not None, "subclass must have self.train"
        assert self.test is not None, "subclass must have self.test"
        assert self.test_labels is not None, "subclass must have self.test_labels"
        return self.train, self.test, self.test_labels

class PSMSegLoader(GetDataset):
    def __init__(self, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = pd.read_csv(data_path + '/train.csv')
        data = data.values[:, 1:]
        data = np.nan_to_num(data)
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = pd.read_csv(data_path + '/test.csv')
        test_data = test_data.values[:, 1:]
        test_data = np.nan_to_num(test_data)
        self.test = self.scaler.transform(test_data)
        self.train = data
        self.val = self.test
        self.test_labels = pd.read_csv(data_path + '/test_label.csv').values[:, 1:]
        self.test_labels = self.test_labels.reshape(-1)

        print("train:", self.train.shape)
        print("test:", self.test.shape)

    def __len__(self):
        """
        Number of images in the object dataset.
        """
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])


class MSLSegLoader(GetDataset):
    def __init__(self, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = np.load(data_path + "/MSL_train.npy")
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/MSL_test.npy")
        self.test = self.scaler.transform(test_data)
        self.train = data
        self.val = self.test
        self.test_labels = np.load(data_path + "/MSL_test_label.npy")

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])
        


class MSL1SegLoader(GetDataset):
    def __init__(self, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.win_size = win_size
        data_path = data_path.replace("MSL1", "MSL")
        self.scaler = StandardScaler()
        data = np.load(data_path + "/MSL_train.npy")[:,0:1]
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/MSL_test.npy")[:,0:1]
        self.test = self.scaler.transform(test_data)
        self.train = data
        self.val = self.test
        self.test_labels = np.load(data_path + "/MSL_test_label.npy")

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])


class SMAPSegLoader(GetDataset):
    def __init__(self, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = np.load(data_path + "/SMAP_train.npy")
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/SMAP_test.npy")
        self.test = self.scaler.transform(test_data)
        self.train = data
        self.val = self.test
        self.test_labels = np.load(data_path + "/SMAP_test_label.npy")

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train": #train and val did not use label
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])


class SMDSegLoader(GetDataset):
    def __init__(self, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = np.load(data_path + "/SMD_train.npy")[:,:]
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/SMD_test.npy")[:,:]
        self.test = self.scaler.transform(test_data)
        self.train = data
        data_len = len(self.train)
        self.val = self.train[(int)(data_len * 0.8):]
        self.test_labels = np.load(data_path + "/SMD_test_label.npy")[:]

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])

        
        
class UCRSegLoader(GetDataset):
    def __init__(self, index, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.index = index
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = np.load(data_path + "/UCR_"+str(index)+"_train.npy")
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/UCR_"+str(index)+"_test.npy")
        self.test = self.scaler.transform(test_data)

        self.train = data
        self.val = self.test
        self.test_labels = np.load(data_path + "/UCR_"+str(index)+"_test_label.npy")
        if self.mode == "val":
            print("train:", self.train.shape)
            print("test:", self.test.shape)

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])    
        

class UCRAUGSegLoader(GetDataset):
    def __init__(self, index, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.index = index
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = np.load(data_path + "/UCR_AUG_"+str(index)+"_train.npy")
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/UCR_AUG_"+str(index)+"_test.npy")
        self.test = self.scaler.transform(test_data)

        self.train = data
        self.val = self.test
        self.test_labels = np.load(data_path + "/UCR_AUG_"+str(index)+"_test_label.npy")
        if self.mode == "val":
            print("train:", self.train.shape)
            print("test:", self.test.shape)

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]) 
        

class NIPS_TS_WaterSegLoader(GetDataset):
    def __init__(self, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = np.load(data_path + "/NIPS_TS_Water_train.npy")
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/NIPS_TS_Water_test.npy")
        self.test = self.scaler.transform(test_data)

        self.train = data
        self.val = self.test
        self.test_labels = np.load(data_path + "/NIPS_TS_Water_test_label.npy")
        print("test:", self.test.shape)
        print("train:", self.train.shape)

    def __len__(self):

        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])        
        
        
        
class NIPS_TS_SwanSegLoader(GetDataset):
    def __init__(self, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = np.load(data_path + "/NIPS_TS_Swan_train.npy")
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/NIPS_TS_Swan_test.npy")
        self.test = self.scaler.transform(test_data)

        self.train = data
        self.val = self.test
        self.test_labels = np.load(data_path + "/NIPS_TS_Swan_test_label.npy")
        print("test:", self.test.shape)
        print("train:", self.train.shape)

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]) 
        


class WADI_SegLoader(GetDataset):
    def __init__(self, data_path, win_size, step, mode="train", ignore=(102,), scaler=False):
        self.mode = mode
        self.step = step
        self.win_size = win_size
        self.scaler = MinMaxScaler()
        data = np.load(data_path + "/WADI_train.npy")
        test_data = np.load(data_path + "/WADI_test.npy")
        
        if ignore is not None:
            data = np.delete(data, ignore, axis=1)
            test_data = np.delete(test_data, ignore, axis=1)

        if scaler:
            data = self.scaler.fit_transform(data)
            test_data = self.scaler.fit_transform(test_data)

        self.train = data
        self.val = self.test = test_data
        self.test_labels = np.load(data_path + "/WADI_test_label.npy")
        print("test:", self.test.shape)
        print("train:", self.train.shape)

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]) 

class NIPS_TS_CCardSegLoader(GetDataset):
    def __init__(self, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = np.load(data_path + "/NIPS_TS_creditcard_train.npy")
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/NIPS_TS_creditcard_test.npy")
        self.test = self.scaler.transform(test_data)

        self.train = data
        self.val = self.test
        self.test_labels = np.load(data_path + "/NIPS_TS_creditcard_test_label.npy")

    def __len__(self):

        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]) 
        
class SMD_Ori_Pikled_SegLoader(GetDataset):
    def __init__(self, index, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.index = index
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = np.load(data_path + "/machine-"+str(index)+"_train.pkl", allow_pickle=True)
        # data = pickle.load(data_path + "/machine-"+str(index)+"_train.pkl")
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/machine-"+str(index)+"_test.pkl", allow_pickle=True)
        # test_data = pickle.load(data_path + "/machine-"+str(index)+"_test.pkl")
        self.test = self.scaler.transform(test_data)

        self.train = data
        self.val = self.test
        self.test_labels = np.load(data_path + "/machine-"+str(index)+"_test_label.pkl", allow_pickle=True)
        # self.test_labels = pickle.load(data_path + "/machine-"+str(index)+"_test_label.pkl")
        if self.mode == "val":
            print("train:", self.train.shape)
            print("test:", self.test.shape)

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]) 
        
        
class SMD_OriSegLoader(GetDataset):
    def __init__(self, index, data_path, win_size, step, mode="train"):
        self.mode = mode
        self.step = step
        self.index = index
        self.win_size = win_size
        self.scaler = StandardScaler()
        data = np.load(data_path + "/SMD_Ori_"+str(index)+"_train.npy")
        self.scaler.fit(data)
        data = self.scaler.transform(data)
        test_data = np.load(data_path + "/SMD_Ori_"+str(index)+"_test.npy")
        self.test = self.scaler.transform(test_data)

        self.train = data
        self.val = self.test
        self.test_labels = np.load(data_path + "/SMD_Ori_"+str(index)+"_test_label.npy")
        if self.mode == "val":
            print("train:", self.train.shape)
            print("test:", self.test.shape)

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.mode == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])         


class SWATSegLoader(Dataset,GetDataset):
    def __init__(self, root_path, win_size, step, flag="train", ignore=(10,), scaler=1):
        self.flag = flag
        self.step = step
        self.win_size = win_size
        self.scaler = StandardScaler()

        train_data = pd.read_csv(os.path.join(root_path, 'swat_train2.csv'))
        test_data = pd.read_csv(os.path.join(root_path, 'swat2.csv'))
        labels = test_data.values[:, -1:]
        train_data = train_data.values[:, :-1]
        test_data = test_data.values[:, :-1]
        if ignore is not None:
            train_data = np.delete(train_data, ignore, axis=1)
            test_data = np.delete(test_data, ignore, axis=1)

        if scaler:
            train_data = self.scaler.fit_transform(train_data)
            test_data = self.scaler.fit_transform(test_data)

        self.train = train_data
        self.test = test_data
        data_len = len(self.train)
        self.val = self.train[(int)(data_len * 0.8):]
        self.test_labels = labels.reshape(-1)
        print("test:", self.test.shape)
        print("train:", self.train.shape)

    def __len__(self):
        """
        Number of images in the object dataset.
        """
        if self.flag == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.flag == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.flag == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.flag == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.flag == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.flag == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])
        


class WAQ_SegLoader(Dataset):
    def __init__(self, root_path, win_size, step, flag="train"):
        self.flag = flag
        self.step = step
        self.win_size = win_size
        self.scaler = StandardScaler()

        train_data = pd.read_csv(os.path.join(root_path, 'WAQ_train.csv'))
        test_data = pd.read_csv(os.path.join(root_path, 'WAQ_test.csv'))
        labels = test_data.values[:, -1:]
        train_data = train_data.values[:, :-1]
        test_data = test_data.values[:, :-1]

        self.scaler.fit(train_data)
        train_data = self.scaler.transform(train_data)
        test_data = self.scaler.transform(test_data)
        self.train = train_data
        self.test = test_data
        data_len = len(self.train)
        self.val = self.train[(int)(data_len * 0.8):]
        self.test_labels = labels
        print("test:", self.test.shape)
        print("train:", self.train.shape)

    def __len__(self):
        """
        Number of images in the object dataset.
        """
        if self.flag == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.flag == 'val'):
            return (self.val.shape[0] - self.win_size) // self.step + 1
        elif (self.flag == 'test'):
            return (self.test.shape[0] - self.win_size) // self.step + 1
        else:
            return (self.test.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.flag == "train":
            return np.float32(self.train[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.flag == 'val'):
            return np.float32(self.val[index:index + self.win_size]), np.zeros(self.win_size)
        elif (self.flag == 'test'):
            return np.float32(self.test[index:index + self.win_size]), np.float32(
                self.test_labels[index:index + self.win_size])
        else:
            return np.float32(self.test[
                              index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])
        
        
def get_loader_segment(index, data_path, batch_size, win_size=100, step=1, mode='train', dataset='MSL', dist=1, ret_data=False):
    if (dataset == 'SMD'):
        dataset = SMDSegLoader(data_path, win_size, step, mode)
    elif (dataset == 'MSL'):
        dataset = MSLSegLoader(data_path, win_size, step, mode)
    elif (dataset == 'MSL1'):
        dataset = MSL1SegLoader(data_path, win_size, step, mode)
    elif (dataset == 'SMAP'):
        dataset = SMAPSegLoader(data_path, win_size, step, mode)
    elif (dataset == 'SWAT'):
        dataset = SWATSegLoader(data_path, win_size, step, mode)
    elif (dataset == 'PSM'):
        dataset = PSMSegLoader(data_path, win_size, step, mode)
    elif (dataset == 'WADI'):
        dataset = WADI_SegLoader(data_path, win_size, step, mode)
    elif (dataset == 'WAQ'):
        dataset = WAQ_SegLoader(data_path, win_size, step, mode)
    elif (dataset == 'UCR'):
        dataset = UCRSegLoader(index, data_path, win_size, step, mode)
    elif (dataset == 'UCR_AUG'):
        dataset = UCRAUGSegLoader(index, data_path, win_size, step, mode)
    elif (dataset == 'NIPS_TS_Water'):
        dataset = NIPS_TS_WaterSegLoader(data_path, win_size, step, mode)
    elif (dataset == 'NIPS_TS_Swan'):
        dataset = NIPS_TS_SwanSegLoader(data_path, win_size, step, mode)
    elif (dataset == 'NIPS_TS_Creditcard'):
        dataset = NIPS_TS_CCardSegLoader(data_path, win_size, step, mode)
    elif (dataset == 'SMD_Ori'):
        dataset = SMD_OriSegLoader(index, data_path, win_size, step, mode)
    elif (dataset == 'SMD_Ori_Pikled'):
        dataset = SMD_Ori_Pikled_SegLoader(index, data_path, win_size, step, mode)


    if ret_data:
        return dataset
    
    shuffle = False
    if mode == 'train':
        shuffle = True

    if dist:
        sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle, drop_last=False)
    else:
        sampler = None

    data_loader = DataLoader(dataset=dataset,
                             batch_size=batch_size,
                             sampler=sampler,
                             shuffle=shuffle,
                             num_workers=8,
                             drop_last=False)
    return data_loader


class SupervisedDataset(Dataset):
    """
    选择split比例划分训练集和测试集
    """
    def __init__(self, data_y, data_labels, win_size, step, mode="train", train_test_split=0.8):
        self.mode = mode
        self.step = step
        self.win_size = win_size
        self.scaler = StandardScaler()
        self.train_y = data_y[:(int)(len(data_y) * train_test_split)]
        self.train_labels = data_labels[:(int)(len(data_labels) * train_test_split)]
        self.test_y = data_y[int(len(data_y) * train_test_split):]
        self.test_labels = data_labels[int(len(data_labels) * train_test_split):]
        self.train_test_split = train_test_split
        self.scaler.fit(self.train_y)
        self.train_y = self.scaler.transform(self.train_y)
        self.test_y = self.scaler.transform(self.test_y)
        self.all_test_y = np.concatenate([self.train_y, self.test_y], axis=0)
        self.all_test_labels = np.concatenate([self.train_labels, self.test_labels], axis=0)
        self.train = self.train_y.astype(np.float32)
        self.train_y = self.train
        self.test = self.test_y.astype(np.float32)
        self.test_y = self.test
        self.test_labels = self.test_labels.astype(np.float32)
        self.train_labels = self.train_labels.astype(np.float32)
        

    def __len__(self):
        if self.mode == "train":
            return (self.train.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test.shape[0] - self.win_size) // self.win_size + 1
        elif (self.mode == 'all_test' or self.mode == 'all'):
            return (self.all_test_y.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.float32(
                self.train_labels[index:index + self.win_size])
        elif (self.mode == 'test'):
            return np.float32(self.test[
                                index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])
        elif (self.mode == 'all_test' or self.mode == 'all'):
            return np.float32(self.all_test_y[
                                index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.all_test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])
        

class RandomSupervisedDataset(Dataset):
    """
    采样anomaly_ratio的比率，保证训练集中有anomaly_ratio比例的异常样本
    固定随机种子seed，保证每次采样结果一致
    """
    def __init__(self, data_y, data_labels, win_size, step=1, mode="train",
                 anomaly_ratio=0.1,
                 train_split_max=0.3,
                 seed=42):
        self.mode = mode
        self.step = step
        # set seed
        # np.random.seed(seed)
        # torch.manual_seed(seed)
        # random.seed(seed)
        # 找到索引i，使得 sum(data_labels[:i])/i >= anomaly_ratio
        idx = int(data_labels.shape[0] * anomaly_ratio)
        train_idx_max = int(data_y.shape[0] * train_split_max)
        while data_labels[:idx].sum() / idx < anomaly_ratio and idx < train_idx_max:
            idx += 1
        self.train_anomaly_ratio = train_anomaly_ratio = data_labels[:idx].sum() / idx
        print(f"train anomaly ratio: {train_anomaly_ratio}, train length: {idx}")
        self.test_anomaly_ratio = test_anomaly_ratio = data_labels[idx:].sum() / (len(data_labels) - idx)
        print(f"test anomaly ratio: {test_anomaly_ratio}, test length: {len(data_labels) - idx}")
        train_y = data_y[:idx]
        train_labels = data_labels[:idx]
        test_y = data_y[idx:]
        test_labels = data_labels[idx:]
        self.win_size = win_size
        self.scaler = StandardScaler()
        self.train_y = train_y
        self.train_labels = train_labels
        self.test_y = test_y
        self.test_labels = test_labels
        self.scaler.fit(self.train_y)
        self.train_y = self.train = self.scaler.transform(self.train_y)
        self.test_y = self.test = self.scaler.transform(self.test_y)
        self.all_test_y = np.concatenate([self.train_y, self.test_y], axis=0)
        self.all_test_labels = np.concatenate([self.train_labels, self.test_labels], axis=0)

    def __len__(self):
        if self.mode == "train":
            return (self.train_y.shape[0] - self.win_size) // self.step + 1
        elif (self.mode == 'test'):
            return (self.test_y.shape[0] - self.win_size) // self.win_size + 1
        elif (self.mode == 'all_test' or self.mode == 'all'):
            return (self.all_test_y.shape[0] - self.win_size) // self.win_size + 1

    def __getitem__(self, index):
        index = index * self.step
        if self.mode == "train":
            return np.float32(self.train[index:index + self.win_size]), np.float32(
                self.train_labels[index:index + self.win_size])
        elif (self.mode == 'test'):
            return np.float32(self.test[
                                index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])
        elif (self.mode == 'all_test' or self.mode == 'all'):
            return np.float32(self.all_test_y[
                                index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32(
                self.all_test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size])
