diff --git a/code/TFC/augmentations.py b/code/TFC/augmentations.py
index 1ec7c28..f232ccbb 100644
--- a/code/TFC/augmentations.py
+++ b/code/TFC/augmentations.py
@@ -58,13 +58,15 @@ def DataTransform_FD(sample, config):
     return aug_F
 
 def remove_frequency(x, pertub_ratio=0.0):
-    mask = torch.cuda.FloatTensor(x.shape).uniform_() > pertub_ratio # maskout_ratio are False
+    # mask = torch.cuda.FloatTensor(x.shape).uniform_() > pertub_ratio # maskout_ratio are False
+    mask = torch.FloatTensor(x.shape).uniform_() > pertub_ratio # maskout_ratio are False
     mask = mask.to(x.device)
     return x*mask
 
 def add_frequency(x, pertub_ratio=0.0):
 
-    mask = torch.cuda.FloatTensor(x.shape).uniform_() > (1-pertub_ratio) # only pertub_ratio of all values are True
+    # mask = torch.cuda.FloatTensor(x.shape).uniform_() > (1-pertub_ratio) # only pertub_ratio of all values are True
+    mask = torch.FloatTensor(x.shape).uniform_() > (1-pertub_ratio) # only pertub_ratio of all values are True
     mask = mask.to(x.device)
     max_amplitude = x.max()
     random_am = torch.rand(mask.shape)*(max_amplitude*0.1)
diff --git a/code/TFC/dataloader.py b/code/TFC/dataloader.py
index e3c940c..9ccc959 100644
--- a/code/TFC/dataloader.py
+++ b/code/TFC/dataloader.py
@@ -1,124 +1,217 @@
-import torch
-from torch.utils.data import DataLoader
-from torch.utils.data import Dataset
-import os
-import numpy as np
-from augmentations import DataTransform_FD, DataTransform_TD
-import torch.fft as fft
-
-def generate_freq(dataset, config):
-    X_train = dataset["samples"]
-    y_train = dataset['labels']
-    # shuffle
-    data = list(zip(X_train, y_train))
-    np.random.shuffle(data)
-    data = data[:10000] # take a subset for testing.
-    X_train, y_train = zip(*data)
-    X_train, y_train = torch.stack(list(X_train), dim=0), torch.stack(list(y_train), dim=0)
-
-    if len(X_train.shape) < 3:
-        X_train = X_train.unsqueeze(2)
-
-    if X_train.shape.index(min(X_train.shape)) != 1:  # make sure the Channels in second dim
-        X_train = X_train.permute(0, 2, 1)
-
-    """Align the TS length between source and target datasets"""
-    X_train = X_train[:, :1, :int(config.TSlength_aligned)] # take the first 178 samples
-
-    if isinstance(X_train, np.ndarray):
-        x_data = torch.from_numpy(X_train)
-    else:
-        x_data = X_train
-
-    """Transfer x_data to Frequency Domain. If use fft.fft, the output has the same shape; if use fft.rfft, 
-    the output shape is half of the time window."""
-
-    x_data_f = fft.fft(x_data).abs() #/(window_length) # rfft for real value inputs.
-    return (X_train, y_train, x_data_f)
-
-class Load_Dataset(Dataset):
-    # Initialize your data, download, etc.
-    def __init__(self, dataset, config, training_mode, target_dataset_size=64, subset=False):
-        super(Load_Dataset, self).__init__()
-        self.training_mode = training_mode
-        X_train = dataset["samples"]
-        y_train = dataset["labels"]
-        # shuffle
-        data = list(zip(X_train, y_train))
-        np.random.shuffle(data)
-        X_train, y_train = zip(*data)
-        X_train, y_train = torch.stack(list(X_train), dim=0), torch.stack(list(y_train), dim=0)
-
-        if len(X_train.shape) < 3:
-            X_train = X_train.unsqueeze(2)
-
-        if X_train.shape.index(min(X_train.shape)) != 1:  # make sure the Channels in second dim
-            X_train = X_train.permute(0, 2, 1)
-
-        """Align the TS length between source and target datasets"""
-        X_train = X_train[:, :1, :int(config.TSlength_aligned)] # take the first 178 samples
-
-        """Subset for debugging"""
-        if subset == True:
-            subset_size = target_dataset_size * 10 #30 #7 # 60*1
-            """if the dimension is larger than 178, take the first 178 dimensions. If multiple channels, take the first channel"""
-            X_train = X_train[:subset_size]
-            y_train = y_train[:subset_size]
-            print('Using subset for debugging, the datasize is:', y_train.shape[0])
-
-        if isinstance(X_train, np.ndarray):
-            self.x_data = torch.from_numpy(X_train)
-            self.y_data = torch.from_numpy(y_train).long()
-        else:
-            self.x_data = X_train
-            self.y_data = y_train
-
-        """Transfer x_data to Frequency Domain. If use fft.fft, the output has the same shape; if use fft.rfft, 
-        the output shape is half of the time window."""
-
-        window_length = self.x_data.shape[-1]
-        self.x_data_f = fft.fft(self.x_data).abs() #/(window_length) # rfft for real value inputs.
-        self.len = X_train.shape[0]
-
-        """Augmentation"""
-        if training_mode == "pre_train":  # no need to apply Augmentations in other modes
-            self.aug1 = DataTransform_TD(self.x_data, config)
-            self.aug1_f = DataTransform_FD(self.x_data_f, config) # [7360, 1, 90]
-
-    def __getitem__(self, index):
-        if self.training_mode == "pre_train":
-            return self.x_data[index], self.y_data[index], self.aug1[index],  \
-                   self.x_data_f[index], self.aug1_f[index]
-        else:
-            return self.x_data[index], self.y_data[index], self.x_data[index], \
-                   self.x_data_f[index], self.x_data_f[index]
-
-    def __len__(self):
-        return self.len
-
-
-def data_generator(sourcedata_path, targetdata_path, configs, training_mode, subset=True):
-    train_dataset = torch.load(os.path.join(sourcedata_path, "train.pt"))
-    finetune_dataset = torch.load(os.path.join(targetdata_path, "train.pt"))  # train.pt
-    test_dataset = torch.load(os.path.join(targetdata_path, "test.pt"))  # test.pt
-    """In pre-training: 
-    train_dataset: [371055, 1, 178] from SleepEEG.    
-    finetune_dataset: [60, 1, 178], test_dataset: [11420, 1, 178] from Epilepsy"""
-
-    # subset = True # if true, use a subset for debugging.
-    train_dataset = Load_Dataset(train_dataset, configs, training_mode, target_dataset_size=configs.batch_size, subset=subset) # for self-supervised, the data are augmented here
-    finetune_dataset = Load_Dataset(finetune_dataset, configs, training_mode, target_dataset_size=configs.target_batch_size, subset=subset)
-    test_dataset = Load_Dataset(test_dataset, configs, training_mode,
-                                target_dataset_size=configs.target_batch_size, subset=False)
-
-    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=configs.batch_size,
-                                               shuffle=True, drop_last=configs.drop_last,
-                                               num_workers=0)
-    finetune_loader = torch.utils.data.DataLoader(dataset=finetune_dataset, batch_size=configs.target_batch_size,
-                                               shuffle=True, drop_last=configs.drop_last,
-                                               num_workers=0)
-    test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=configs.target_batch_size,
-                                              shuffle=True, drop_last=False,
-                                              num_workers=0)
-
-    return train_loader, finetune_loader, test_loader
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+import os
+import numpy as np
+import pandas as pd
+from augmentations import DataTransform_FD, DataTransform_TD
+import torch.fft as fft
+
+def generate_freq(dataset, config):
+    X_train = dataset["samples"]
+    y_train = dataset['labels']
+    # shuffle
+    data = list(zip(X_train, y_train))
+    np.random.shuffle(data)
+    data = data[:10000] # take a subset for testing.
+    X_train, y_train = zip(*data)
+    X_train, y_train = torch.stack(list(X_train), dim=0), torch.stack(list(y_train), dim=0)
+
+    if len(X_train.shape) < 3:
+        X_train = X_train.unsqueeze(2)
+
+    if X_train.shape.index(min(X_train.shape)) != 1:  # make sure the Channels in second dim
+        X_train = X_train.permute(0, 2, 1)
+
+    """Align the TS length between source and target datasets"""
+    X_train = X_train[:, :1, :int(config.TSlength_aligned)] # take the first 178 samples
+
+    if isinstance(X_train, np.ndarray):
+        x_data = torch.from_numpy(X_train)
+    else:
+        x_data = X_train
+
+    """Transfer x_data to Frequency Domain. If use fft.fft, the output has the same shape; if use fft.rfft, 
+    the output shape is half of the time window."""
+
+    x_data_f = fft.fft(x_data).abs() #/(window_length) # rfft for real value inputs.
+    return (X_train, y_train, x_data_f)
+
+class Load_Dataset(Dataset):
+    # Initialize your data, download, etc.
+    def __init__(self, dataset, config, training_mode, target_dataset_size=64, subset=False):
+        super(Load_Dataset, self).__init__()
+        self.training_mode = training_mode
+        X_train = dataset["samples"]
+        y_train = dataset["labels"]
+        # shuffle
+        data = list(zip(X_train, y_train))
+        np.random.shuffle(data)
+        X_train, y_train = zip(*data)
+        X_train, y_train = torch.stack(list(X_train), dim=0), torch.stack(list(y_train), dim=0)
+
+        if len(X_train.shape) < 3:
+            X_train = X_train.unsqueeze(2)
+
+        if X_train.shape.index(min(X_train.shape)) != 1:  # make sure the Channels in second dim
+            X_train = X_train.permute(0, 2, 1)
+
+        """Align the TS length between source and target datasets"""
+        X_train = X_train[:, :1, :int(config.TSlength_aligned)] # take the first 178 samples
+
+        """Subset for debugging"""
+        if subset == True:
+            subset_size = target_dataset_size * 10 #30 #7 # 60*1
+            """if the dimension is larger than 178, take the first 178 dimensions. If multiple channels, take the first channel"""
+            X_train = X_train[:subset_size]
+            y_train = y_train[:subset_size]
+            print('Using subset for debugging, the datasize is:', y_train.shape[0])
+
+        if isinstance(X_train, np.ndarray):
+            self.x_data = torch.from_numpy(X_train)
+            self.y_data = torch.from_numpy(y_train).long()
+        else:
+            self.x_data = X_train
+            self.y_data = y_train
+
+        """Transfer x_data to Frequency Domain. If use fft.fft, the output has the same shape; if use fft.rfft, 
+        the output shape is half of the time window."""
+
+        window_length = self.x_data.shape[-1]
+        self.x_data_f = fft.fft(self.x_data).abs() #/(window_length) # rfft for real value inputs.
+        self.len = X_train.shape[0]
+
+        """Augmentation"""
+        if training_mode == "pre_train":  # no need to apply Augmentations in other modes
+            self.aug1 = DataTransform_TD(self.x_data, config)
+            self.aug1_f = DataTransform_FD(self.x_data_f, config) # [7360, 1, 90]
+
+    def __getitem__(self, index):
+        if self.training_mode == "pre_train":
+            return self.x_data[index], self.y_data[index], self.aug1[index],  \
+                   self.x_data_f[index], self.aug1_f[index]
+        else:
+            return self.x_data[index], self.y_data[index], self.x_data[index], \
+                   self.x_data_f[index], self.x_data_f[index]
+
+    def __len__(self):
+        return self.len
+
+
+def data_generator(sourcedata_path, targetdata_path, configs, training_mode, subset=True):
+    train_dataset = torch.load(os.path.join(sourcedata_path, "train.pt"))
+    finetune_dataset = torch.load(os.path.join(targetdata_path, "train.pt"))  # train.pt
+    test_dataset = torch.load(os.path.join(targetdata_path, "test.pt"))  # test.pt
+    """In pre-training: 
+    train_dataset: [371055, 1, 178] from SleepEEG.    
+    finetune_dataset: [60, 1, 178], test_dataset: [11420, 1, 178] from Epilepsy"""
+
+    # subset = True # if true, use a subset for debugging.
+    train_dataset = Load_Dataset(train_dataset, configs, training_mode, target_dataset_size=configs.batch_size, subset=subset) # for self-supervised, the data are augmented here
+    finetune_dataset = Load_Dataset(finetune_dataset, configs, training_mode, target_dataset_size=configs.target_batch_size, subset=subset)
+    test_dataset = Load_Dataset(test_dataset, configs, training_mode,
+                                target_dataset_size=configs.target_batch_size, subset=False)
+
+    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=configs.batch_size,
+                                               shuffle=True, drop_last=configs.drop_last,
+                                               num_workers=0)
+    finetune_loader = torch.utils.data.DataLoader(dataset=finetune_dataset, batch_size=configs.target_batch_size,
+                                               shuffle=True, drop_last=configs.drop_last,
+                                               num_workers=0)
+    test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=configs.target_batch_size,
+                                              shuffle=True, drop_last=False,
+                                              num_workers=0)
+
+    return train_loader, finetune_loader, test_loader
+
+
+class Load_Dataset_motion(Dataset):
+    # Initialize your data, download, etc.
+    def __init__(self, dataset, config, training_mode, target_dataset_size=64):
+        super(Load_Dataset_motion, self).__init__()
+        self.training_mode = training_mode
+        ts_data = dataset['ts_data']
+        fft_data = dataset['fft_data']
+        # self.seq_len = dataset.shape[1] # (count, seq_len, fea_num)
+        X_train = torch.from_numpy(ts_data)
+        X_train_fft = torch.from_numpy(fft_data)
+
+        if len(X_train.shape) < 3:
+            X_train = X_train.unsqueeze(2)
+        if len(X_train_fft.shape) < 3:
+            X_train_fft = X_train_fft.unsqueeze(2)
+
+        if X_train.shape.index(min(X_train.shape)) != 1:  # make sure the Channels in second dim
+            X_train = X_train.permute(0, 2, 1)
+        if X_train_fft.shape.index(min(X_train_fft.shape)) != 1:  # make sure the Channels in second dim
+            X_train_fft = X_train_fft.permute(0, 2, 1)
+
+        self.len = X_train.shape[0]
+
+        """Augmentation"""
+        self.x_data = X_train
+        self.x_data_f = X_train_fft
+        if training_mode == "pre_train":  # no need to apply Augmentations in other modes
+            self.aug1 = DataTransform_TD(self.x_data, config)
+            self.aug1_f = DataTransform_FD(self.x_data_f, config) # [7360, 1, 90]
+
+    def __getitem__(self, index):
+        if self.training_mode == "pre_train":
+            return self.x_data[index], self.aug1[index], self.x_data_f[index], self.aug1_f[index]
+        else:
+            return self.x_data[index], self.y_data[index], self.x_data[index], \
+                   self.x_data_f[index], self.x_data_f[index]
+
+    def __len__(self):
+        return self.len
+
+
+def data_generator_motion(taindata_path_set, valdata_path_set, configs, training_mode, sample_count=-1, sample_count_val=-1):
+    train_data_dir = taindata_path_set['train_data_dir']
+    train_motion_names_dir = taindata_path_set['train_motion_names_dir']
+    train_fft_dir = taindata_path_set['train_fft_dir']
+    val_data_dir = valdata_path_set['val_data_dir']
+    val_motion_names_dir = valdata_path_set['val_motion_names_dir']
+    val_fft_dir = valdata_path_set['val_fft_dir']
+
+    print('loading train_ts')
+    train_ts = np.load(train_data_dir)
+    if sample_count > 0: 
+        train_ts = train_ts[:sample_count, :, :]
+    train_motion_names = pd.read_parquet(train_motion_names_dir)
+    if sample_count > 0: 
+        train_motion_names = train_motion_names.iloc[:sample_count]
+    print('loading train_fft')
+    train_fft = np.load(train_fft_dir)
+    if sample_count > 0: 
+        train_fft = train_fft[:sample_count, :, :]
+
+    print('loading val_ts')
+    val_ts = np.load(val_data_dir)
+    if sample_count_val > 0: 
+        val_ts = val_ts[:sample_count_val, :, :]
+    val_motion_names = pd.read_parquet(val_motion_names_dir)
+    if sample_count_val > 0:
+        val_motion_names = val_motion_names.iloc[:sample_count_val]
+    print('loading val_fft')
+    val_fft = np.load(val_fft_dir)
+    if sample_count_val > 0:
+        val_fft = val_fft[:sample_count_val, :, :]
+    
+    train_ts = np.nan_to_num(train_ts, neginf=0)
+    train_fft = np.nan_to_num(train_fft, neginf=0)
+    val_ts = np.nan_to_num(val_ts, neginf=0)
+    val_fft = np.nan_to_num(val_fft, neginf=0)
+
+    train_dataset = {'ts_data': train_ts, 'fft_data': train_fft}
+    val_dataset = {'ts_data': val_ts, 'fft_data': val_fft}
+
+    train_dataset = Load_Dataset_motion(train_dataset, configs, training_mode, target_dataset_size=configs.batch_size) # for self-supervised, the data are augmented here
+    val_dataset = Load_Dataset_motion(val_dataset, configs, training_mode, target_dataset_size=configs.batch_size)
+
+    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=configs.batch_size, 
+                                               shuffle=True, drop_last=configs.drop_last, num_workers=0)
+    val_loader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=configs.batch_size,
+                                             shuffle=True, drop_last=configs.drop_last, num_workers=0)
+
+    return train_loader, val_loader, train_motion_names, val_motion_names
diff --git a/code/TFC/loss.py b/code/TFC/loss.py
index cb8dcbc..b58e5f9 100644
--- a/code/TFC/loss.py
+++ b/code/TFC/loss.py
@@ -1,193 +1,193 @@
-import torch
-import torch.nn.functional as F
-import numpy as np
-
-class NTXentLoss(torch.nn.Module):
-    def __init__(self, device, batch_size, temperature, use_cosine_similarity):
-        super(NTXentLoss, self).__init__()
-        self.batch_size = batch_size
-        self.temperature = temperature
-        self.device = device
-        self.softmax = torch.nn.Softmax(dim=-1)
-        self.mask_samples_from_same_repr = self._get_correlated_mask().type(torch.bool)
-        self.similarity_function = self._get_similarity_function(use_cosine_similarity)
-        self.criterion = torch.nn.CrossEntropyLoss(reduction="sum")
-
-    def _get_similarity_function(self, use_cosine_similarity):
-        if use_cosine_similarity:
-            self._cosine_similarity = torch.nn.CosineSimilarity(dim=-1)
-            return self._cosine_simililarity
-        else:
-            return self._dot_simililarity
-
-    def _get_correlated_mask(self):
-        diag = np.eye(2 * self.batch_size)
-        l1 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=-self.batch_size)
-        l2 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=self.batch_size)
-        mask = torch.from_numpy((diag + l1 + l2))
-        mask = (1 - mask).type(torch.bool)
-        return mask.to(self.device)
-
-    @staticmethod
-    def _dot_simililarity(x, y):
-        v = torch.tensordot(x.unsqueeze(1), y.T.unsqueeze(0), dims=2)
-        # x shape: (N, 1, C)
-        # y shape: (1, C, 2N)
-        # v shape: (N, 2N)
-        return v
-
-    def _cosine_simililarity(self, x, y):
-        # x shape: (N, 1, C)
-        # y shape: (1, 2N, C)
-        # v shape: (N, 2N)
-        v = self._cosine_similarity(x.unsqueeze(1), y.unsqueeze(0))
-        return v
-
-    def forward(self, zis, zjs):
-        representations = torch.cat([zjs, zis], dim=0)
-
-        similarity_matrix = self.similarity_function(representations, representations)
-
-        # filter out the scores from the positive samples
-        l_pos = torch.diag(similarity_matrix, self.batch_size)
-        r_pos = torch.diag(similarity_matrix, -self.batch_size)
-        positives = torch.cat([l_pos, r_pos]).view(2 * self.batch_size, 1)
-
-        negatives = similarity_matrix[self.mask_samples_from_same_repr].view(2 * self.batch_size, -1)
-
-        logits = torch.cat((positives, negatives), dim=1)
-        logits /= self.temperature
-
-        """Criterion has an internal one-hot function. Here, make all positives as 1 while all negatives as 0. """
-        labels = torch.zeros(2 * self.batch_size).to(self.device).long()
-        loss = self.criterion(logits, labels)
-
-        return loss / (2 * self.batch_size)
-
-class NTXentLoss_poly(torch.nn.Module):
-
-    def __init__(self, device, batch_size, temperature, use_cosine_similarity):
-        super(NTXentLoss_poly, self).__init__()
-        self.batch_size = batch_size
-        self.temperature = temperature
-        self.device = device
-        self.softmax = torch.nn.Softmax(dim=-1)
-        self.mask_samples_from_same_repr = self._get_correlated_mask().type(torch.bool)
-        self.similarity_function = self._get_similarity_function(use_cosine_similarity)
-        self.criterion = torch.nn.CrossEntropyLoss(reduction="sum")
-
-    def _get_similarity_function(self, use_cosine_similarity):
-        if use_cosine_similarity:
-            self._cosine_similarity = torch.nn.CosineSimilarity(dim=-1)
-            return self._cosine_simililarity
-        else:
-            return self._dot_simililarity
-
-    def _get_correlated_mask(self):
-        diag = np.eye(2 * self.batch_size)
-        l1 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=-self.batch_size)
-        l2 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=self.batch_size)
-        mask = torch.from_numpy((diag + l1 + l2))
-        mask = (1 - mask).type(torch.bool)
-        return mask.to(self.device)
-
-    @staticmethod
-    def _dot_simililarity(x, y):
-        v = torch.tensordot(x.unsqueeze(1), y.T.unsqueeze(0), dims=2)
-        # x shape: (N, 1, C)
-        # y shape: (1, C, 2N)
-        # v shape: (N, 2N)
-        return v
-
-    def _cosine_simililarity(self, x, y):
-        # x shape: (N, 1, C)
-        # y shape: (1, 2N, C)
-        # v shape: (N, 2N)
-        v = self._cosine_similarity(x.unsqueeze(1), y.unsqueeze(0))
-        return v
-
-    def forward(self, zis, zjs):
-        representations = torch.cat([zjs, zis], dim=0)
-
-        similarity_matrix = self.similarity_function(representations, representations)
-
-        # filter out the scores from the positive samples
-        l_pos = torch.diag(similarity_matrix, self.batch_size)
-        r_pos = torch.diag(similarity_matrix, -self.batch_size)
-        positives = torch.cat([l_pos, r_pos]).view(2 * self.batch_size, 1)
-
-        negatives = similarity_matrix[self.mask_samples_from_same_repr].view(2 * self.batch_size, -1)
-
-        logits = torch.cat((positives, negatives), dim=1)
-        logits /= self.temperature
-
-        """Criterion has an internal one-hot function. Here, make all positives as 1 while all negatives as 0. """
-        labels = torch.zeros(2 * self.batch_size).to(self.device).long()
-        CE = self.criterion(logits, labels)
-
-        onehot_label = torch.cat((torch.ones(2 * self.batch_size, 1),torch.zeros(2 * self.batch_size, negatives.shape[-1])),dim=-1).to(self.device).long()
-        # Add poly loss
-        pt = torch.mean(onehot_label* torch.nn.functional.softmax(logits,dim=-1))
-
-        epsilon = self.batch_size
-        # loss = CE/ (2 * self.batch_size) + epsilon*(1-pt) # replace 1 by 1/self.batch_size
-        loss = CE / (2 * self.batch_size) + epsilon * (1/self.batch_size - pt)
-        # loss = CE / (2 * self.batch_size)
-
-        return loss
-
-
-
-
-class hierarchical_contrastive_loss(torch.nn.Module):
-
-    def __init__(self, device):
-        super(hierarchical_contrastive_loss, self).__init__()
-        self.device = device
-
-    def instance_contrastive_loss(self, z1, z2):
-        B, T = z1.size(0), z1.size(1)
-        if B == 1:
-            return z1.new_tensor(0.)
-        z = torch.cat([z1, z2], dim=0)  # 2B x T x C
-        z = z.transpose(0, 1)  # T x 2B x C
-        sim = torch.matmul(z, z.transpose(1, 2))  # T x 2B x 2B
-        logits = torch.tril(sim, diagonal=-1)[:, :, :-1]  # T x 2B x (2B-1)
-        logits += torch.triu(sim, diagonal=1)[:, :, 1:]
-        logits = -F.log_softmax(logits, dim=-1)
-
-        i = torch.arange(B)
-        loss = (logits[:, i, B + i - 1].mean() + logits[:, B + i, i].mean()) / 2
-        return loss
-
-
-    def temporal_contrastive_loss(self, z1, z2):
-        B, T = z1.size(0), z1.size(1)
-        if T == 1:
-            return z1.new_tensor(0.)
-        z = torch.cat([z1, z2], dim=1)  # B x 2T x C
-        sim = torch.matmul(z, z.transpose(1, 2))  # B x 2T x 2T
-        logits = torch.tril(sim, diagonal=-1)[:, :, :-1]  # B x 2T x (2T-1)
-        logits += torch.triu(sim, diagonal=1)[:, :, 1:]
-        logits = -F.log_softmax(logits, dim=-1)
-
-        t = torch.arange(T)
-        loss = (logits[:, t, T + t - 1].mean() + logits[:, T + t, t].mean()) / 2
-        return loss
-    def forward(self, z1, z2, alpha=0.5, temporal_unit=0):
-        loss = torch.tensor(0., device=self.device) #, device=z1.device
-        d = 0
-        while z1.size(1) > 1:
-            if alpha != 0:
-                loss += alpha * self.instance_contrastive_loss(z1, z2)
-            if d >= temporal_unit:
-                if 1 - alpha != 0:
-                    loss += (1 - alpha) * self.temporal_contrastive_loss(z1, z2)
-            d += 1
-            z1 = F.max_pool1d(z1.transpose(1, 2), kernel_size=2).transpose(1, 2)
-            z2 = F.max_pool1d(z2.transpose(1, 2), kernel_size=2).transpose(1, 2)
-        if z1.size(1) == 1:
-            if alpha != 0:
-                loss += alpha * self.instance_contrastive_loss(z1, z2)
-            d += 1
-        return loss / d
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+class NTXentLoss(torch.nn.Module):
+    def __init__(self, device, batch_size, temperature, use_cosine_similarity):
+        super(NTXentLoss, self).__init__()
+        self.batch_size = batch_size
+        self.temperature = temperature
+        self.device = device
+        self.softmax = torch.nn.Softmax(dim=-1)
+        self.mask_samples_from_same_repr = self._get_correlated_mask().type(torch.bool)
+        self.similarity_function = self._get_similarity_function(use_cosine_similarity)
+        self.criterion = torch.nn.CrossEntropyLoss(reduction="sum")
+
+    def _get_similarity_function(self, use_cosine_similarity):
+        if use_cosine_similarity:
+            self._cosine_similarity = torch.nn.CosineSimilarity(dim=-1)
+            return self._cosine_simililarity
+        else:
+            return self._dot_simililarity
+
+    def _get_correlated_mask(self):
+        diag = np.eye(2 * self.batch_size)
+        l1 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=-self.batch_size)
+        l2 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=self.batch_size)
+        mask = torch.from_numpy((diag + l1 + l2))
+        mask = (1 - mask).type(torch.bool)
+        return mask.to(self.device)
+
+    @staticmethod
+    def _dot_simililarity(x, y):
+        v = torch.tensordot(x.unsqueeze(1), y.T.unsqueeze(0), dims=2)
+        # x shape: (N, 1, C)
+        # y shape: (1, C, 2N)
+        # v shape: (N, 2N)
+        return v
+
+    def _cosine_simililarity(self, x, y):
+        # x shape: (N, 1, C)
+        # y shape: (1, 2N, C)
+        # v shape: (N, 2N)
+        v = self._cosine_similarity(x.unsqueeze(1), y.unsqueeze(0))
+        return v
+
+    def forward(self, zis, zjs):
+        representations = torch.cat([zjs, zis], dim=0)
+
+        similarity_matrix = self.similarity_function(representations, representations)
+
+        # filter out the scores from the positive samples
+        l_pos = torch.diag(similarity_matrix, self.batch_size)
+        r_pos = torch.diag(similarity_matrix, -self.batch_size)
+        positives = torch.cat([l_pos, r_pos]).view(2 * self.batch_size, 1)
+
+        negatives = similarity_matrix[self.mask_samples_from_same_repr].view(2 * self.batch_size, -1)
+
+        logits = torch.cat((positives, negatives), dim=1)
+        logits /= self.temperature
+
+        """Criterion has an internal one-hot function. Here, make all positives as 1 while all negatives as 0. """
+        labels = torch.zeros(2 * self.batch_size).to(self.device).long()
+        loss = self.criterion(logits, labels)
+
+        return loss / (2 * self.batch_size)
+
+class NTXentLoss_poly(torch.nn.Module):
+
+    def __init__(self, device, batch_size, temperature, use_cosine_similarity):
+        super(NTXentLoss_poly, self).__init__()
+        self.batch_size = batch_size
+        self.temperature = temperature
+        self.device = device
+        self.softmax = torch.nn.Softmax(dim=-1)
+        self.mask_samples_from_same_repr = self._get_correlated_mask().type(torch.bool)
+        self.similarity_function = self._get_similarity_function(use_cosine_similarity)
+        self.criterion = torch.nn.CrossEntropyLoss(reduction="sum")
+
+    def _get_similarity_function(self, use_cosine_similarity):
+        if use_cosine_similarity:
+            self._cosine_similarity = torch.nn.CosineSimilarity(dim=-1)
+            return self._cosine_simililarity
+        else:
+            return self._dot_simililarity
+
+    def _get_correlated_mask(self):
+        diag = np.eye(2 * self.batch_size)
+        l1 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=-self.batch_size)
+        l2 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=self.batch_size)
+        mask = torch.from_numpy((diag + l1 + l2))
+        mask = (1 - mask).type(torch.bool)
+        return mask.to(self.device)
+
+    @staticmethod
+    def _dot_simililarity(x, y):
+        v = torch.tensordot(x.unsqueeze(1), y.T.unsqueeze(0), dims=2)
+        # x shape: (N, 1, C)
+        # y shape: (1, C, 2N)
+        # v shape: (N, 2N)
+        return v
+
+    def _cosine_simililarity(self, x, y):
+        # x shape: (N, 1, C)
+        # y shape: (1, 2N, C)
+        # v shape: (N, 2N)
+        v = self._cosine_similarity(x.unsqueeze(1), y.unsqueeze(0))
+        return v
+
+    def forward(self, zis, zjs):
+        representations = torch.cat([zjs, zis], dim=0)
+
+        similarity_matrix = self.similarity_function(representations, representations)
+
+        # filter out the scores from the positive samples
+        l_pos = torch.diag(similarity_matrix, self.batch_size)
+        r_pos = torch.diag(similarity_matrix, -self.batch_size)
+        positives = torch.cat([l_pos, r_pos]).view(2 * self.batch_size, 1)
+
+        negatives = similarity_matrix[self.mask_samples_from_same_repr].view(2 * self.batch_size, -1)
+
+        logits = torch.cat((positives, negatives), dim=1)
+        logits /= self.temperature
+
+        """Criterion has an internal one-hot function. Here, make all positives as 1 while all negatives as 0. """
+        labels = torch.zeros(2 * self.batch_size).to(self.device).long()
+        CE = self.criterion(logits, labels)
+
+        onehot_label = torch.cat((torch.ones(2 * self.batch_size, 1),torch.zeros(2 * self.batch_size, negatives.shape[-1])),dim=-1).to(self.device).long()
+        # Add poly loss
+        pt = torch.mean(onehot_label* torch.nn.functional.softmax(logits,dim=-1))
+
+        epsilon = self.batch_size
+        # loss = CE/ (2 * self.batch_size) + epsilon*(1-pt) # replace 1 by 1/self.batch_size
+        loss = CE / (2 * self.batch_size) + epsilon * (1/self.batch_size - pt)
+        # loss = CE / (2 * self.batch_size)
+
+        return loss
+
+
+
+
+class hierarchical_contrastive_loss(torch.nn.Module):
+
+    def __init__(self, device):
+        super(hierarchical_contrastive_loss, self).__init__()
+        self.device = device
+
+    def instance_contrastive_loss(self, z1, z2):
+        B, T = z1.size(0), z1.size(1)
+        if B == 1:
+            return z1.new_tensor(0.)
+        z = torch.cat([z1, z2], dim=0)  # 2B x T x C
+        z = z.transpose(0, 1)  # T x 2B x C
+        sim = torch.matmul(z, z.transpose(1, 2))  # T x 2B x 2B
+        logits = torch.tril(sim, diagonal=-1)[:, :, :-1]  # T x 2B x (2B-1)
+        logits += torch.triu(sim, diagonal=1)[:, :, 1:]
+        logits = -F.log_softmax(logits, dim=-1)
+
+        i = torch.arange(B)
+        loss = (logits[:, i, B + i - 1].mean() + logits[:, B + i, i].mean()) / 2
+        return loss
+
+
+    def temporal_contrastive_loss(self, z1, z2):
+        B, T = z1.size(0), z1.size(1)
+        if T == 1:
+            return z1.new_tensor(0.)
+        z = torch.cat([z1, z2], dim=1)  # B x 2T x C
+        sim = torch.matmul(z, z.transpose(1, 2))  # B x 2T x 2T
+        logits = torch.tril(sim, diagonal=-1)[:, :, :-1]  # B x 2T x (2T-1)
+        logits += torch.triu(sim, diagonal=1)[:, :, 1:]
+        logits = -F.log_softmax(logits, dim=-1)
+
+        t = torch.arange(T)
+        loss = (logits[:, t, T + t - 1].mean() + logits[:, T + t, t].mean()) / 2
+        return loss
+    def forward(self, z1, z2, alpha=0.5, temporal_unit=0):
+        loss = torch.tensor(0., device=self.device) #, device=z1.device
+        d = 0
+        while z1.size(1) > 1:
+            if alpha != 0:
+                loss += alpha * self.instance_contrastive_loss(z1, z2)
+            if d >= temporal_unit:
+                if 1 - alpha != 0:
+                    loss += (1 - alpha) * self.temporal_contrastive_loss(z1, z2)
+            d += 1
+            z1 = F.max_pool1d(z1.transpose(1, 2), kernel_size=2).transpose(1, 2)
+            z2 = F.max_pool1d(z2.transpose(1, 2), kernel_size=2).transpose(1, 2)
+        if z1.size(1) == 1:
+            if alpha != 0:
+                loss += alpha * self.instance_contrastive_loss(z1, z2)
+            d += 1
+        return loss / d
diff --git a/code/TFC/model.py b/code/TFC/model.py
index 0482852..e5de8c1 100644
--- a/code/TFC/model.py
+++ b/code/TFC/model.py
@@ -1,60 +1,60 @@
-from torch import nn
-import torch
-from torch.nn import TransformerEncoder, TransformerEncoderLayer
-
-"""Two contrastive encoders"""
-class TFC(nn.Module):
-    def __init__(self, configs):
-        super(TFC, self).__init__()
-
-        encoder_layers_t = TransformerEncoderLayer(configs.TSlength_aligned, dim_feedforward=2*configs.TSlength_aligned, nhead=2, )
-        self.transformer_encoder_t = TransformerEncoder(encoder_layers_t, 2)
-
-        self.projector_t = nn.Sequential(
-            nn.Linear(configs.TSlength_aligned, 256),
-            nn.BatchNorm1d(256),
-            nn.ReLU(),
-            nn.Linear(256, 128)
-        )
-
-        encoder_layers_f = TransformerEncoderLayer(configs.TSlength_aligned, dim_feedforward=2*configs.TSlength_aligned,nhead=2,)
-        self.transformer_encoder_f = TransformerEncoder(encoder_layers_f, 2)
-
-        self.projector_f = nn.Sequential(
-            nn.Linear(configs.TSlength_aligned, 256),
-            nn.BatchNorm1d(256),
-            nn.ReLU(),
-            nn.Linear(256, 128)
-        )
-
-
-    def forward(self, x_in_t, x_in_f):
-        """Use Transformer"""
-        x = self.transformer_encoder_t(x_in_t)
-        h_time = x.reshape(x.shape[0], -1)
-
-        """Cross-space projector"""
-        z_time = self.projector_t(h_time)
-
-        """Frequency-based contrastive encoder"""
-        f = self.transformer_encoder_f(x_in_f)
-        h_freq = f.reshape(f.shape[0], -1)
-
-        """Cross-space projector"""
-        z_freq = self.projector_f(h_freq)
-
-        return h_time, z_time, h_freq, z_freq
-
-
-"""Downstream classifier only used in finetuning"""
-class target_classifier(nn.Module):
-    def __init__(self, configs):
-        super(target_classifier, self).__init__()
-        self.logits = nn.Linear(2*128, 64)
-        self.logits_simple = nn.Linear(64, configs.num_classes_target)
-
-    def forward(self, emb):
-        emb_flat = emb.reshape(emb.shape[0], -1)
-        emb = torch.sigmoid(self.logits(emb_flat))
-        pred = self.logits_simple(emb)
-        return pred
+from torch import nn
+import torch
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+
+"""Two contrastive encoders"""
+class TFC(nn.Module):
+    def __init__(self, configs):
+        super(TFC, self).__init__()
+
+        encoder_layers_t = TransformerEncoderLayer(configs.TSlength_aligned, dim_feedforward=2*configs.TSlength_aligned, nhead=configs.transformer_nhead, )
+        self.transformer_encoder_t = TransformerEncoder(encoder_layers_t, configs.transformer_num_layers)
+
+        self.projector_t = nn.Sequential(
+            nn.Linear(configs.TSlength_aligned*configs.input_channels, configs.embedding_len*2),
+            nn.BatchNorm1d(configs.embedding_len*2),
+            nn.ReLU(),
+            nn.Linear(configs.embedding_len*2, configs.embedding_len)
+        )
+
+        encoder_layers_f = TransformerEncoderLayer(configs.TSlength_aligned, dim_feedforward=2*configs.TSlength_aligned,nhead=configs.transformer_nhead,)
+        self.transformer_encoder_f = TransformerEncoder(encoder_layers_f, configs.transformer_num_layers)
+
+        self.projector_f = nn.Sequential(
+            nn.Linear(configs.TSlength_aligned*configs.input_channels, configs.embedding_len*2),
+            nn.BatchNorm1d(configs.embedding_len*2),
+            nn.ReLU(),
+            nn.Linear(configs.embedding_len*2, configs.embedding_len)
+        )
+
+
+    def forward(self, x_in_t, x_in_f):
+        """Use Transformer"""
+        x = self.transformer_encoder_t(x_in_t)
+        h_time = x.reshape(x.shape[0], -1)
+
+        """Cross-space projector"""
+        z_time = self.projector_t(h_time)
+
+        """Frequency-based contrastive encoder"""
+        f = self.transformer_encoder_f(x_in_f)
+        h_freq = f.reshape(f.shape[0], -1)
+
+        """Cross-space projector"""
+        z_freq = self.projector_f(h_freq)
+
+        return h_time, z_time, h_freq, z_freq
+
+
+"""Downstream classifier only used in finetuning"""
+class target_classifier(nn.Module):
+    def __init__(self, configs):
+        super(target_classifier, self).__init__()
+        self.logits = nn.Linear(2*128, 64)
+        self.logits_simple = nn.Linear(64, configs.num_classes_target)
+
+    def forward(self, emb):
+        emb_flat = emb.reshape(emb.shape[0], -1)
+        emb = torch.sigmoid(self.logits(emb_flat))
+        pred = self.logits_simple(emb)
+        return pred
diff --git a/code/TFC/trainer.py b/code/TFC/trainer.py
index 04844c0..8a187f8 100644
--- a/code/TFC/trainer.py
+++ b/code/TFC/trainer.py
@@ -14,94 +14,36 @@ def one_hot_encoding(X):
     b = np.eye(n_values)[X]
     return b
 
-def Trainer(model,  model_optimizer, classifier, classifier_optimizer, train_dl, valid_dl, test_dl, device,
-            logger, config, experiment_log_dir, training_mode):
-    # Start training
-    logger.debug("Training started ....")
+def Trainer(model,  model_optimizer, train_dl, valid_dl, device, config, training_mode, save_model_dir, save_model_or_checkpoints):
+    print("Training started ....\n")
 
     criterion = nn.CrossEntropyLoss()
     scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(model_optimizer, 'min')
+    train_loss_set = []
+    val_loss_set = []
     if training_mode == 'pre_train':
-        print('Pretraining on source dataset')
         for epoch in range(1, config.num_epoch + 1):
-            # Train and validate
-            """Train. In fine-tuning, this part is also trained???"""
             train_loss = model_pretrain(model, model_optimizer, criterion, train_dl, config, device, training_mode)
-            logger.debug(f'\nPre-training Epoch : {epoch}', f'Train Loss : {train_loss:.4f}')
-
-        os.makedirs(os.path.join(experiment_log_dir, "saved_models"), exist_ok=True)
-        chkpoint = {'model_state_dict': model.state_dict()}
-        torch.save(chkpoint, os.path.join(experiment_log_dir, "saved_models", f'ckp_last.pt'))
-        print('Pretrained model is stored at folder:{}'.format(experiment_log_dir+'saved_models'+'ckp_last.pt'))
-
-    """Fine-tuning and Test"""
-    if training_mode != 'pre_train':
-        """fine-tune"""
-        print('Fine-tune on Fine-tuning set')
-        performance_list = []
-        total_f1 = []
-        KNN_f1 = []
-        global emb_finetune, label_finetune, emb_test, label_test
-
-        for epoch in range(1, config.num_epoch + 1):
-            logger.debug(f'\nEpoch : {epoch}')
-
-            valid_loss, emb_finetune, label_finetune, F1 = model_finetune(model, model_optimizer, valid_dl, config,
-                                  device, training_mode, classifier=classifier, classifier_optimizer=classifier_optimizer)
-            scheduler.step(valid_loss)
-
-
-            # save best fine-tuning model""
-            global arch
-            arch = 'sleepedf2eplipsy'
-            if len(total_f1) == 0 or F1 > max(total_f1):
-                print('update fine-tuned model')
-                os.makedirs('experiments_logs/finetunemodel/', exist_ok=True)
-                torch.save(model.state_dict(), 'experiments_logs/finetunemodel/' + arch + '_model.pt')
-                torch.save(classifier.state_dict(), 'experiments_logs/finetunemodel/' + arch + '_classifier.pt')
-            total_f1.append(F1)
-
-            # evaluate on the test set
-            """Testing set"""
-            logger.debug('Test on Target datasts test set')
-            model.load_state_dict(torch.load('experiments_logs/finetunemodel/' + arch + '_model.pt'))
-            classifier.load_state_dict(torch.load('experiments_logs/finetunemodel/' + arch + '_classifier.pt'))
-            test_loss, test_acc, test_auc, test_prc, emb_test, label_test, performance = model_test(model, test_dl, config, device, training_mode,
-                                                             classifier=classifier, classifier_optimizer=classifier_optimizer)
-            performance_list.append(performance)
-
-            """Use KNN as another classifier; it's an alternation of the MLP classifier in function model_test. 
-            Experiments show KNN and MLP may work differently in different settings, so here we provide both. """
-            # train classifier: KNN
-            neigh = KNeighborsClassifier(n_neighbors=5)
-            neigh.fit(emb_finetune, label_finetune)
-            knn_acc_train = neigh.score(emb_finetune, label_finetune)
-            # print('KNN finetune acc:', knn_acc_train)
-            representation_test = emb_test.detach().cpu().numpy()
-
-            knn_result = neigh.predict(representation_test)
-            knn_result_score = neigh.predict_proba(representation_test)
-            one_hot_label_test = one_hot_encoding(label_test)
-            # print(classification_report(label_test, knn_result, digits=4))
-            # print(confusion_matrix(label_test, knn_result))
-            knn_acc = accuracy_score(label_test, knn_result)
-            precision = precision_score(label_test, knn_result, average='macro', )
-            recall = recall_score(label_test, knn_result, average='macro', )
-            F1 = f1_score(label_test, knn_result, average='macro')
-            auc = roc_auc_score(one_hot_label_test, knn_result_score, average="macro", multi_class="ovr")
-            prc = average_precision_score(one_hot_label_test, knn_result_score, average="macro")
-            print('KNN Testing: Acc=%.4f| Precision = %.4f | Recall = %.4f | F1 = %.4f | AUROC= %.4f | AUPRC=%.4f'%
-                  (knn_acc, precision, recall, F1, auc, prc))
-            KNN_f1.append(F1)
-        logger.debug("\n################## Best testing performance! #########################")
-        performance_array = np.array(performance_list)
-        best_performance = performance_array[np.argmax(performance_array[:,0], axis=0)]
-        print('Best Testing Performance: Acc=%.4f| Precision = %.4f | Recall = %.4f | F1 = %.4f | AUROC= %.4f '
-              '| AUPRC=%.4f' % (best_performance[0], best_performance[1], best_performance[2], best_performance[3],
-                                best_performance[4], best_performance[5]))
-        print('Best KNN F1', max(KNN_f1))
-
-    logger.debug("\n################## Training is Done! #########################")
+            train_loss_set.append(train_loss)
+            val_loss = model_pretrain_val(model, model_optimizer, criterion, valid_dl, config, device, training_mode)
+            val_loss_set.append(val_loss)
+            print(f'Pre-training Epoch : {epoch}', f'Train Loss : {train_loss:.4f}')
+            print(f'Pre-training Epoch : {epoch}', f'Val Loss : {val_loss:.4f}\n')
+
+            if save_model_or_checkpoints.lower() == 'checkpoints': 
+                # chkpoint = {'model_state_dict': model.state_dict()}
+                torch.save(model.state_dict(), f"{save_model_dir}/model_{epoch}.pt")
+
+        if save_model_or_checkpoints.lower() == 'model': 
+            # chkpoint = {'model_state_dict': model.state_dict()}
+            torch.save(model.state_dict(), f"{save_model_dir}/model_final_epoch.pt")
+        # os.makedirs(os.path.join(experiment_log_dir, "saved_models"), exist_ok=True)
+        # chkpoint = {'model_state_dict': model.state_dict()}
+        # torch.save(chkpoint, os.path.join(experiment_log_dir, "saved_models", f'ckp_last.pt'))
+        # print('Pretrained model is stored at folder:{}'.format(experiment_log_dir+'saved_models'+'ckp_last.pt'))
+    
+    print("Training End")
+    return train_loss_set, val_loss_set
 
 def model_pretrain(model, model_optimizer, criterion, train_loader, config, device, training_mode,):
     total_loss = []
@@ -111,8 +53,8 @@ def model_pretrain(model, model_optimizer, criterion, train_loader, config, devi
     # optimizer
     model_optimizer.zero_grad()
 
-    for batch_idx, (data, labels, aug1, data_f, aug1_f) in enumerate(train_loader):
-        data, labels = data.float().to(device), labels.long().to(device) # data: [128, 1, 178], labels: [128]
+    for batch_idx, (data, aug1, data_f, aug1_f) in enumerate(train_loader):
+        data = data.float().to(device)
         aug1 = aug1.float().to(device)  # aug1 = aug2 : [128, 1, 178]
         data_f, aug1_f = data_f.float().to(device), aug1_f.float().to(device)  # aug1 = aug2 : [128, 1, 178]
 
@@ -123,7 +65,7 @@ def model_pretrain(model, model_optimizer, criterion, train_loader, config, devi
         """Compute Pre-train loss"""
         """NTXentLoss: normalized temperature-scaled cross entropy loss. From SimCLR"""
         nt_xent_criterion = NTXentLoss_poly(device, config.batch_size, config.Context_Cont.temperature,
-                                       config.Context_Cont.use_cosine_similarity) # device, 128, 0.2, True
+                                            config.Context_Cont.use_cosine_similarity) # device, 128, 0.2, True
 
         loss_t = nt_xent_criterion(h_t, h_t_aug)
         loss_f = nt_xent_criterion(h_f, h_f_aug)
@@ -139,171 +81,46 @@ def model_pretrain(model, model_optimizer, criterion, train_loader, config, devi
         loss.backward()
         model_optimizer.step()
 
-    print('Pretraining: overall loss:{}, l_t: {}, l_f:{}, l_c:{}'.format(loss, loss_t, loss_f, l_TF))
+    # print('Pretraining: train loss:{:.3f}, l_t: {:.3f}, l_f:{:.3f}, l_c:{:.3f}'.format(loss, loss_t, loss_f, l_TF))
 
     ave_loss = torch.tensor(total_loss).mean()
 
     return ave_loss
 
 
-def model_finetune(model, model_optimizer, val_dl, config, device, training_mode, classifier=None, classifier_optimizer=None):
-    global labels, pred_numpy, fea_concat_flat
-    model.train()
-    classifier.train()
-
+def model_pretrain_val(model, model_optimizer, criterion, val_loader, config, device, training_mode,):
     total_loss = []
-    total_acc = []
-    total_auc = []  # it should be outside of the loop
-    total_prc = []
-
-    criterion = nn.CrossEntropyLoss()
-    outs = np.array([])
-    trgs = np.array([])
-    feas = np.array([])
-
-    for data, labels, aug1, data_f, aug1_f in val_dl:
-        # print('Fine-tuning: {} of target samples'.format(labels.shape[0]))
-        data, labels = data.float().to(device), labels.long().to(device)
-        data_f = data_f.float().to(device)
-        aug1 = aug1.float().to(device)
-        aug1_f = aug1_f.float().to(device)
-
-        """if random initialization:"""
-        model_optimizer.zero_grad()  # The gradients are zero, but the parameters are still randomly initialized.
-        classifier_optimizer.zero_grad()  # the classifier is newly added and randomly initialized
-
-        """Produce embeddings"""
-        h_t, z_t, h_f, z_f = model(data, data_f)
-        h_t_aug, z_t_aug, h_f_aug, z_f_aug = model(aug1, aug1_f)
-        nt_xent_criterion = NTXentLoss_poly(device, config.target_batch_size, config.Context_Cont.temperature,
-                                            config.Context_Cont.use_cosine_similarity)
-        loss_t = nt_xent_criterion(h_t, h_t_aug)
-        loss_f = nt_xent_criterion(h_f, h_f_aug)
-        l_TF = nt_xent_criterion(z_t, z_f)
-
-        l_1, l_2, l_3 = nt_xent_criterion(z_t, z_f_aug), nt_xent_criterion(z_t_aug, z_f), \
-                        nt_xent_criterion(z_t_aug, z_f_aug)
-        loss_c = (1 + l_TF - l_1) + (1 + l_TF - l_2) + (1 + l_TF - l_3) #
-
-
-        """Add supervised classifier: 1) it's unique to finetuning. 2) this classifier will also be used in test."""
-        fea_concat = torch.cat((z_t, z_f), dim=1)
-        predictions = classifier(fea_concat)
-        fea_concat_flat = fea_concat.reshape(fea_concat.shape[0], -1)
-        loss_p = criterion(predictions, labels)
-
-        lam = 0.1
-        loss = loss_p + l_TF + lam*(loss_t + loss_f)
-
-        acc_bs = labels.eq(predictions.detach().argmax(dim=1)).float().mean()
-        onehot_label = F.one_hot(labels)
-        pred_numpy = predictions.detach().cpu().numpy()
-
-        try:
-            auc_bs = roc_auc_score(onehot_label.detach().cpu().numpy(), pred_numpy, average="macro", multi_class="ovr" )
-        except:
-            auc_bs = np.float(0)
-        prc_bs = average_precision_score(onehot_label.detach().cpu().numpy(), pred_numpy)
-
-        total_acc.append(acc_bs)
-        total_auc.append(auc_bs)
-        total_prc.append(prc_bs)
-        total_loss.append(loss.item())
-        loss.backward()
-        model_optimizer.step()
-        classifier_optimizer.step()
-
-        if training_mode != "pre_train":
-            pred = predictions.max(1, keepdim=True)[1]  # get the index of the max log-probability
-            outs = np.append(outs, pred.cpu().numpy())
-            trgs = np.append(trgs, labels.data.cpu().numpy())
-            feas = np.append(feas, fea_concat_flat.data.cpu().numpy())
-
-    feas = feas.reshape([len(trgs), -1])  # produce the learned embeddings
-
-    labels_numpy = labels.detach().cpu().numpy()
-    pred_numpy = np.argmax(pred_numpy, axis=1)
-    precision = precision_score(labels_numpy, pred_numpy, average='macro', )
-    recall = recall_score(labels_numpy, pred_numpy, average='macro', )
-    F1 = f1_score(labels_numpy, pred_numpy, average='macro', )
-    ave_loss = torch.tensor(total_loss).mean()
-    ave_acc = torch.tensor(total_acc).mean()
-    ave_auc = torch.tensor(total_auc).mean()
-    ave_prc = torch.tensor(total_prc).mean()
-
-    print(' Finetune: loss = %.4f| Acc=%.4f | Precision = %.4f | Recall = %.4f | F1 = %.4f| AUROC=%.4f | AUPRC = %.4f'
-          % (ave_loss, ave_acc*100, precision * 100, recall * 100, F1 * 100, ave_auc * 100, ave_prc *100))
-
-    return ave_loss, feas, trgs, F1
-
-
-def model_test(model,  test_dl, config,  device, training_mode, classifier=None, classifier_optimizer=None):
     model.eval()
-    classifier.eval()
-
-    total_loss = []
-    total_acc = []
-    total_auc = []
-    total_prc = []
-
-    criterion = nn.CrossEntropyLoss() # the loss for downstream classifier
-    outs = np.array([])
-    trgs = np.array([])
-    emb_test_all = []
+    global val_loss, val_loss_t, val_loss_f, val_l_TF, val_loss_c, data_test, data_f_test
 
+    # no model update 
     with torch.no_grad():
-        labels_numpy_all, pred_numpy_all = np.zeros(1), np.zeros(1)
-        for data, labels, _,data_f, _ in test_dl:
-            data, labels = data.float().to(device), labels.long().to(device)
-            data_f = data_f.float().to(device)
+        for batch_idx, (data, aug1, data_f, aug1_f) in enumerate(val_loader):
+            data, aug1 = data.float().to(device), aug1.float().to(device) # [batch, fea_count, ts_count]
+            data_f, aug1_f = data_f.float().to(device), aug1_f.float().to(device)  # aug1 = aug2 : [128, 1, 178]
 
-            """Add supervised classifier: 1) it's unique to finetuning. 2) this classifier will also be used in test"""
+            """Produce embeddings"""
             h_t, z_t, h_f, z_f = model(data, data_f)
-            fea_concat = torch.cat((z_t, z_f), dim=1)
-            predictions_test = classifier(fea_concat)
-            fea_concat_flat = fea_concat.reshape(fea_concat.shape[0], -1)
-            emb_test_all.append(fea_concat_flat)
+            h_t_aug, z_t_aug, h_f_aug, z_f_aug = model(aug1, aug1_f)
 
-            loss = criterion(predictions_test, labels)
-            acc_bs = labels.eq(predictions_test.detach().argmax(dim=1)).float().mean()
-            onehot_label = F.one_hot(labels)
-            pred_numpy = predictions_test.detach().cpu().numpy()
-            labels_numpy = labels.detach().cpu().numpy()
-            try:
-                auc_bs = roc_auc_score(onehot_label.detach().cpu().numpy(), pred_numpy,
-                                   average="macro", multi_class="ovr")
-            except:
-                auc_bs = np.float(0)
-            prc_bs = average_precision_score(onehot_label.detach().cpu().numpy(), pred_numpy, average="macro")
-            pred_numpy = np.argmax(pred_numpy, axis=1)
+            """Compute Pre-train loss"""
+            """NTXentLoss: normalized temperature-scaled cross entropy loss. From SimCLR"""
+            nt_xent_criterion = NTXentLoss_poly(device, config.batch_size, config.Context_Cont.temperature,
+                                        config.Context_Cont.use_cosine_similarity) # device, 128, 0.2, True
 
-            total_acc.append(acc_bs)
-            total_auc.append(auc_bs)
-            total_prc.append(prc_bs)
+            val_loss_t = nt_xent_criterion(h_t, h_t_aug)
+            val_loss_f = nt_xent_criterion(h_f, h_f_aug)
+            val_l_TF = nt_xent_criterion(z_t, z_f) # this is the initial version of TF loss
 
-            total_loss.append(loss.item())
-            pred = predictions_test.max(1, keepdim=True)[1]  # get the index of the max log-probability
-            outs = np.append(outs, pred.cpu().numpy())
-            trgs = np.append(trgs, labels.data.cpu().numpy())
-            labels_numpy_all = np.concatenate((labels_numpy_all, labels_numpy))
-            pred_numpy_all = np.concatenate((pred_numpy_all, pred_numpy))
-    labels_numpy_all = labels_numpy_all[1:]
-    pred_numpy_all = pred_numpy_all[1:]
+            l_1, l_2, l_3 = nt_xent_criterion(z_t, z_f_aug), nt_xent_criterion(z_t_aug, z_f), nt_xent_criterion(z_t_aug, z_f_aug)
+            val_loss_c = (1 + val_l_TF - l_1) + (1 + val_l_TF - l_2) + (1 + val_l_TF - l_3)
 
-    # print('Test classification report', classification_report(labels_numpy_all, pred_numpy_all))
-    # print(confusion_matrix(labels_numpy_all, pred_numpy_all))
-    precision = precision_score(labels_numpy_all, pred_numpy_all, average='macro', )
-    recall = recall_score(labels_numpy_all, pred_numpy_all, average='macro', )
-    F1 = f1_score(labels_numpy_all, pred_numpy_all, average='macro', )
-    acc = accuracy_score(labels_numpy_all, pred_numpy_all, )
+            lam = 0.2
+            val_loss = lam*(val_loss_t + val_loss_f) + val_l_TF
+            total_loss.append(val_loss.item())
 
-    total_loss = torch.tensor(total_loss).mean()
-    total_acc = torch.tensor(total_acc).mean()
-    total_auc = torch.tensor(total_auc).mean()
-    total_prc = torch.tensor(total_prc).mean()
+        # print('Pretraining: validation loss:{:.3f}, l_t: {:.3f}, l_f:{:.3f}, l_c:{:.3f}'.format(val_loss, val_loss_t, val_loss_f, val_l_TF))
 
-    performance = [acc * 100, precision * 100, recall * 100, F1 * 100, total_auc * 100, total_prc * 100]
-    print('MLP Testing: Acc=%.4f| Precision = %.4f | Recall = %.4f | F1 = %.4f | AUROC= %.4f | AUPRC=%.4f'
-          % (acc*100, precision * 100, recall * 100, F1 * 100, total_auc*100, total_prc*100))
-    emb_test_all = torch.concat(tuple(emb_test_all))
-    return total_loss, total_acc, total_auc, total_prc, emb_test_all, trgs, performance
+    ave_loss = torch.tensor(total_loss).mean()
+
+    return ave_loss
diff --git a/code/config_files/Epilepsy_Configs.py b/code/config_files/Epilepsy_Configs.py
index 127af8b..f4aece1 100644
--- a/code/config_files/Epilepsy_Configs.py
+++ b/code/config_files/Epilepsy_Configs.py
@@ -13,7 +13,7 @@ def __init__(self):
         self.features_len_f = 24 # 13 #self.features_len   # the output results in time domain
 
         # training configs
-        self.num_epoch = 40 # 40
+        self.num_epoch = 2 # 40
         
         # optimizer parameters
         self.beta1 = 0.9
@@ -26,6 +26,8 @@ def __init__(self):
         self.batch_size = 32 #64 #  128
         self.target_batch_size = 16 # the size of target dataset (the # of samples used to fine-tune).
 
+        self.TSlength_aligned = 178
+
         self.Context_Cont = Context_Cont_configs()
         self.TC = TC()
         self.augmentation = augmentations()
diff --git a/code/config_files/MotionSample_Configs.py b/code/config_files/MotionSample_Configs.py
new file mode 100644
index 0000000..ef6d98c
--- /dev/null
+++ b/code/config_files/MotionSample_Configs.py
@@ -0,0 +1,44 @@
+class Config(object):
+    def __init__(self):
+        # model configs
+        self.input_channels = 8  # feature count
+        self.transformer_nhead = 2
+        self.transformer_num_layers = 2
+        self.embedding_len = 160  # final embedding len = embedding_len*2
+
+        # training configs
+        self.num_epoch = 2
+        
+        # optimizer parameters
+        self.beta1 = 0.9
+        self.beta2 = 0.99
+        self.lr = 3e-4  # original lr: 3e-4
+
+        # data parameters
+        self.drop_last = True
+        self.batch_size = 2 
+        self.TSlength_aligned = 900  # sequence length 15Hz * 60second
+        
+        self.Context_Cont = Context_Cont_configs()
+        self.TC = TC()
+        self.augmentation = augmentations()
+
+
+class augmentations(object):
+    def __init__(self):
+        self.jitter_scale_ratio = 0.001
+        self.jitter_ratio = 0.001
+        self.max_seg = 5
+
+
+class Context_Cont_configs(object):
+    def __init__(self):
+        self.temperature = 0.2
+        self.use_cosine_similarity = True
+        self.use_cosine_similarity_f = True
+
+
+class TC(object):
+    def __init__(self):
+        self.hidden_dim = 100
+        self.timesteps = 10
\ No newline at end of file
diff --git a/code/config_files/Motion_Configs.py b/code/config_files/Motion_Configs.py
new file mode 100644
index 0000000..6d5a75e
--- /dev/null
+++ b/code/config_files/Motion_Configs.py
@@ -0,0 +1,44 @@
+class Config(object):
+    def __init__(self):
+        # model configs
+        self.input_channels = 8  # feature count
+        self.transformer_nhead = 2
+        self.transformer_num_layers = 2
+        self.embedding_len = 160  # final embedding len = embedding_len*2
+
+        # training configs
+        self.num_epoch = 10
+        
+        # optimizer parameters
+        self.beta1 = 0.9
+        self.beta2 = 0.99
+        self.lr = 3e-4 # original lr: 3e-4
+
+        # data parameters
+        self.drop_last = True
+        self.batch_size = 16 
+        self.TSlength_aligned = 900  # sequence length 15Hz * 60second
+        
+        self.Context_Cont = Context_Cont_configs()
+        self.TC = TC()
+        self.augmentation = augmentations()
+
+
+class augmentations(object):
+    def __init__(self):
+        self.jitter_scale_ratio = 0.001
+        self.jitter_ratio = 0.001
+        self.max_seg = 5
+
+
+class Context_Cont_configs(object):
+    def __init__(self):
+        self.temperature = 0.2
+        self.use_cosine_similarity = True
+        self.use_cosine_similarity_f = True
+
+
+class TC(object):
+    def __init__(self):
+        self.hidden_dim = 100
+        self.timesteps = 10
\ No newline at end of file