data/data_loader.py

import h5py
import os, sys
import numpy as np
import json
import torch
from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler
import pandas as pd
import matplotlib.pyplot as plt
from collections import Counter, defaultdict
import cv2
from multiprocessing import Manager
import random 
from scipy.signal import find_peaks
import librosa
import pickle
from tqdm import tqdm
import math

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from configs.global_cfg import TARGET_RESAMPLE, DATASET_PATH, LABELS_PATH, EPOCH_LEN_MAX, EPOCH_LEN_MIN, DEMO_MAPS, PID_PREFIX_MAP

CHNGRP_PICK = {
    'EEG': ['C3-M2','C3-A2','C4-M1'],
    'EOG': ['E1-M2','EOG-Diff','Left-A2'],
    'ECG': ['ECG'],
    'AIRFLOW': ['AIRFLOW'],
    'EMG': ['EMG','Chin1-Chin2'],
    'SaO2': ['SaO2'],
    'ABD': ['ABD'], 
    'RESP': ['AIRFLOW'],  
}


class Pretrain_Dataset(Dataset):
    def __init__(self, 
                data_root,
                data_list,                 
                transforms=None,                 
                split='train',
                mode=None,
                num_classes=5,
                modalities=["EEG"],
                sig_len_target=60,
                ):

        super().__init__() 
                       
        self.data_root = data_root
        self.transforms = transforms                                       
        self.split = split      
        if mode is None:
            self.mode = split
        else:
            self.mode = mode
        #self.data_list = sorted(os.listdir(data_path))
        self.data_list = data_list
        self.total_num_epochs = 0
        self.total_num_subjects = len(self.data_list)
        
        self.subject_epoch_dict = {}
        self.all_labels = []        
        self.modalities = modalities
        self.sig_len_target = sig_len_target
        self.datasets = []

        # What datasets are used? phy, kiss, shhs1,...
        self.datasets = set([d[0] for d in self.data_list])        
        
        self.data_dict = {}
        self.stats = {}
        self.sig_info_dict = {}        
        self.chn_grp_map = dict()


        for dataset in self.datasets:
            for m in self.modalities:                
                chn_grp_dir = os.path.join(self.data_root, DATASET_PATH[dataset], m)
                if not os.path.exists(chn_grp_dir):
                    chn_grp_dir = os.path.join(self.data_root, DATASET_PATH[dataset], self.split, m)
                chns = [f.split('.')[0] for f in os.listdir(chn_grp_dir) if f.endswith('mmap')]
                self.chn_grp_map[f'{dataset}/{m}'] = chns
                # self.chn_grp_map[f'{dataset}/{m}'] = CHNGRP_MAP[f'{dataset}/{m}']
                for chn in chns:
                    info_path = os.path.join(chn_grp_dir, f'{chn}.pickle')
                    with open(info_path, 'rb') as f:
                        sig_info = pickle.load(f)
                    
                    self.sig_info_dict[f'{dataset}/{m}/{chn}'] = sig_info
                    sum_mu, sum_musq, cnt = 0.0, 0.0, 0.0
                    for k, v in sig_info['sig_info'].items():
                        if (dataset, k) in self.data_list:
                            start, end = v['pos']
                            mu, musq = v['stats']['sum'], v['stats']['sumsq']
                            sum_mu += mu
                            sum_musq += musq
                            cnt += (end - start)                            

                    self.data_dict[f'{dataset}/{m}/{chn}'] = np.memmap(os.path.join(chn_grp_dir, f'{chn}.mmap'), dtype="float32", 
                                                                       mode="r", 
                                                                       shape=sig_info['data_shape'])
                    sig_len = sig_info['data_shape'][1]  # Assume signal shape (sum(num_epochs), sig_len_epoch)
                    self.stats[f'{dataset}/{m}/{chn}'] = {'mean': sum_mu / cnt / sig_len if cnt > 0 else 0, 
                                                          'std': np.sqrt(sum_musq / cnt / sig_len - (sum_mu / cnt / sig_len) ** 2) if cnt > 0 else 1.0}
                
        self.index_map = self.create_index()        
        
    def __len__(self):
        return self.total_num_epochs
    def __repr__(self):
        return "Video_Frames_Dataset : Total Subjects : {}, Total Epochs : {}".format(self.total_num_subjects, self.total_num_epochs)

    def _get_channel(self, dataset, modality):
        """Select a channel for a modality based on the current split."""
        chns = self.chn_grp_map[f'{dataset}/{modality}']
        if self.mode == 'train':
            return random.choice(chns)
        for chn in chns:
            if chn in CHNGRP_PICK[modality]:
                return chn
        return chns[0]
    
    def _get_signal_for_modality(self, dataset, subject_name, m):
        """Get signal data for a specific modality."""
        
        chn = self._get_channel(dataset, m)        
        sig_info = self.sig_info_dict[f'{dataset}/{m}/{chn}']
        sub_start, sub_end = sig_info['sig_info'][subject_name]['pos']
        sub_epochs = sub_end - sub_start
        signals = self.data_dict[f'{dataset}/{m}/{chn}'][sub_start:sub_end]
        sig_len_epoch = signals.shape[1]
        
        return signals, chn, sig_len_epoch, sub_epochs

    def _load_signal_epoch(self, dataset, subject_name, modality, epoch, shift=0):
        """Load and slice a single modality epoch from mmap storage."""
        
        signals, chn, sig_len_epoch, _ = self._get_signal_for_modality(dataset, subject_name, modality)

        sig_len_epoch = signals.shape[1]
        hz = sig_len_epoch // 30
        signals_reshape = np.reshape(signals, (-1, 1))

        start = max(0, sig_len_epoch * epoch - shift * hz)
        end = start + self.sig_len_target * hz
        if signals_reshape.shape[0] < end:
            end = signals_reshape.shape[0]
            start = end - self.sig_len_target * hz

        return chn, signals_reshape[start:end], sig_len_epoch, hz

    def _signal_to_tensor(self, signal_epoch, sig_len_epoch, dataset, modality, chn):
        """Convert a sliced epoch to a normalized tensor."""
        hz = sig_len_epoch // 30
        if modality == 'SaO2':
            target_signal = torch.tensor(signal_epoch, dtype=torch.float32)
            if hz != TARGET_RESAMPLE:
                target_signal = target_signal.repeat_interleave(TARGET_RESAMPLE, dim=0)
            return (torch.clamp(target_signal, 0.7, 1.0) - 0.7) / 0.3

        std = self.stats[f'{dataset}/{modality}/{chn}']['std']
        signal_epoch = signal_epoch * 1 / std
        if hz != TARGET_RESAMPLE:
            signal_epoch = librosa.resample(signal_epoch, orig_sr=sig_len_epoch // 30, target_sr=TARGET_RESAMPLE,
                                            axis=0, res_type='soxr_qq')
        return torch.tensor(signal_epoch, dtype=torch.float32)

    def _build_modality_tensor(self, dataset, subject_name, epoch, modality, shift=0):
        """Shared path for reading, slicing, and normalizing a modality."""
        chn, signal_epoch, sig_len_epoch, _ = self._load_signal_epoch(dataset, subject_name, modality, epoch, shift=shift)
        return self._signal_to_tensor(signal_epoch, sig_len_epoch, dataset, modality, chn)
    
    def create_index(self):
        '''
        Create index for the dataset
        '''
        index_map = []

        for dataset, subject_name in self.data_list:
            m = self.modalities[0]
            chn = self.chn_grp_map[f'{dataset}/{m}'][0]            
            sig_info = self.sig_info_dict[f'{dataset}/{m}/{chn}']            
            num_epochs = sig_info['sig_info'][subject_name]['pos'][1] - sig_info['sig_info'][subject_name]['pos'][0]
            # if num_epochs <= EPOCH_LEN_MIN:
            #     print(f"Warning: num_epochs less than EPOCH_LEN_MIN for {dataset}, {subject_name}. Num_epochs: {num_epochs}, skip this row")
            #     continue
            for epoch in range(num_epochs):
                index_map.append((dataset, subject_name, epoch))

        self.total_num_epochs = len(index_map)

        return index_map
    
    def __getitem__(self, idx, modalities_override=None):
       
        # Return frames, label, (subject_name, epoch)
        dataset, subject_name, epoch = self.index_map[idx]    

        inputs = dict()               
        modalities_to_use = modalities_override if modalities_override is not None else self.modalities        
        
        
        shift_max = self.sig_len_target - 30
        if shift_max > 0:
            shift = np.random.randint(0, shift_max)  # (in seconds)
        elif shift_max == 0:
            shift = 0
        else:
            shift = np.random.randint(shift_max, 0)  # (in seconds)
        
        # inputs["label"] = torch.tensor(label, dtype=torch.long)        
        for m in modalities_to_use:
            inputs[m] = self._build_modality_tensor(dataset, subject_name, epoch, m, shift=shift)

        return inputs
            
    
    



class SFT_Dataset(Pretrain_Dataset):
    def __init__(self,                 
                data_root,                
                data_list,       
                ft_task='sleep_staging',          
                transforms=None, 
                split='train',
                mode=None, 
                num_classes=5,
                modalities=["EEG"], 
                sig_len_target=30                
                ):
        super().__init__(data_root, data_list, transforms, split, mode, num_classes, modalities, sig_len_target)
        self.ft_task = ft_task
        self.num_classes = num_classes
        
        # if self.ft_task == 'sleep_staging' and self.num_classes == 4:
        #     self.label_change = {0:0, 1:1, 2:1, 3:2, 4:3}
            
        if self.ft_task == 'apnea_seg' and self.num_classes == 2:
            self.label_change = {0:0, 1:1, 2:1}
            
        self.get_labels()

    def __len__(self):        
        return self.total_num_epochs

        
    def get_labels(self):        
        self.all_labels = []
        shift = (self.sig_len_target - 30) // 2
        prev_subject_name = ''   

        for dataset, subject_name, epoch in tqdm(self.index_map):
            if not subject_name == prev_subject_name:            
                suffix = 'ah' if self.ft_task == 'apnea_seg' else 'ann'
                label_dir = os.path.join(self.data_root, LABELS_PATH[self.ft_task][dataset])
                label_path = os.path.join(label_dir, f'{subject_name}_{suffix}.npy')
                if not os.path.exists(label_path):
                    label_path = os.path.join(label_dir, self.split, f'{subject_name}_{suffix}.npy')
                labels = np.load(label_path).astype(np.int64) # (num_epochs, 1)
            
            if self.ft_task == 'apnea_seg':
                labels = labels.reshape(-1,)                
                labels_slice = labels[epoch*30 : (epoch+1)*30]
                
                if self.num_classes == 2:
                    labels_slice = np.array([self.label_change[lbl] for lbl in labels_slice])
                
                self.all_labels.append(labels_slice)
                
            else:             
                self.all_labels.append(labels[epoch])
            prev_subject_name = subject_name

    def get_all_labels(self):
        return self.all_labels        


    def __getitem__(self, idx, modalities_override=None):        
        # Return frames, label, (subject_name, epoch)

        
        dataset, subject_name, epoch = self.index_map[idx]    
        label = self.all_labels[idx]                
    
        inputs = dict()               
        modalities_to_use = modalities_override if modalities_override is not None else self.modalities
        
        shift = (self.sig_len_target - 30) // 2
        
        for m in modalities_to_use:
            inputs[m] = self._build_modality_tensor(dataset, subject_name, epoch, m, shift=shift)

        # Return an empty boolean tensor for mask so default_collate can batch samples.
        # Downstream training code treats an empty tensor as 'no mask'.
        return inputs, torch.tensor(label, dtype=torch.long), torch.tensor([], dtype=torch.bool), torch.tensor([], dtype=torch.bool), subject_name # No mask


class Seq_Dataset(Pretrain_Dataset):
    def __init__(self,                 
                data_root,                
                data_list,                       
                transforms=None, 
                split='train',
                mode=None, 
                num_classes=5,
                modalities=["EEG"], 
                sig_len_target=30,
                pretrain_macro=False,
                seq_truncate=False,
                ft_task='sleep_staging',
                dual_views=False,             
                ):
        super().__init__(data_root, data_list, transforms, split, mode, num_classes, modalities, sig_len_target)

        self.num_classes = num_classes        
        self.get_demographics()
        self.epoch_accum = 0
        self.seq_truncate = seq_truncate
        self.max_num_epochs = EPOCH_LEN_MAX if seq_truncate else 1440
        self.min_num_epochs = EPOCH_LEN_MIN if seq_truncate else 0
        self.pretrain_macro = pretrain_macro
        self.dual_views = dual_views
        self.ft_task = ft_task
        
        self.create_index_full()      
        print("Max length:", self.max_num_epochs)
        
        if not self.pretrain_macro:
            self.get_labels()

    def reshuffle_full_night_index(self):
        """
        Reshuffle the full night index. Call this at the start of each epoch
        to get a different ordering of samples.
        """
        self.epoch_accum += 1
        self.create_index_full()
        if not self.pretrain_macro:
            self.get_labels()
    
    def _validate_demographics(self, dataset, subject_name):
        """Validate demographics data for a subject. Returns (is_valid, age, bmi)."""
        demo = self.demographics.get(f'{dataset}/{subject_name}', {})
        
        age = demo.get('Age', None)
        sex = demo.get('Sex', None)
        bmi = demo.get('BMI', None)
        
        # Validate age
        if age is None or (isinstance(age, float) and math.isnan(age)):
            print(f"Warning: Age is missing for {dataset}, {subject_name}. Skipping subject")
            return False, None, None
        
        # Validate sex
        if sex is None:
            print(f"Warning: Sex is missing for {dataset}, {subject_name}. Skipping subject")
            return False, None, None
        
        # Validate BMI
        if bmi is None:
            print(f"Warning: BMI is missing for {dataset}, {subject_name}. Skipping subject")
            return False, None, None
        
        try:
            bmi = float(bmi)
        except (ValueError, TypeError):
            print(f"Warning: BMI is not a number for {dataset}, {subject_name}. Skipping subject")
            return False, None, None
        
        if isinstance(bmi, float) and math.isnan(bmi):
            print(f"Warning: BMI is NaN for {dataset}, {subject_name}. Skipping subject")
            return False, None, None
        
        return True, age, bmi
    
    def create_index_full(self):
        '''Create index for the dataset'''
        index_map, len_map, ages, bmis = [], [], [], []

        for dataset, subject_name in self.data_list:
            m = self.modalities[0]
            chn = self.chn_grp_map[f'{dataset}/{m}'][0]            
            sig_info = self.sig_info_dict[f'{dataset}/{m}/{chn}']            
            num_epochs = sig_info['sig_info'][subject_name]['pos'][1] - sig_info['sig_info'][subject_name]['pos'][0]
            
            if not self.seq_truncate:
                self.max_num_epochs = max(num_epochs, self.max_num_epochs)            
            
            if num_epochs <= self.min_num_epochs:
                print(f"Warning: Sequence length < EPOCH_LEN_MIN for {dataset}, {subject_name}. Num_epochs: {num_epochs}, skipping")
                continue
            
            if not self.pretrain_macro:
                # For finetuning, just collect subjects
                index_map.append((dataset, subject_name))
                continue
            
            # For pretraining with macro demographics, validate demographics
            is_valid, age, bmi = self._validate_demographics(dataset, subject_name)
            if not is_valid:
                continue
            
            ages.append(age)
            bmis.append(bmi)
            len_map.append(num_epochs)
            index_map.append((dataset, subject_name))
        
        if not self.pretrain_macro:
            self.index_map = index_map
            self.total_subjects = len(index_map)
            return

        # For pretraining: reorder by demographics and sequence length
        self._reorder_index_by_demographics(index_map, len_map, ages, bmis)
    
    def _reorder_index_by_demographics(self, index_map, len_map, ages, bmis):
        """Reorder index map by demographics and sequence length for balanced batching."""
        zipped_pairs = list(zip(index_map, len_map, ages, bmis))
        
        # Alternate sorting by BMI (even epochs) and age (odd epochs)
        if self.epoch_accum % 2 == 0:
            sorted_pairs = sorted(zipped_pairs, key=lambda x: x[3])  # Sort by BMI
        else:
            sorted_pairs = sorted(zipped_pairs, key=lambda x: x[2])  # Sort by age
        
        index_map = [pair[0] for pair in sorted_pairs]
        len_map = [pair[1] for pair in sorted_pairs]
        
        self.total_subjects = len(index_map)
        
        # Group by age and apply secondary sorting within groups
        num_age_grp = 4
        age_grp_size = self.total_subjects // num_age_grp
        age_groups = []
        residuals = []
        
        for i in range(num_age_grp):
            start_idx = i * age_grp_size
            end_idx = (i + 1) * age_grp_size
            index_map_grp = index_map[start_idx:end_idx]
            len_map_grp = len_map[start_idx:end_idx]
            
            # Sort by length: ascending for epochs 1,2; descending for 0,3
            reverse = (self.epoch_accum % 4 == 0 or self.epoch_accum % 4 == 3)
            index_map_grp_sorted = [pair[0] for pair in sorted(
                zip(index_map_grp, len_map_grp), 
                key=lambda x: x[1], 
                reverse=reverse
            )]
            
            # Group by 4 subjects
            len_sub4 = len(index_map_grp_sorted) // 4
            sub4 = [index_map_grp_sorted[4*k:4*(k+1)] for k in range(len_sub4)]
            age_groups.append(sub4)
            residuals.extend(index_map_grp_sorted[4*len_sub4:])
            
            if i == num_age_grp - 1 and end_idx < self.total_subjects:
                residuals.extend(index_map[end_idx:])
        
        # Interleave groups and add residuals
        final_index_map = []
        for items in zip(*age_groups):
            flattened = [item for sublist in items for item in sublist]
            final_index_map.extend(flattened)
        final_index_map.extend(residuals)
        
        assert len(final_index_map) == self.total_subjects, \
            f"Index map length mismatch: {len(final_index_map)} vs {self.total_subjects}"
        
        self.index_map = final_index_map
        self.total_subjects = len(self.index_map)
        

    def __len__(self):
        return self.total_subjects
        
    def get_demographics(self):
        """Load demographic data for all datasets."""
        self.demographics = dict()
        

        
        # Datasets to skip
        skip_datasets = {'isruc', 'cfs', 'sof'}
        
        dataset_list = list(set([d[0] for d in self.data_list]))
        for dataset in tqdm(dataset_list):
            if any(s in dataset for s in skip_datasets):
                continue
            
            # Get demographics file path
            demo_file = DEMO_MAPS.get(dataset, f'{dataset}_demographics.pkl')
            demo_path = os.path.join(self.data_root, 'demographics', demo_file)
            
            if not os.path.exists(demo_path):
                print(f"Warning: Demographics file not found: {demo_path}")
                continue
            
            with open(demo_path, 'rb') as f:
                demo_data = pickle.load(f)

            for pid, demo in demo_data.items():
                if dataset in PID_PREFIX_MAP:
                    prefix = PID_PREFIX_MAP[dataset]
                    if dataset in ['mros1', 'mros1-val', 'mros1-te', 'mros2', 'mros2-val', 'mros2-te']:
                        pid = f'{prefix}-{pid.lower()}'
                    elif dataset in ['mesa', 'mesa-val', 'mesa-te']:
                        pid = f'{prefix}-{pid:04d}'
                    elif dataset == 'cfs':
                        pid = f'{dataset}-visit5-{pid}'
                    else:
                        pid = f'{dataset}-{pid}'
                
                self.demographics[f'{dataset}/{pid}'] = demo 
        
    def get_labels(self):        
        self.all_labels = []
        shift = (self.sig_len_target - 30) // 2
        prev_subject_name = ''                
        
        if self.ft_task == 'diagnosis':
            # Disease prediction
            events_csv = 'datasets/#_SHHS/datasets/shhs-cvd-summary-dataset-0.20.0.csv'
            event_df = pd.read_csv(events_csv)

            is_event_cols = ['angina','cvd_death','chf','chd_death','mi','stroke']
            event_time_cols = ['ang_date','cvd_dthdt','chf_date','chd_dthdt','mi_date','stk_date']

            # 1. 데이터 병합 (censdate 포함)
            # 날짜 보정을 위해 'censdate' 열을 함께 가져옵니다.
            df_filtered = pd.concat([
                event_df[['nsrrid', 'censdate']], 
                event_df[is_event_cols], 
                event_df[event_time_cols]
            ], axis=1).set_index('nsrrid')

            # 2. 정합성 검사 및 보정
            for event_col, time_col in zip(is_event_cols, event_time_cols):
                # 조건 A: 이벤트가 1 이상인데 날짜가 NaN이거나 0인 경우
                # 이 경우 이벤트를 0으로 변경하고 날짜는 censdate를 참조하게 함
                invalid_mask = (df_filtered[event_col] >= 1) & (df_filtered[time_col].isna() | (df_filtered[time_col] == 0))
                # invalid_mask = (df_filtered[event_col] >= 1) & ((df_filtered[time_col] == 0))
                df_filtered.loc[invalid_mask, event_col] = 0
                
                # 조건 B: 날짜열의 값이 없거나(NaN) 0인 모든 경우에 대해 'censdate' 값으로 채움
                # fillna를 사용하여 NaN을 채우고, 0인 값은 replace를 통해 처리할 수 있습니다.
                df_filtered[time_col] = df_filtered[time_col].fillna(df_filtered['censdate'])
                df_filtered.loc[df_filtered[time_col] == 0, time_col] = df_filtered['censdate']

            # 3. 최종 결측치 처리 (censdate마저 없는 경우를 대비한 방어 코드)
            df_filtered[is_event_cols] = df_filtered[is_event_cols].fillna(0)
            # 만약 censdate도 NaN이라서 여전히 결측치라면 1e6으로 최종 처리
            df_filtered[event_time_cols] = df_filtered[event_time_cols].fillna(365*6)
            
            df_filtered[is_event_cols] = np.minimum(df_filtered[is_event_cols], 1)

            study_ids = set(df_filtered.index)            
            

            labels_dict = {}
            for study_id in study_ids:
                row = df_filtered.loc[study_id]
                
                # 중복 ID가 있을 경우를 대비해 첫 번째 행 선택
                if isinstance(row, pd.DataFrame):
                    row = row.iloc[0]
                    
                is_event_row = [int(val) for val in row[is_event_cols]]
                event_time_row = [float(val) for val in row[event_time_cols]]
                
                labels_dict[str(study_id)] = {
                    "is_event": is_event_row,
                    "event_time": event_time_row
                }            
        
        for sub_info in tqdm(self.index_map):            

            dataset, subject = sub_info
            
            if self.ft_task == 'diagnosis':
                if not dataset == 'shhs1':
                    raise NotImplementedError("Diagnosis task is only implemented for shhs1 dataset.")
                study_id = subject.split('-')[1]
                self.all_labels.append(labels_dict[study_id])                  
                
            else:
                suffix = 'ah' if self.ft_task == 'apnea_seg' else 'ann'
                if self.ft_task in ['age','sex','ahi']:
                    ft_task = 'sleep_staging'
                else:
                    ft_task = self.ft_task
                label_dir = os.path.join(self.data_root, LABELS_PATH[ft_task][dataset])
                label_path = os.path.join(label_dir, f'{subject}_{suffix}.npy')
                if not os.path.exists(label_path):
                    label_path = os.path.join(label_dir, self.split, f'{subject}_{suffix}.npy')
                labels = np.load(label_path).astype(np.int64) # (num_epochs, 1)
                labels = labels.reshape(-1,)
                self.all_labels.append(labels)
        
        
    def get_all_labels(self):
        return self.all_labels
    
    def _process_signal(self, signal_epochs, sig_len_epoch, dataset, m, chn):
        """Convert raw signal to target tensor with appropriate normalization."""
        signals_reshape = np.reshape(signal_epochs, (-1, 1))
        Hz = sig_len_epoch // 30
        
        if m == 'SaO2':
            resampled = signals_reshape
            target_signal = torch.tensor(resampled, dtype=torch.float32)
            if Hz != TARGET_RESAMPLE:
                target_signal = target_signal.repeat_interleave(TARGET_RESAMPLE, dim=0)
            target_signal = (torch.clamp(target_signal, 0.7, 1.0) - 0.7) / 0.3
        else:
            std = self.stats[f'{dataset}/{m}/{chn}']['std']
            signals_reshape = signals_reshape * 1 / std
            if Hz != TARGET_RESAMPLE:
                resampled = librosa.resample(signals_reshape, orig_sr=sig_len_epoch//30, 
                                            target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq')
            else:
                resampled = signals_reshape
            target_signal = torch.tensor(resampled, dtype=torch.float32)
        
        return target_signal
    
    def _shift_signals(self, signals, shift_epochs, shift_fwd):
        """Shift signals by specified number of epochs."""
        if shift_epochs == 0:
            return signals
        if shift_fwd:
            return np.concatenate([signals[shift_epochs:], signals[-shift_epochs:]])
        return np.concatenate([signals[:shift_epochs], signals[:-shift_epochs]])
    
    def _pad_signals(self, signals, sub_epochs):
        """Pad signals to max_num_epochs."""
        if sub_epochs >= self.max_num_epochs:
            return signals[:self.max_num_epochs]
        pad_len = self.max_num_epochs - sub_epochs
        padding = np.zeros_like(signals[-1:]).repeat(pad_len, axis=0)
        return np.concatenate([signals, padding], axis=0)
    
    def _build_demo_dict(self, dataset, subject_name):
        """Build demographic dictionary for a subject."""
        demo = {}
        demo_key = f'{dataset}/{subject_name}'
        
        if self.pretrain_macro:
            age = self.demographics[demo_key]['Age']
            sex = self.demographics[demo_key]['Sex']
            bmi = self.demographics[demo_key]['BMI']
            demo['age'] = float(age)
            demo['sex'] = sex
            demo['bmi'] = float(bmi)
        else:
            # Finetuning: include optional demographics
            demo_values = self.demographics.get(demo_key, {})
            age = demo_values.get('Age', -1)
            sex = demo_values.get('Sex', -1)
            bmi = demo_values.get('BMI', -1)
            ahi = demo_values.get('AHI', -1)
            
            # Convert to float and handle NaN
            demo['age'] = float(age) if age is not None and not (isinstance(age, float) and math.isnan(age)) else -1.0
            demo['sex'] = sex
            demo['bmi'] = float(bmi) if bmi is not None and not (isinstance(bmi, float) and math.isnan(bmi)) else -1.0
            demo['ahi'] = float(ahi) if ahi is not None and not (isinstance(ahi, float) and math.isnan(ahi)) else -1.0
        
        return demo
   
    def __getitem__(self, idx, modalities_override=None):
        """Get a full-night sequence of signals with optional dual views."""
        dataset, subject_name = self.index_map[idx]
        modalities_to_use = modalities_override or self.modalities
        
        # Build demographic dictionary
        demo = self._build_demo_dict(dataset, subject_name)
        
        # Initialize label and mask
        label = None
        if not self.pretrain_macro:
            label = self.all_labels[idx]
        
        # Get signal info from first modality to create mask
        m_first = modalities_to_use[0]
        _, _, _, sub_epochs = self._get_signal_for_modality(dataset, subject_name, m_first)
        
        # Update label and create mask based on truncation/padding
        if not self.pretrain_macro and self.ft_task in ['age', 'sex', 'ahi', 'sleep_staging']:
            if sub_epochs >= self.max_num_epochs:
                label = label[:self.max_num_epochs]
                mask = torch.ones(self.max_num_epochs, dtype=torch.bool)
            else:
                label = np.concatenate([label, label[-1:].repeat(self.max_num_epochs - sub_epochs, axis=0)])
                mask = torch.zeros(self.max_num_epochs, dtype=torch.bool)
                mask[:sub_epochs] = True
        else:
            if sub_epochs >= self.max_num_epochs:
                mask = torch.ones(self.max_num_epochs, dtype=torch.bool)
            else:
                mask = torch.zeros(self.max_num_epochs, dtype=torch.bool)
                mask[:sub_epochs] = True
        
        # Process signals for dual or single views
        if self.dual_views:
            inputs_view1, inputs_view2 = self._process_dual_views(dataset, subject_name, modalities_to_use, sub_epochs)
            inputs_out = (inputs_view1, inputs_view2)
        else:
            inputs_out = self._process_single_view(dataset, subject_name, modalities_to_use, sub_epochs)
        
        # Return based on task type
        if self.pretrain_macro:
            return inputs_out, torch.tensor([], dtype=torch.long), mask, demo, subject_name
        elif self.ft_task == 'diagnosis':
            is_event = torch.tensor(label["is_event"])
            event_time = torch.tensor(label["event_time"], dtype=torch.float32)
            return inputs_out, is_event, event_time, mask, demo, subject_name
        else:
            return inputs_out, torch.tensor(label, dtype=torch.long), mask, demo, subject_name
    
    def _process_single_view(self, dataset, subject_name, modalities_to_use, sub_epochs):
        """Process signals for single view."""
        inputs = {}
        for m in modalities_to_use:
            signals, chn, sig_len_epoch, _ = self._get_signal_for_modality(dataset, subject_name, m)
            signals = self._pad_signals(signals, sub_epochs)
            inputs[m] = self._process_signal(signals, sig_len_epoch, dataset, m, chn)
        return inputs
    
    def _process_dual_views(self, dataset, subject_name, modalities_to_use, sub_epochs):
        """
        Used for Macro-Encoder alternative training methods (SimCLR and DINO)
        Process signals for dual views with random augmentation.
        """
        # Generate random shift parameters for two views
        shift_epochs_1 = np.random.randint(0, 60)
        shift_fwd_1 = np.random.rand() > 0.5
        shift_epochs_2 = np.random.randint(0, 60)
        shift_fwd_2 = np.random.rand() > 0.5
        
        # Random modality dropout
        zero_modalities = np.random.choice(len(modalities_to_use), 6, replace=False)
        k1 = np.random.randint(6)
        zero_modalities_1 = zero_modalities[:k1]
        zero_modalities_2 = zero_modalities[k1:]
        
        inputs_view1, inputs_view2 = {}, {}
        
        for i, m in enumerate(modalities_to_use):
            signals, chn, sig_len_epoch, _ = self._get_signal_for_modality(dataset, subject_name, m)
            
            # Create two shifted views
            signals_v1 = self._shift_signals(signals, shift_epochs_1, shift_fwd_1)
            signals_v2 = self._shift_signals(signals, shift_epochs_2, shift_fwd_2)
            
            # Apply random modality dropout
            if i in zero_modalities_1:
                signals_v1 = np.zeros_like(signals_v1)
            if i in zero_modalities_2:
                signals_v2 = np.zeros_like(signals_v2)
            
            # Pad and process
            signals_v1 = self._pad_signals(signals_v1, sub_epochs)
            signals_v2 = self._pad_signals(signals_v2, sub_epochs)
            
            inputs_view1[m] = self._process_signal(signals_v1, sig_len_epoch, dataset, m, chn)
            inputs_view2[m] = self._process_signal(signals_v2, sig_len_epoch, dataset, m, chn)
        
        return inputs_view1, inputs_view2
        
            
        
            

def collate_fn_random_modalities(batch_indices, dataset):
    # Randomly select two modalities for this batch
    all_modalities = dataset.modalities
    selected_modalities = sorted(random.sample(all_modalities, 2))
    # selected_modalities = all_modalities  # Use all modalities for now
    # Fetch data for each index with the selected modalities
    batch = [dataset.__getitem__(idx, modalities_override=selected_modalities) for idx in batch_indices]    
    # inputs = zip(*batch)
    inputs = batch
    inputs_ = {m: [] for m in selected_modalities}
    for m in selected_modalities:
        for item in inputs:        
            inputs_[m].append(item[m])
    inputs = {m: torch.stack(inputs_[m], dim=0) for m in selected_modalities}
    return inputs

class Pretrain_Dataset_Wrapper(Dataset):
    def __init__(self, dataset):
        self.dataset = dataset
    def __len__(self):
        return len(self.dataset)
    def __getitem__(self, idx):
        return idx  # Just return the index

# build loader with sample
def build_loader(data_root,                     
                data_list=None, 
                batch_size=1, 
                shuffle=True, 
                transforms=None, 
                sampler='weighted', 
                num_workers=8, 
                proportion = 1.0, 
                split='train',
                mode=None,                
                num_classes=5,                 
                modalities=["EEG"],
                pretrain_micro=False,
                pretrain_macro=False,
                sig_len_target=30,
                ft_task='sleep_staging',
                full_night=False,                
                seq_truncate=False,
                dual_views=False,
                ):
    """_summary_

    Args:
        data_path (_type_): _description_
        mode (str, optional): _description_. Defaults to 'epochs'.
        batch_size (int, optional): _description_. Defaults to 1.
        shuffle (bool, optional): _description_. Defaults to True.
        transforms (_type_, optional): _description_. Defaults to None.
        sampler (str, optional): _description_. Defaults to 'weighted'.
        train (bool, optional): _description_. Defaults to True.

    Returns:
        _type_: _description_
    """
    if pretrain_micro:
        # Create dataset    
        print("Using Pretrain Dataset")
        dataset = Pretrain_Dataset(data_root,
                                        data_list,                                         
                                        transforms, 
                                        mode=mode,
                                        split=split,
                                        num_classes=num_classes,
                                        modalities=modalities,
                                        sig_len_target=sig_len_target
                                        )
      
        print("Dataset created")
        loader = DataLoader(dataset, 
                            batch_size=batch_size,                             
                            num_workers=num_workers, 
                            shuffle=shuffle,
                            # collate_fn=lambda batch_indices: collate_fn_random_modalities(batch_indices, dataset),
                            prefetch_factor=2)

    else:        
        if full_night:
            print("Using sequence dataset")        
            dataset = Seq_Dataset(data_root,                                     
                                    data_list,                                         
                                    transforms=transforms,
                                    split=split,
                                    mode=mode, 
                                    num_classes=num_classes,
                                    modalities=modalities,
                                    sig_len_target=sig_len_target,
                                    ft_task=ft_task,
                                    pretrain_macro=pretrain_macro,
                                    seq_truncate=seq_truncate,
                                    dual_views=dual_views
                                    )      
        else:
            print("Using Epoch-level Fine-tuning Dataset. Task:", ft_task)
            dataset = SFT_Dataset(data_root,                                     
                                    data_list,                                         
                                    ft_task=ft_task,
                                    transforms=transforms,
                                    split=split,
                                    mode=mode, 
                                    num_classes=num_classes,
                                    modalities=modalities,
                                    sig_len_target=sig_len_target                                    
                                    )      
            
        print("Dataset created")    
        
        if sampler=='weighted':
            # Balanced Sampler
            labels = dataset.get_all_labels()
            
            # class_count = Counter(labels)
            class_count = np.zeros(num_classes, dtype=np.int64)
            unique, counts = np.unique(labels, return_counts=True)
            for i, cnt in zip(unique, counts):            
                class_count[i] += cnt        
            print("Class count:", class_count)

            class_weights = 1.0 / torch.tensor(class_count, dtype=torch.float)
            # Length : min(class_count)
            sample_weights  = [class_weights[i] for i in labels]
            sample_weights = torch.tensor(sample_weights, dtype=torch.float)
            print("Class weights:", sample_weights)
            # Length : min(class_count)
            sampler = torch.utils.data.sampler.WeightedRandomSampler(sample_weights, int(len(labels)*proportion), replacement=True)
            loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, num_workers=num_workers, prefetch_factor=2)
            
        elif sampler == 'random':
            # print("Random sampler, num_samples : ", int(len(dataset)*proportion), "Length dataset", len(dataset))
            # Proportion does not look effective....
            # sampler = RandomSampler(dataset, replacement=True, num_samples=int(len(dataset)*proportion))
            labels = dataset.get_all_labels()
            
            class_count = np.zeros(num_classes, dtype=np.int64)
            unique, counts = np.unique(labels, return_counts=True)        
            for i, cnt in zip(unique, counts):            
                class_count[i] += cnt 
            print("Class count:", class_count)

            # class_weights = torch.tensor(class_count, dtype=torch.float)        
            class_weights = torch.ones(num_classes, dtype=torch.float) / num_classes
            sample_weights  = [1.0 for i in labels]
            sampler = torch.utils.data.sampler.WeightedRandomSampler(sample_weights, int(len(labels)*proportion), replacement=True)
            # loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, num_workers=num_workers, prefetch_factor=2)
            loader = DataLoader(dataset, 
                                batch_size=batch_size, 
                                sampler=sampler, 
                                num_workers=num_workers, 
                                collate_fn=lambda batch_indices: collate_fn_random_modalities(batch_indices, dataset),
                                prefetch_factor=2)
        elif sampler == 'distributed': 
            from torch.utils.data import DistributedSampler
            # 2. Manually create the DistributedSampler with shuffle=False
            sampler = DistributedSampler(
                dataset, 
                shuffle=shuffle  # <--- The key override
            )              
            loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, num_workers=num_workers, prefetch_factor=2, drop_last=True)

        elif sampler is None:               
            loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, prefetch_factor=2)
        else:
            assert False, "Not Implemented loader"    
    return loader