models/classification.py

from functools import partial
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from models.encoders import MultiEncoderClassification
from models.transformer_modules import RevIN
from models.mamba_modules import SeqEncoder
from models.fusion import FusionUpsampleBlock
from models.utils import flip_padded_sequences
from einops import rearrange
from configs.global_cfg import TARGET_RESAMPLE


class ClassificationMultiConv1d(nn.Module):
    """ Masked Autoencoder with VisionTransformer backbone
    """    

    def __init__(self, sig_len=6000, patch_size=10,
                 embed_dim=[32,64,128], depth=[2,2,4], specific_depth=[2,2], num_heads=8,
                 fusion_heads=8, fusion_depth=2,
                 mlp_ratio=4, norm_layer=nn.LayerNorm, num_classes=5,
                 modalities=["EEG","EOG"], num_experts=2, num_experts_per=1,  
                 use_shared=True, use_moe=True, fusion_method='concat', 
                 use_revin=False, pos_enc='sinusoidal',  ft_task='sleep_staging'):
        
        super().__init__()     
    
        self.sig_len = sig_len
        self.patch_size = patch_size
        self.modalities = modalities
        
        self.num_experts = num_experts
        self.num_experts_per = num_experts_per
        self.use_shared = use_shared
        self.use_moe = use_moe
        self.fusion_method = fusion_method
        self.embed_dim = embed_dim
        self.depth = depth
        self.pos_enc = pos_enc

        self.encoder = MultiEncoderClassification(sig_len=sig_len, patch_size=patch_size,
                 embed_dim=embed_dim, depth=depth, num_heads=num_heads, specific_depth=specific_depth,
                 mlp_ratio=mlp_ratio, norm_layer=norm_layer,
                 modalities=self.modalities, num_experts=num_experts, num_experts_per=num_experts_per,
                 use_shared=use_shared, use_moe=use_moe, pos_enc=pos_enc)
        # self.head = nn.Linear(embed_dim[-1], num_classes) if num_classes > 0 else nn.Identity()                
        self.num_patches = self.encoder.num_patches        

        self.fuser = FusionUpsampleBlock(dim_list=embed_dim[2:], 
                                         fusion_method=fusion_method, 
                                         fusion_depth=fusion_depth, 
                                         fusion_heads=fusion_heads,
                                         mlp_ratio=mlp_ratio, 
                                         norm_layer=norm_layer, 
                                         pos_enc=pos_enc, 
                                         use_shared=use_shared, 
                                         use_upsample=True, # if ft_task=='apnea_seg' else False, 
                                         modalities=self.modalities)
    
        self.ft_task = ft_task
        if self.ft_task == 'sleep_staging':
            post_fusion_dim = embed_dim[3] * len(self.modalities)
            self.head = nn.Linear(post_fusion_dim, num_classes) if num_classes > 0 else nn.Identity()
        elif self.ft_task == 'apnea_seg':
            post_fusion_dim = embed_dim[2] * len(self.modalities)
            
            self.patch_per_sec = TARGET_RESAMPLE // (self.patch_size * 2) # Encoder downsamples by 2 for consistency in pretrainng and finetuning
            self.head = nn.Linear(post_fusion_dim, num_classes)  # 1D conv for sequence labeling
            

        self.use_revin = use_revin
        if self.use_revin:
            self.revin = nn.ModuleDict()
            for m in self.modalities:
                if m != 'SaO2':
                    self.revin[m] = RevIN(num_features=1, affine=False)      
        


        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            # we use xavier_uniform following official JAX ViT:
            torch.nn.init.xavier_uniform_(m.weight)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)

        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0) 

    def forward_features(self, inputs):


        modalities_to_use = sorted(list(inputs.keys()))        
        
        if self.use_revin:
            for i, m in enumerate(modalities_to_use):                
                if m == 'SaO2':                    
                    inputs[m] = inputs[m][:, ::TARGET_RESAMPLE, :] # N,L* TARGET_RESAMPLE,1 -> N,L,1
                else:
                    inputs[m] = self.revin[m](inputs[m], mode='norm')
        
        latents_shared, latents_specific, gating_scores = self.encoder(inputs, modalities_to_use)
        
        latents_fused, upsampled = self.fuser(latents_shared, latents_specific)

        if self.ft_task == 'apnea_seg':
            latents = torch.stack(upsampled, dim=0) # (M, N, L, D * 2 or 4)
            latents = rearrange(latents, 'm n (l s) d -> n l s (m d)', s=self.patch_per_sec)
            latents = latents.mean(2)
            lengths = latents.shape[1]
            if lengths > 30:
                shift = (lengths - 30) // 2
                latents = latents[:, shift:shift+30, :]
        else:
            # latents = torch.concat(latents_cat, dim=-1) # (M, N, L, D * 2 or 4)                        
            latents = torch.stack(latents_fused, dim=0)
            latents = rearrange(latents, 'm n l d -> n m l d')
            latents = latents.mean(2)
            latents = rearrange(latents, 'n m d -> n (m d)')
        return latents

    def forward(self, sig):
        x = self.forward_features(sig)
        x = self.head(x)
        return x    

class ClassificationMultiConv1dMamba(nn.Module):
    """ Masked Autoencoder with VisionTransformer backbone
    """    

    def __init__(self, sig_len=6000, patch_size=10, 
                 embed_dim=[32,64,128], depth=[2,2,4], specific_depth=[2,2], num_heads=8,
                 seq_dim=512, seq_depth=2,
                 mlp_ratio=4, norm_layer=nn.LayerNorm, num_classes=5,
                 modalities=["EEG","EOG"], num_experts=2, num_experts_per=1,  
                 use_shared=True, use_moe=True, fusion_method='concat', fusion_depth=2, fusion_heads=8,
                 use_revin=False, pos_enc='sinusoidal',
                 ft_task='sleep_staging', d_state=16, d_conv=4, expand=2, 
                 bidi=True, pretrain=False, use_micro=False,
                 demo_keys=['age','bmi','sex'], cycle_len=180, use_lstm=False, use_transformer=False,
                 simclr=False, simclr_temp=0.1,
                 dino=False, dino_student_temp=0.1, dino_teacher_temp=0.04,
                 dino_center_momentum=0.9, dino_momentum=0.996,
                 ):
        
        super().__init__()     
        # Consolidate simple config assignments to reduce repetition
        _cfg = dict(
            sig_len=sig_len,
            patch_size=patch_size,
            modalities=list(modalities),
            num_experts=num_experts,
            num_experts_per=num_experts_per,
            use_shared=use_shared,
            use_moe=use_moe,
            fusion_method=fusion_method,
            embed_dim=embed_dim,
            pos_enc=pos_enc,
            pretrain=pretrain,
            ft_task=ft_task,
            use_micro=use_micro,
            demo_keys=list(demo_keys),
            use_lstm=use_lstm,
            cycle_len=cycle_len,
            use_transformer=use_transformer,
            simclr=simclr,
            simclr_temp=simclr_temp,
            dino=dino,
            dino_student_temp=dino_student_temp,
            dino_teacher_temp=dino_teacher_temp,
            dino_center_momentum=dino_center_momentum,
            dino_momentum=dino_momentum,
        )
        for _k, _v in _cfg.items():
            setattr(self, _k, _v)
        
        # Build core modules
        self.encoder = MultiEncoderClassification(sig_len=sig_len, patch_size=patch_size,
                                          embed_dim=embed_dim, depth=depth, num_heads=num_heads, specific_depth=specific_depth,
                                          mlp_ratio=mlp_ratio, norm_layer=norm_layer,
                                          modalities=self.modalities, num_experts=num_experts, num_experts_per=num_experts_per,
                                          use_shared=use_shared, use_moe=use_moe, pos_enc=pos_enc)

        self.num_patches = self.encoder.num_patches

        self.fuser = FusionUpsampleBlock(dim_list=embed_dim[2:], 
                                   fusion_method=fusion_method, 
                                   fusion_depth=fusion_depth, 
                                   fusion_heads=fusion_heads,
                                   mlp_ratio=mlp_ratio, 
                                   norm_layer=norm_layer, 
                                   pos_enc=pos_enc, 
                                   use_shared=use_shared, 
                                   use_upsample=True, 
                                   modalities=self.modalities)

        # ReVin (optional)
        self.use_revin = use_revin
        if self.use_revin:
            self.revin = nn.ModuleDict()
            for m in self.modalities:
                if m != 'SaO2':
                    self.revin[m] = RevIN(num_features=1, affine=False)

        # Heads and sequence modules
        self._build_heads(embed_dim, num_classes, bidi, seq_dim, seq_depth, d_state, d_conv, expand,
                          use_micro, use_lstm, use_transformer, norm_layer)
                
        self.apply(self._init_weights)
        if self.pretrain and self.dino:
            self.teacher_encoder = None
            self.teacher_fuser = None
            self.teacher_proj = None
            self.teacher_sequence_encoder = None

        

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            torch.nn.init.xavier_uniform_(m.weight)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)

        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0) 

    # --- Helper builders -------------------------------------------------
    def _build_heads(self, embed_dim, num_classes, bidi, seq_dim, seq_depth, d_state, d_conv, expand,
                     use_micro, use_lstm, use_transformer, norm_layer=nn.LayerNorm):
        # micro head
        if use_micro:
            self.micro_head = nn.Linear(embed_dim[-1] * len(self.modalities), num_classes)
            return

        # sequence / macro heads
        self.bidi = bidi
        self.proj = nn.Linear(embed_dim[-1] * len(self.modalities), seq_dim)
        self.sequence_encoder = SeqEncoder(seq_dim=seq_dim, depth=seq_depth, d_state=d_state, d_conv=d_conv,
                                          expand=expand, bidi=bidi, use_transformer=use_transformer,
                                          norm_layer=norm_layer)

        out_dim = seq_dim * 2 if self.bidi else seq_dim
        out_dim += embed_dim[-1] * len(self.modalities)

        if self.pretrain:
            # Projection head for contrastive learning
            self.contr_head = nn.Sequential(
                nn.Linear(seq_dim, seq_dim),
                nn.ReLU(),
                nn.Linear(seq_dim, seq_dim)
            )
            self.koleo_weight = 0.01
            self.tau = 0.1
            self.upsilon = 0.5

        else:
            if use_lstm:
                if self.ft_task == 'diagnosis':
                    self.head = DiagnosisFinetuneFullLSTMCOXPHWithDemo(input_dim=out_dim, embed_dim=128, num_classes=num_classes)
                else:
                    self.head = SleepEventLSTMClassifier(input_dim=out_dim, embed_dim=128, num_classes=num_classes, num_heads=4)
            else:
                self.head = nn.Linear(out_dim, num_classes)

    

    def forward_features(self, inputs, mask=None):
        modalities_to_use = sorted(list(inputs.keys()))

        m = modalities_to_use[0]
        N, L, _  = inputs[m].shape
              
        
        if mask is not None:
            nonmask_idx = torch.nonzero(mask.view(-1,)).squeeze(-1)  # N*L_masked

        for i, m in enumerate(modalities_to_use):                            
            inputs[m] = inputs[m].view(-1, self.sig_len, inputs[m].shape[-1])  # add channel dim
            
            if self.use_revin:
                if m == 'SaO2':                    
                    inputs[m] = inputs[m][:, ::TARGET_RESAMPLE, :] # N,L* TARGET_RESAMPLE,1 -> N,L,1
                else:
                    inputs[m] = self.revin[m](inputs[m], mode='norm')
                
                if mask is not None:                    
                    inputs[m] = torch.gather(inputs[m], dim=0, index=nonmask_idx.unsqueeze(-1).unsqueeze(-1).repeat(1, inputs[m].shape[1], inputs[m].shape[2]))
        
        latents_shared, latents_specific, _ = self.encoder(inputs, modalities_to_use)
        
        latents_fused, upsampled = self.fuser(latents_shared, latents_specific)

        latents = torch.stack(latents_fused, dim=0) # (M, N, L, D)
        latents = rearrange(latents, 'm n l d -> n l (m d)')
        latents_micro = latents       
        
        # latents_micro = torch.stack(latents_fused, dim=0) # (M, N, L, D)
        # latents_micro = rearrange(latents_micro, 'm n l d -> n l (m d)')

        if not self.pretrain:
            latents_micro = latents_micro.mean(1)
            dummy2 = torch.zeros((N*(L // self.sig_len), latents_micro.shape[-1]), device=latents_micro.device, dtype=latents_micro.dtype)
            dummy2.scatter_(dim=0, index=nonmask_idx.unsqueeze(-1).repeat(1, latents_micro.shape[-1]), src=latents_micro)
            latents_micro = dummy2
            latents_micro = rearrange(latents_micro, '(n1 n2) d -> n1 n2 d', n1=N)

        # Finetuning with micro head only (ablation)
        if self.use_micro:
            return None, latents_micro
        
        else:        
            latents = self.proj(latents)
            latents = latents.mean(1)     
            
            dummy = torch.zeros((N*(L // self.sig_len), latents.shape[-1]), device=latents.device, dtype=latents.dtype)
            dummy.scatter_(dim=0, index=nonmask_idx.unsqueeze(-1).repeat(1, latents.shape[-1]), src=latents)
            latents = dummy
            latents = rearrange(latents, '(n1 n2) d -> n1 n2 d', n1=N)                     
            
            latents = self.sequence_encoder(latents, mask=mask)
            
            return latents, latents_micro
        
    def forward(self, sig, mask=None, demo=None):
        if self.pretrain and self.dino:
            return self._forward_dino(sig, mask)
            

        if self.pretrain and self.simclr:            
           return self._forward_simclr(sig, mask)

        feats = self.forward_features(sig, mask=mask)
        
        x_macro, x_micro = feats
        
        # Pretraining
        if self.pretrain:        
            if self.bidi:
                xf, xr = x_macro                
                loss_fwd = self.dgcl_loss(xf, demo, mask)
                loss_bwd = self.dgcl_loss(xr, demo, mask)
                loss = (loss_fwd + loss_bwd) / 2            
            else:                
                loss = self.dgcl_loss(x_macro, demo, mask)
            return loss
        
        # Inference / Finetuning
        if not self.use_micro and self.bidi:
            xf, xr = x_macro
            xr = flip_padded_sequences(xr, mask)
            x_macro = torch.cat([xf, xr], dim=-1)
            
        
        B = x_micro.shape[0]
        lengths = mask.sum(dim=1).long()  # (B,)        
        
        if self.use_micro:
            if self.ft_task in ['diagnosis', 'age', 'sex', 'ahi']:
                x_micro = torch.stack([x_micro[i, :lengths[i]].mean(dim=0) for i in range(B)])
            x = self.micro_head(x_micro)
        else:
            x_mami = torch.cat([x_macro, x_micro], dim=-1)
            if self.use_lstm:
                x = self.head(x_mami, mask, demo)                                
            else:
                if self.ft_task in ['diagnosis', 'age', 'sex', 'ahi']:
                    x_mami = torch.stack([x_mami[i, :lengths[i]].mean(dim=0) for i in range(B)])
                x = self.head(x_mami)
        return x
            
        
        
        # if self.ft_task in ['diagnosis', 'age' ,'sex', 'ahi']:            
        #     B = x_micro.shape[0]
        #     lengths = mask.sum(dim=1).long()  # (B,)
        #     if self.use_micro:                
        #         x_micro = torch.stack([x_micro[i, :lengths[i]].mean(dim=0) for i in range(B)])                
        #         x = self.micro_head(x_micro)
        #     else:
        #         x_mami = torch.cat([x_macro, x_micro], dim=-1)
        #         if self.use_lstm:
        #             x = self.head(x_mami, mask, demo)                                
        #         else:
        #             x_mami = torch.stack([x_mami[i, :lengths[i]].mean(dim=0) for i in range(B)])
        #             x = self.head(x_mami)
        #     return x
        # else:            
        #     if self.use_micro:
        #         x = self.micro_head(x_micro)
        #     else:
        #         x_mami = torch.cat([x_macro, x_micro], dim=-1)
        #         # x = self.head(x_mami)            
        #         if self.use_lstm:
        #             x = self.head(x_mami, mask)            
        #         else:
        #             x = self.head(x_mami)
        #     x = rearrange(x, 'n1 n2 d -> n1 d n2')
        #     return x

    def dgcl_loss(self, xf, demo, mask):
        """
        Calculate DGCL loss based on the difference of age.
        When the corresponding ages between two sequences are distant, the representation should be distant
        In constrast, if ages are close, the representation should be similar.
        """
        
        N, L, D = xf.shape
        num_cycle = L // self.cycle_len
        
        age = demo['age'] 
        sex = demo['sex'] 
        bmi = demo['bmi']

        mu_age, std_age = 56.071731, 16.131597
        mu_bmi, std_bmi = 27.095034, 4.943150
        age = (age - mu_age) / std_age # Normalize
        
        age_t = age.unsqueeze(-1)
        age = age.unsqueeze(0)
        distance_age = (age - age_t).abs()
        
        distance = torch.zeros_like(distance_age)

        if 'age' in self.demo_keys:
            distance += distance_age

        if 'bmi' in self.demo_keys:
            # BMI adjustment
            bmi = (bmi - mu_bmi) / std_bmi
            bmi_t = bmi.unsqueeze(-1)
            bmi = bmi.unsqueeze(0)        
            distance_bmi = (bmi - bmi_t).abs()
            distance += distance_bmi
        
        if 'age' in self.demo_keys and 'bmi' in self.demo_keys:
            distance = distance / 2

        if 'sex' in self.demo_keys:
            diff_sex_matrix = 1.0 - (sex.unsqueeze(0) == sex.unsqueeze(1)).float() # Same sex: 0, different sex: 1                   
            distance += diff_sex_matrix
        
        
        distance += torch.eye(N, device=age.device) * 1e8
        
        mask_seg = mask.view(N, -1, self.cycle_len).any(-1)  # (N, num_seg)
        distance = distance.unsqueeze(0).repeat(num_cycle,1,1)  # (num_cycle, N, N) Same for all cycles
        mask1 = mask_seg.permute(1,0)  # (num_cycle, N)        
        label = F.softmax((- distance / self.upsilon), dim=-1) # (num_cycle, N, N)
        
        if mask is not None: 
            xf = xf.masked_fill(~mask.unsqueeze(-1), 0.0)                     
            
        min_length = mask.sum(dim=1).min()
        num_seg = min_length // self.cycle_len # number of cycles that can be fully formed for all sequences in the batch

        if self.use_transformer:
            cycle_mode = 'mean'
        else:
            cycle_mode = 'last'
        seq1 = self._get_cycle_values(xf, mask, mode=cycle_mode)  # (N, num_seg + 1, D)

        # if self.use_transformer:           
        #     seq1 = xf.view(N, -1, self.cycle_len, D).mean(2)  # (N, num_seg, D)
        #     seq1 = self.contr_head(seq1)
        #     seq1 = F.normalize(seq1, p=2, dim=-1)  #
        # else:
        #     # Last feature
        #     if not min_length == L:
        #         mask_len = mask.sum(dim=1)
        #         seq1_last = xf[torch.arange(N, device=xf.device), mask_len - 1]
        #         seq1_last = self.contr_head(seq1_last)
        #         seq1_last = F.normalize(seq1_last, p=2, dim=-1)    

        #     if num_seg < 1:            
        #         seq1 = seq1_last.unsqueeze(1)
        #     else:
        #         min_length = num_seg * self.cycle_len

        #         seq_trunc = xf[:, :min_length, :]
        #         seq1 = seq_trunc[:,(self.cycle_len-1)::self.cycle_len]
        #         B_temp, S_temp, D_temp = seq1.shape
        #         seq1 = seq1.reshape(-1, D_temp)  # flatten batch and segments
        #         seq1 = self.contr_head(seq1)
        #         seq1 = seq1.reshape(B_temp, S_temp, -1)  # reshape back
        #         seq1 = F.normalize(seq1, p=2, dim=-1)  # (N, num_seg, D)
            
        #         # Last feature
        #         if not min_length == L:                
        #             seq1 = torch.cat([seq1, seq1_last.unsqueeze(1)], dim=1)  # (N, num_seg+1, D)
        
        
        koleo_loss = 0.0
        for i in range(seq1.shape[1]):            
            koleo_loss += self.koleo_loss(seq1[i])        
        
        koleo_loss = koleo_loss / seq1.shape[1]
        # Patient by patient similarity
        seq1 = seq1.permute(1,0,2).contiguous()  # (num_seg + 1, N, D)        
        sims_pos = seq1 @ seq1.permute(0,2,1)  - torch.eye(N, device=age.device).unsqueeze(0) * 1e8 # (num_seg + 1, N, N)                
        sims = sims_pos / self.tau  # (num_seg + 1, N, N)   
        # print("## Similarity", sims[0,0])     
        
        loss1, cnt = 0, 0
        # For each segment, calculate the loss only among valid sequences 
        # (with enough length) to avoid noise from padding
        for i in range(num_seg):            
            # Skip any segments with too few valid sequences.
            if mask1[i].sum() <= 4:
                continue            
            sims1 = sims[i][mask1[i]].masked_fill(~mask1[i], -1e8)
            label1 = label[i][mask1[i]].masked_fill(~mask1[i], 0.0)
            loss1 += nn.CrossEntropyLoss(reduction='mean')(sims1, label1)
            cnt += 1
        
        # Additionally, calculate loss for the last segment.
        # num_seg is the number of full segments.
        if not self.use_transformer and not min_length == L:
            loss1 += nn.CrossEntropyLoss(reduction='mean')(sims[-1], label[0])        
            cnt += 1
            
        loss1 = loss1 / cnt if cnt > 0 else loss1        
        loss1 = loss1 + self.koleo_weight * koleo_loss        
         
        return loss1
    
    # def init_teacher(self):
    #     raise RuntimeError("init_teacher was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.")
    

    # def _forward_dino(self, sig, mask):
    #     raise RuntimeError("_forward_dino was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.")
    
    # def _forward_simclr(self, sig, mask):
    #     raise RuntimeError("_forward_simclr was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use SimCLR.")
    
    # def _simclr_representation(self, sig, mask):
    #     raise RuntimeError("_simclr_representation was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use SimCLR.")
        
    def _get_cycle_values(self, xf, mask, mode='last'):
        """
        Get the last/mean features of each cycle.
        "mode" is either 'last' or 'mean'. 'last' is used in the paper (Mamba), 
        but 'mean' can be an alternative (Transformer)
        """
        if mask is not None: 
            xf = xf.masked_fill(~mask.unsqueeze(-1), 0.0)                     
            
        min_length = mask.sum(dim=1).min()
        num_seg = min_length // self.cycle_len
        N, L, D = xf.shape

        if mode == 'mean':
            seq1 = xf.view(N, -1, self.cycle_len, D).mean(2)  # (N, num_seg, D)
            seq1 = self.contr_head(seq1)
            seq1 = F.normalize(seq1, p=2, dim=-1) 
            return seq1
        
        # Last feature
        if not min_length == L:
            mask_len = mask.sum(dim=1)
            seq1_last = xf[torch.arange(N, device=xf.device), mask_len - 1]
            seq1_last = self.contr_head(seq1_last)
            seq1_last = F.normalize(seq1_last, p=2, dim=-1)          

        if num_seg < 1:            
            seq1 = seq1_last.unsqueeze(1)
        else:
            min_length = num_seg * self.cycle_len

            seq_trunc = xf[:, :min_length, :]
            seq1 = seq_trunc[:,(self.cycle_len-1)::self.cycle_len]
            B_temp, S_temp, D_temp = seq1.shape
            seq1 = seq1.reshape(-1, D_temp)  # flatten batch and segments
            seq1 = self.contr_head(seq1)
            seq1 = seq1.reshape(B_temp, S_temp, -1)  # reshape back
            seq1 = F.normalize(seq1, p=2, dim=-1)  # (N, num_seg, D)
        
            # Last feature
            if not min_length == L:                
                seq1 = torch.cat([seq1, seq1_last.unsqueeze(1)], dim=1)  # (N, num_seg+1, D)
        return seq1

    # def _masked_mean(self, x, mask):
    #     if mask is None:
    #         return x.mean(1)
    #     mask = mask.to(dtype=x.dtype)
    #     denom = mask.sum(dim=1, keepdim=True).clamp(min=1.0)
    #     return (x * mask.unsqueeze(-1)).sum(dim=1) / denom

    # def _dino_representation(self, sig, mask, use_teacher=False):
    #     raise RuntimeError("_dino_representation was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.")


    # def _forward_features_with_modules(self, inputs, mask, encoder, fuser, proj, sequence_encoder):
    #     raise RuntimeError("_forward_features_with_modules was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO teacher forward.")

    # def dino_loss(self, student_out, teacher_out, fwd=False):
    #     raise RuntimeError("dino_loss was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.")

    # @torch.no_grad()
    # def _update_dino_center(self, teacher_out_f, teacher_out_r):
    #     raise RuntimeError("_update_dino_center was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.")

    # @torch.no_grad()
    # def update_teacher(self, momentum=None):
    #     raise RuntimeError("update_teacher was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.")

    # def _freeze_teacher(self):
    #     raise RuntimeError("_freeze_teacher was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use DINO.")


    # def simclr_loss(self, z1, z2):
    #     raise RuntimeError("simclr_loss was moved to PretrainingWrapper. Wrap the backbone with PretrainingWrapper to use SimCLR.")


    def koleo_loss(self, x, ys=None, eps=1e-8):
        """
        Compute KoLeo (Kozachenko-Leonenko) regularization loss.
        Encourages uniform distribution of representations on the hypersphere.
        
        Args:
            x: Tensor of shape (N, D) where N is batch size and D is feature dimension
            eps: Small constant for numerical stability
        
        Returns:
            KoLeo loss (scalar)
        """
        # Normalize representations to unit hypersphere
        x = F.normalize(x, p=2, dim=-1)
        
        # Compute pairwise distances
        # Using efficient computation: ||a - b||^2 = ||a||^2 + ||b||^2 - 2<a,b>
        # Since normalized, ||a||^2 = ||b||^2 = 1, so ||a - b||^2 = 2 - 2<a,b>
        gram_matrix = torch.matmul(x, x.t())  # (N, N)
        sq_distances = 2 - 2 * gram_matrix  # (N, N)
        sq_distances = torch.clamp(sq_distances, min=0)  # numerical stability
        
        # Set diagonal to infinity (exclude self-distance)
        sq_distances = sq_distances + torch.eye(x.shape[0], device=x.device) * 1e10

        if ys is not None:
            for y in ys:
                y = F.normalize(y, p=2, dim=-1)
                y_distance = (x - y).pow(2).sum(dim=-1, keepdim=True)  # (N, 1)
                sq_distances = torch.cat([sq_distances, y_distance], dim=1)  # (N, N + len(ys))
        
        # Find minimum distance for each sample (nearest neighbor)
        min_sq_distances, _ = torch.min(sq_distances, dim=1)  # (N,)
        
        # KoLeo loss is negative log of minimum distances
        # We want to maximize minimum distances -> minimize negative log
        distances = torch.sqrt(min_sq_distances + eps)
        koleo_loss = -torch.log(distances + eps).mean()
        
        return koleo_loss

    
class SleepEventLSTMClassifier(nn.Module):
    def __init__(self, 
                 input_dim=384*6+512*2,
                 embed_dim=128, 
                 num_heads=4, 
                 num_layers=1, 
                 num_classes=5, 
                 pooling_head=4, 
                 dropout=0.3, 
                 max_seq_length=128):
        super().__init__()
        
        # Define spatial pooling
        # self.spatial_pooling = AttentionPooling(embed_dim, num_heads=pooling_head, dropout=dropout)

        # Set max sequence length
        # if max_seq_length is None:
        #     max_seq_length = 20000
            
        # self.positional_encoding = PositionalEncoding(max_seq_length, embed_dim)
        self.layer_norm = nn.LayerNorm(input_dim)
        self.proj = nn.Linear(input_dim, embed_dim)

        # Transformer encoder for spatial modeling
        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dropout=dropout, batch_first=True, norm_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        # LSTM for temporal modeling
        lstm_dropout = dropout if num_layers > 1 else 0.0
        self.lstm = nn.LSTM(input_size=embed_dim, hidden_size=embed_dim//2, num_layers=num_layers, batch_first=True, dropout=lstm_dropout, bidirectional=True)
        
        # Fully connected layer for classification
        self.fc = nn.Linear(embed_dim, num_classes)

    def forward(self, x, mask, demo=None):
        # B, C, S, E = x.shape
        
        # Rearrange for spatial pooling
        # x = rearrange(x, 'b c s e -> (b s) c e')
        
        # # Prepare the mask for spatial pooling
        # mask_spatial = mask[:, :, 0]
        # mask_spatial = mask_spatial.unsqueeze(1).expand(-1, S, -1)
        # mask_spatial = rearrange(mask_spatial, 'b t c -> (b t) c')
        
        # # Ensure the mask is boolean
        # if mask_spatial.dtype != torch.bool:
        #     mask_spatial = mask_spatial.to(dtype=torch.bool)

        # # Apply spatial pooling
        # x = self.spatial_pooling(x, mask_spatial)
        
        # Reshape to (B, S, E) after pooling
        # x = x.view(B, S, E)

        # Apply positional encoding and layer normalization
        # x = self.positional_encoding(x)
        x = self.layer_norm(x)
        x = self.proj(x)

        # Apply transformer encoder for spatial modeling
        # mask_temporal = mask[:, 0, :]
        mask_temporal = ~mask # 1 should be padding, 0 should be valid values
        x = self.transformer_encoder(x, src_key_padding_mask=mask_temporal)

        # Apply LSTM for temporal modeling
        x, _ = self.lstm(x)  # Shape: (B, S, E)

        # Apply the final fully connected layer for classification
        x = self.fc(x)  # Shape: (B, S, num_classes)

        return x  # Return mask along temporal dimension



# From SleepFM-clinical
import torch.nn.utils.rnn as rnn_utils
class DiagnosisFinetuneFullLSTMCOXPHWithDemo(nn.Module):
    def __init__(self, input_dim, embed_dim=128, num_layers=2, num_classes=6, dropout=0.1):
        super(DiagnosisFinetuneFullLSTMCOXPHWithDemo, self).__init__()
        
        self.layer_norm = nn.LayerNorm(input_dim)
        self.proj = nn.Linear(input_dim, embed_dim)        
        self.lstm = nn.LSTM(
            input_size=embed_dim, 
            hidden_size=embed_dim // 2, 
            num_layers=num_layers, 
            batch_first=True, 
            bidirectional=True
        )
        self.demo_embedding = nn.Sequential(
            nn.Linear(2, embed_dim // 4),
            nn.ReLU(),
            nn.Dropout(dropout)
        )

        # Update the input size of the disease head to account for embedded demo features
        self.head = nn.Linear(embed_dim + embed_dim // 4, num_classes)  # embed_dim // 4 from demo features

    def forward(self, x, mask, demo):
        # B, C, S, E = x.shape # C: Modalities
        
        # # Rearrange for spatial pooling
        # x = rearrange(x, 'b c s e -> (b s) c e')

        # mask_spatial = mask[:, :, 0]
        # mask_spatial = mask_spatial.unsqueeze(1).expand(-1, S, -1)
        # mask_spatial = rearrange(mask_spatial, 'b t c -> (b t) c')

        # # Ensure the mask is boolean
        # mask_spatial = mask_spatial.to(dtype=torch.bool)

        # # Apply spatial pooling
        # x = self.spatial_pooling(x, mask_spatial)
        # x = x.view(B, S, E)

        # # Rearrange mask for temporal pooling
        # # mask_temporal = mask[:, 0, :]  # Use the mask along the temporal dimension

        # # Compute the lengths for packing based on the mask
        # lengths = (mask_temporal == 0).sum(dim=1).cpu()


        x = self.layer_norm(x)
        x = self.proj(x)
        # lengths = (mask == 1).sum(dim=1).cpu()  # (B,)
        lengths = mask.sum(dim=1).long().cpu()  # (B,)
        B = x.shape[0]
        # Pack the sequence for LSTM
        packed_x = rnn_utils.pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False)

        # Pass the packed sequence through the LSTM
        packed_out, _ = self.lstm(packed_x)

        # Unpack the sequence
        x, _ = rnn_utils.pad_packed_sequence(packed_out, batch_first=True)

        # Apply mean pooling only over valid lengths
        x = torch.stack([x[i, :lengths[i]].mean(dim=0) for i in range(B)])
        
        age = (demo['age'] - 63.14)/11.22
        age = torch.clamp(age, -3.0, 3.0)
        
        sex = demo['sex'] - 1 # 0 or 1

        demo_features = torch.cat([age.unsqueeze(1), sex.unsqueeze(1)], dim=1)
        demo_features = demo_features.to(x.device, dtype=x.dtype)
        # Embed the demo features
        demo_embed = self.demo_embedding(demo_features)  # (B, embed_dim // 4)

        # Concatenate the demo features with the LSTM output
        x = torch.cat([x, demo_embed], dim=1)  # (B, embed_dim + embed_dim // 4)

        # Apply the disease head to get hazards
        hazards = self.head(x)
        
        return hazards


class DiagnosisFinetuneDemoOnly(nn.Module):
    def __init__(self, embed_dim=128, num_layers=2, num_classes=6, dropout=0.1):
        super().__init__()        
        
        self.demo_embedding = nn.Sequential(
            nn.Linear(2, embed_dim // 4),
            nn.ReLU(),
            nn.Dropout(dropout)
        )

        # Update the input size of the disease head to account for embedded demo features
        self.head = nn.Linear(embed_dim // 4, num_classes)  # embed_dim // 4 from demo features

    def forward(self, x, mask, demo):        
                
        age = (demo['age'] - 63.14)/11.22
        age = torch.clamp(age, -3.0, 3.0)
        
        sex = demo['sex'] - 1 # 0 or 1

        demo_features = torch.cat([age.unsqueeze(1), sex.unsqueeze(1)], dim=1)
        demo_features = demo_features.to(x.device, dtype=x.dtype)
        # Embed the demo features
        demo_embed = self.demo_embedding(demo_features)  # (B, embed_dim // 4)
        
        # Apply the disease head to get hazards
        hazards = self.head(demo_embed)
        
        return hazards