import torch.nn as nn
from basicsr.archs.mair_arch import MaIR

def buildMaIR_Small(upscale=2):
    return MaIR(img_size=64,
                   patch_size=1,
                   in_chans=3,
                   embed_dim=60,
                   depths=(6, 6, 6, 6),
                   mlp_ratio=1.6,
                   ssm_ratio=1.4,
                   drop_rate=0.,
                   norm_layer=nn.LayerNorm,
                   patch_norm=True,
                   use_checkpoint=False,
                   upscale=upscale,
                   img_range=1.,
                   upsampler='pixelshuffledirect',
                   resi_connection='1conv')

def buildMaIR_Tiny(upscale=2):
    return MaIR(img_size=64,
                   patch_size=1,
                   in_chans=3,
                   embed_dim=60,
                   depths=(6, 6, 6, 6),
                   mlp_ratio=1.6,
                   ssm_ratio=1.1,
                   d_state=1,
                   drop_rate=0.,
                   norm_layer=nn.LayerNorm,
                   patch_norm=True,
                   use_checkpoint=False,
                   upscale=upscale,
                   img_range=1.,
                   upsampler='pixelshuffledirect',
                   resi_connection='1conv')


def buildMaIR(upscale=2):
    return MaIR(img_size=64,
                   patch_size=1,
                   in_chans=3,
                   embed_dim=180,
                   depths=(6, 6, 6, 6, 6, 6),
                   mlp_ratio=2.,
                   drop_rate=0.,
                   norm_layer=nn.LayerNorm,
                   patch_norm=True,
                   use_checkpoint=False,
                   upscale=upscale,
                   img_range=1.,
                   upsampler='pixelshuffle',
                   resi_connection='1conv')

def buildMaIR_SR(upscale=2):
    return MaIR(img_size=64,
                 patch_size=1,
                 in_chans=3,
                 embed_dim=180,
                 depths=(6, 6, 6, 6, 6, 6),
                 drop_rate=0.,
                 d_state=16,
                 ssm_ratio=2.0,
                 mlp_ratio=2.5,

                 drop_path_rate=0.1,
                 norm_layer=nn.LayerNorm,
                 patch_norm=True,
                 use_checkpoint=False,
                 upscale=upscale,
                 img_range=1.,
                 upsampler='pixelshuffle',
                 resi_connection='1conv',
                 dynamic_ids=False,
                 scan_len=4,
                 batch_size=1,
                )