from functools import partial
import torch.nn as nn

from models.mae import MultiMaskedAutoencoderConv1d
from models.classification import ClassificationMultiConv1d, ClassificationMultiConv1dMamba
from models.transformer_modules import RMSNorm

def convmae1d_multi(**kwargs):
    model = MultiMaskedAutoencoderConv1d(        
        embed_dim=[32, 64, 128, 384], depth=[4, 4], specific_depth=[2, 2], num_heads=8,
        decoder_embed_dim=64, decoder_depth=2, decoder_num_heads=4,
        fusion_method='cross_attn', fusion_heads=8, fusion_depth=4,        
        mlp_ratio=4, norm_layer=partial(RMSNorm, eps=1e-8), pos_enc='rope', **kwargs)
    return model

def classification_multi(**kwargs):
    model = ClassificationMultiConv1d(                
        embed_dim=[32, 64, 128, 384], depth=[4, 4], specific_depth=[2, 2], num_heads=8,
        fusion_method='cross_attn', fusion_heads=8, fusion_depth=4,        
        mlp_ratio=4, norm_layer=partial(RMSNorm, eps=1e-8), pos_enc='rope', **kwargs)
    return model

def classification_multi_full_night(**kwargs):
    model = ClassificationMultiConv1dMamba(                
        embed_dim=[32, 64, 128, 384], depth=[4, 4], specific_depth=[2, 2], num_heads=8,
        fusion_method='cross_attn', fusion_heads=8, fusion_depth=4,
        seq_dim=512, seq_depth=1, 
        mlp_ratio=4, norm_layer=partial(RMSNorm, eps=1e-8), pos_enc='rope', **kwargs)
    return model