# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.
import torch
import torch.nn as nn
import torch.nn.functional as functional
from functools import partial
from torch import Tensor
from typing import Optional

from timm.models.vision_transformer import VisionTransformer, _cfg
# from timm.models.registry import register_model
# from timm.models.layers import trunc_normal_, lecun_normal_

from timm.models import register_model
from timm.layers import trunc_normal_, lecun_normal_

from timm.layers import DropPath, to_2tuple

# from timm.models.layers import DropPath, to_2tuple
from timm.models.vision_transformer import _load_weights

import math

from collections import namedtuple

from mamba_ssm.modules.mamba_simple import Mamba
from mamba_ssm.utils.generation import GenerationMixin
from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf

from rope import *
import random
import sys

try:
    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
except ImportError:
    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None


# layer_norm_fn and rms_norm_fn both are normalization method

__all__ = [
    'vim_tiny_patch16_224', 'vim_small_patch16_224', 'vim_base_patch16_224',
    'vim_tiny_patch16_384', 'vim_small_patch16_384', 'vim_base_patch16_384',
]


'''
in the original script(ft-vim-s.sh)
img_size = 224, 
patch_size = 16, 
stride = 8, 
in_chans = 3, 
embed_dim = 768
------------------------------------
self.img_size: (224, 224)
self.patch_size: (16, 16)
self.grid_size: (27, 27)
self.num_patches: 729
self.flatten: True
self.norm: nn.Identity()
'''
class PatchEmbed(nn.Module):
    """ 2D Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, stride=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = ((img_size[0] - patch_size[0]) // stride + 1, (img_size[1] - patch_size[1]) // stride + 1)
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.flatten = flatten

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride)
        # if the norm_layer is not none or null, the self.norm = norm_layer(embed_dim)
        # otherwise, self.norm = nn.Identity()
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
        

    def forward(self, x):
        B, C, H, W = x.shape
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        
        x = self.proj(x)
        print("This is the shape after the CNN", x.shape)

        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
            print("This is the shape after the flatten:", x.shape)

        x = self.norm(x)
        return x

class PatchEmbed_spectrogram(nn.Module):
    """ 2D spectrogram to Patch Embedding
    """
    def __init__(self, img_size_f = 128, img_size_t = 64, patch_size=6, stride=3, in_chans=12, embed_dim=432, flatten=True):
        super().__init__()
        # img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)

        self.img_size_f = img_size_f
        self.img_size_t = img_size_t

        self.patch_size = patch_size
        self.grid_size = ((img_size_f - patch_size[0]) // stride + 1, (img_size_t - patch_size[1]) // stride + 1)
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.flatten = flatten

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride)
        # if the norm_layer is not none or null, the self.norm = norm_layer(embed_dim)
        # otherwise, self.norm = nn.Identity()
        # self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
        

    def forward(self, x):
        B, C, H, W = x.shape
        assert H == self.img_size_f and W == self.img_size_t, \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size_f}*{self.img_size_t})."
        
        x = self.proj(x)

        # This is the shape after the CNN torch.Size([1, 432, 41, 20])
        # print("This is the shape after the CNN", x.shape)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        # This is the shape after the flatten: torch.Size([1, 820, 432])
        # print("This is the shape after the flatten:", x.shape)
        return x


class CNN_layers(nn.Module):

    def __init__(self, embed_size = 384):
        super().__init__()
        self.multiple_cnn = nn.Sequential( 
            nn.Conv1d(12, 128, kernel_size=14, stride=3, padding=2, bias=False),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),

            nn.Conv1d(128, embed_size, kernel_size=15, stride=4, padding=100, bias=False),
            nn.BatchNorm1d(embed_size),
            nn.ReLU(inplace=True))
        
    def forward(self, x):
        # print("This is the shape of enconder:(before)", x.shape)
        # This is the shape of enconder:(before) torch.Size([44, 12, 8192])
        x = self.multiple_cnn(x)
        x = x.transpose(1, 2)
        # print("This is the shape of enconder:(after)", x.shape)
        # This is the shape of enconder:(after) torch.Size([44, 729, 384])

        # sys.exit()
        return x 

class CNN_layers_shortcut(nn.Module):

    def __init__(self, embed_size = 384):
        super().__init__()
        self.multiple_cnn1 = nn.Sequential( 
            nn.Conv1d(12, 128, kernel_size=14, stride=3, padding=2, bias=False),
            nn.BatchNorm1d(128),
            nn.ReLU(inplace=True),

            nn.Conv1d(128, embed_size, kernel_size=15, stride=4, padding=100, bias=False),
            nn.BatchNorm1d(embed_size),
            nn.ReLU(inplace=True))

        self.multiple_cnn2 = nn.Sequential(
            nn.Conv1d(in_channels=embed_size, out_channels=embed_size, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm1d(embed_size),
            nn.ReLU(inplace=True),

            nn.Conv1d(in_channels=embed_size, out_channels=embed_size, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm1d(embed_size),
            nn.ReLU(inplace=True)

            )
        
    def forward(self, x):
        # print("This is the shape of enconder:(before)", x.shape)
        # This is the shape of enconder:(before) torch.Size([44, 12, 8192])
        x = self.multiple_cnn1(x)

        shortcut = x

        x = self.multiple_cnn2(x)
        x = x + shortcut

        x = x.transpose(1, 2)
        # print("This is the shape of enconder:(after)", x.shape)
        # This is the shape of enconder:(after) torch.Size([44, 729, 384])
        # sys.exit()
        return x 

# class ECG_Patch_embedding(nn.Module):

#     def __init__(self, embed_size = 384):
#         super().__init__()
#         self.multiple_cnn = nn.Sequential( 
#             nn.Conv1d(12, embed_size, kernel_size=16, stride=8),
#             )
        
#     def forward(self, x):
#         # print("This is the shape of enconder:(before)", x.shape)
#         # This is the shape of enconder:(before) torch.Size([44, 12, 8192])
#         x = self.multiple_cnn(x)
#         x = x.transpose(1, 2)
#         # print("This is the shape of enconder:(after)", x.shape)
#         # This is the shape of enconder:(after) torch.Size([44, 1023, 384])

#         # sys.exit()
#         return x 

# class LeadCombiner(nn.Module):
#     def __init__(self, lead, out_ch):
#         super(LeadCombiner, self).__init__()
#         self.conv2_1 = nn.Conv1d(in_channels=lead * out_ch,
#                                  out_channels=out_ch,
#                                  kernel_size=1,
#                                  bias=False)
#         self.bn2_1 = nn.BatchNorm1d(out_ch)

#         self.conv2_2 = nn.Conv1d(in_channels=lead * out_ch,
#                                  out_channels=out_ch,
#                                  kernel_size=1,
#                                  bias=False)
#         self.bn2_2 = nn.BatchNorm1d(out_ch)

#         self.pool1 = nn.AdaptiveMaxPool1d(output_size=1)
#         self.pool2 = nn.AdaptiveMaxPool1d(output_size=1)

#     def forward(self, x):
#         # this is the shape of x: torch.Size([78, 128, 12, 128]
#         x1 = rearrange(x, 'b c l t -> b (c l) t')
#         x1 = functional.leaky_relu(self.bn2_1(self.conv2_1(x1)))

#         x2 = rearrange(x, 'b c l t -> b (t l) c')
#         x2 = functional.leaky_relu(self.bn2_2(self.conv2_2(x2)))

#         x1 = functional.dropout(x1, p=0.5, training=self.training)
#         x2 = functional.dropout(x2, p=0.5, training=self.training)

#         x1 = self.pool1(x1).squeeze(2)
#         x2 = self.pool2(x2).squeeze(2)

#         x = torch.cat([x1, x2], dim=1)
#         return x


# '''
# changed the stride and added one more layer
# '''
# class CNN_layers(nn.Module):

#     def __init__(self, embed_size = 384):
#         super().__init__()
#         self.multiple_cnn = nn.Sequential( 
#             nn.Conv1d(12, 128, kernel_size=15, stride=2, padding=2, bias=False),
#             nn.BatchNorm1d(128),
#             nn.ReLU(inplace=True),

#             nn.Conv1d(128, embed_size, kernel_size=15, stride=2, padding=100, bias=False),
#             nn.BatchNorm1d(embed_size),
#             nn.ReLU(inplace=True),
            
#             nn.Conv1d(384, embed_size, kernel_size=15, stride=3, padding=100, bias=False),
#             nn.BatchNorm1d(embed_size),
#             nn.ReLU(inplace=True)
#             )
        
#     def forward(self, x):
#         # print("This is the shape of enconder:(before)", x.shape)
#         # This is the shape of enconder:(before) torch.Size([44, 12, 8192])
#         x = self.multiple_cnn(x)
#         x = x.transpose(1, 2)
        

#         return x

# class D2_CNN_layers(nn.Module):
#     def __init__(self, embed_size = 384):
#         super().__init__()
#         self.multiple_cnn = nn.Sequential(
#             nn.Conv2d(in_channels=1, out_channels=128, kernel_size=(3, 15), padding=(1, 7), stride=(1, 2), bias=False),
#             nn.BatchNorm2d(128),
#             nn.LeakyReLU(inplace=True),

#             nn.Conv2d(in_channels=128, out_channels=384, kernel_size=(3, 15), padding=(1, 7), stride=(1, 2), bias=False),
#             nn.BatchNorm2d(embed_size),
#             nn.LeakyReLU(inplace=True),

#             nn.Conv2d(in_channels=384, out_channels=384, kernel_size=(3, 15), padding=(1, 7), stride=(1, 30), bias=False),
#             nn.BatchNorm2d(embed_size),
#             nn.LeakyReLU(inplace=True),

#         )

#     def forward(self, x):
#         x = x.unsqueeze(1)
#         # print("This is the input shape:", x.shape)
#         x = self.multiple_cnn(x)
#         x = torch.flatten(x, start_dim=-2)
#         x = x.transpose(1, 2)
#         return x

def broadcat(tensors, dim=-1):
    num_tensors = len(tensors)
    shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
    assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions'
    shape_len = list(shape_lens)[0]
    dim = (dim + shape_len) if dim < 0 else dim
    dims = list(zip(*map(lambda t: list(t.shape), tensors)))
    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
    assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), 'invalid dimensions for broadcastable concatenation'
    max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
    expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
    expanded_dims.insert(dim, (dim, dims[dim]))
    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
    tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
    return torch.cat(tensors, dim=dim)

def rotate_half(x):
    x = rearrange(x, '... (d r) -> ... d r', r=2)
    x1, x2 = x.unbind(dim=-1)
    x = torch.stack((-x2, x1), dim=-1)
    return rearrange(x, '... d r -> ... (d r)')

# Adapted Rotary Embedding for 1D time-series (ECG) with 1024 tokens
class TimeSeriesRotaryEmbeddingFast(nn.Module):
    def __init__(
        self,
        dim,
        seq_len=1024,  # Updated to 1024 tokens
        custom_freqs=None,
        freqs_for='lang',  # Suitable for 1D sequential data
        theta=10000,
        max_freq=10,
        num_freqs=1,
    ):
        super().__init__()
        if custom_freqs:
            freqs = custom_freqs
        elif freqs_for == 'lang':
            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
        elif freqs_for == 'pixel':
            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
        elif freqs_for == 'constant':
            freqs = torch.ones(num_freqs).float()
        else:
            raise ValueError(f'unknown modality {freqs_for}')

        # 1D sequence for 1024 tokens
        t = torch.arange(seq_len).float()  # [0, 1, ..., 1023]

        freqs = torch.einsum('..., f -> ... f', t, freqs)
        freqs = repeat(freqs, '... n -> ... (n r)', r=2)  # Doubles the dimension for rotation

        # Shape: (1024, dim)
        freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
        freqs_sin = freqs.sin().view(-1, freqs.shape[-1])

        self.register_buffer("freqs_cos", freqs_cos)  # Shape: (1024, dim)
        self.register_buffer("freqs_sin", freqs_sin)  # Shape: (1024, dim)

        print('======== shape of rope freq', self.freqs_cos.shape, '========')

    def forward(self, t):
        # t: (batch_size, seq_len, embed_dim)
        # Apply rotation to the entire sequence
        return t * self.freqs_cos + rotate_half(t) * self.freqs_sin



'''
dim: 384
mixer_cls: mixer_cla is an instance of mamba.
drop_path = 0.
norm_cls = nn.LayerNorm
fused_add_norm = True, 
residual_in_fp32 = True
'''
class Block(nn.Module):
    def __init__(
        self, dim, mixer_cls, norm_cls=nn.LayerNorm, fused_add_norm=False, residual_in_fp32=False, drop_path=0.,
    ):
        """
        Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection"

        This Block has a slightly different structure compared to a regular
        prenorm Transformer block.
        The standard block is: LN -> MHA/MLP -> Add.
        [Ref: https://arxiv.org/abs/2002.04745]
        Here we have: Add -> LN -> Mixer, returning both
        the hidden_states (output of the mixer) and the residual.
        This is purely for performance reasons, as we can fuse add and LayerNorm.
        The residual needs to be provided (except for the very first block).
        """
        super().__init__()
        
        self.residual_in_fp32 = residual_in_fp32
        self.fused_add_norm = fused_add_norm
        self.mixer = mixer_cls(dim)
        self.norm = norm_cls(dim)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        
        # fused_add_norm true
        if self.fused_add_norm:
            assert RMSNorm is not None, "RMSNorm import fails"
            assert isinstance(self.norm, (nn.LayerNorm, RMSNorm)), "Only LayerNorm and RMSNorm are supported for fused_add_norm"

    def forward(self, hidden_states: Tensor, residual: Optional[Tensor] = None, inference_params=None):

        r"""Pass the input through the encoder layer.

        Args:
            hidden_states: the sequence to the encoder layer (required).
            residual: hidden_states = Mixer(LN(residual))
        """
        
        if not self.fused_add_norm:
            if residual is None:
                residual = hidden_states
            else:
                residual = residual + self.drop_path(hidden_states)
            
            hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
            if self.residual_in_fp32:
                residual = residual.to(torch.float32)
        
        # since the self.fused_add_norm is true, the code is going below
        # fused_add_norm_fn = layer_norm_fn
        ###########
        # hidden_states: Tensor
        # self.norm.weight = torch.nn.LayerNorm.weight
        # self.norm.bias = torch.nn.LayerNorm.bias
        # residual: Optional[Tensor] = None
        # prenorm=True
        # residual_in_fp32 = True
        # eps = torch.nn.LayerNorm.eps
        
        else:
            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm, RMSNorm) else layer_norm_fn
            if residual is None:
                hidden_states, residual = fused_add_norm_fn(
                    hidden_states,
                    self.norm.weight,
                    self.norm.bias,
                    residual=residual,
                    prenorm=True,
                    residual_in_fp32=self.residual_in_fp32,
                    eps=self.norm.eps,
                )
            else:
                hidden_states, residual = fused_add_norm_fn(
                    self.drop_path(hidden_states),
                    self.norm.weight,
                    self.norm.bias,
                    residual=residual,
                    prenorm=True,
                    residual_in_fp32=self.residual_in_fp32,
                    eps=self.norm.eps,
                )  

        # inference_params=None
        hidden_states = self.mixer(hidden_states, inference_params=inference_params)
        return hidden_states, residual

    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
        print("the code is going through allocate_inference_cache in the block")

        return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)

'''
from torch.nn.ModuleList()

embed_dim=384, (embedding dimension)
device: None 
dtype: None
ssm_cfg: None
norm_epsilon: float = 1e-5
rms_norm: bool = False
residual_in_fp32 = True
fused_add_norm = True
if_bimamba = False
bimamba_type = "v2"
inter_dpr: [0.0, 0.0, 0.004347826354205608, ...,0.09565217792987823, 0.10000000149011612]
if_devide_out = True
init_layer_scale = None
layer_idx = i
'''

def create_block(
    d_model,
    ssm_cfg=None,
    norm_epsilon=1e-5,
    drop_path=0.,
    rms_norm=False,
    residual_in_fp32=False,
    fused_add_norm=False,
    layer_idx=None,
    device=None,
    dtype=None,
    if_bimamba=False,
    bimamba_type="none",
    if_devide_out=False,
    init_layer_scale=None,
    block_name = "default_value"
):
    if if_bimamba:
        bimamba_type = "v1"

    if ssm_cfg is None:
        ssm_cfg = {}

    factory_kwargs = {"device": device, "dtype": dtype}

    if block_name == "VisionMamba":
        mixer_cls = partial(Mamba, layer_idx=layer_idx, bimamba_type=bimamba_type, if_devide_out=if_devide_out, init_layer_scale=init_layer_scale, **ssm_cfg, **factory_kwargs)
    elif block_name == "OriginalMamba":
        mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)
    else:
        raise ValueError(f"No matching condition for value: {block_name}")
    
    

    # rms_norm = False
    norm_cls = partial(nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs)
    
    block = Block(
        d_model,
        mixer_cls,
        norm_cls=norm_cls,
        drop_path=drop_path,
        fused_add_norm=fused_add_norm,
        residual_in_fp32=residual_in_fp32,
    )
    block.layer_idx = layer_idx
    return block


# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
def _init_weights(
    module,
    n_layer,
    initializer_range=0.02,  # Now only used for embedding layer.
    rescale_prenorm_residual=True,
    n_residuals_per_layer=1,  # Change to 2 if we have MLP
):
    if isinstance(module, nn.Linear):
        if module.bias is not None:
            if not getattr(module.bias, "_no_reinit", False):
                nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.normal_(module.weight, std=initializer_range)

    if rescale_prenorm_residual:
        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
        #
        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
        for name, p in module.named_parameters():
            if name in ["out_proj.weight", "fc2.weight"]:
                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
                # We need to reinit p since this code could be called multiple times
                # Having just p *= scale would repeatedly scale it down
                nn.init.kaiming_uniform_(p, a=math.sqrt(5))
                with torch.no_grad():
                    p /= math.sqrt(n_residuals_per_layer * n_layer)


def segm_init_weights(m):
    if isinstance(m, nn.Linear):
        trunc_normal_(m.weight, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.Conv2d):
        # NOTE conv was left to pytorch default in my original init
        lecun_normal_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
        nn.init.zeros_(m.bias)
        nn.init.ones_(m.weight)

'''
below is the 'ft-vim-s.sh':

patch_size=16, (just patch size) 
stride=8,  (just stride)
embed_dim=384, (embedding dimension)
depth=24, (?) the number of block
rms_norm = True, (?)
residual_in_fp32 = True, (?) 
fused_add_norm = True, (?)
final_pool_type = 'mean', (?)
if_abs_pos_embed = True, (?)
if_rope = False, (?)
if_rope_residual = False, (?) 
bimamba_type = "v2", (?)
if_cls_token = True, (?)
if_devide_out = True, (?)
use_middle_cls_token = True, 
**kwargs
'''

class VisionMamba(nn.Module):
    def __init__(self, 
                 img_size=224, 
                 patch_size=16, 
                 stride=16,
                 depth=24, 
                 embed_dim=192, 
                 channels=3, 
                 num_classes=26,
                 ssm_cfg=None, 
                 drop_rate=0.,
                 drop_path_rate=0,
                 norm_epsilon: float = 1e-5, 
                 rms_norm: bool = False, 
                 initializer_cfg=None,
                 fused_add_norm=False,
                 residual_in_fp32=False,
                 device=None,
                 dtype=None,
                 ft_seq_len=None,
                 pt_hw_seq_len=14,
                 if_bidirectional=False,
                 final_pool_type='none',
                 if_abs_pos_embed=False,
                 if_rope=False,
                 if_rope_residual=False,
                 flip_img_sequences_ratio=-1.,
                 if_bimamba=False,
                 bimamba_type="none",
                 if_cls_token=False,
                 if_devide_out=False,
                 init_layer_scale=None,
                 use_double_cls_token=False,
                 use_middle_cls_token=False,
                 **kwargs):
        # print("The program is coming the init")
        factory_kwargs = {"device": device, "dtype": dtype}
        # factory_kwargs: {'device': None, 'dtype': None}
        # add factory_kwargs into kwargs
        block_name = kwargs.get('block', 'default_value')

        kwargs.update(factory_kwargs) 

        super().__init__()
        self.residual_in_fp32 = residual_in_fp32
        self.fused_add_norm = fused_add_norm
        self.if_bidirectional = if_bidirectional
        self.final_pool_type = final_pool_type
        self.if_abs_pos_embed = if_abs_pos_embed
        self.if_rope = if_rope
        self.if_rope_residual = if_rope_residual
        self.flip_img_sequences_ratio = flip_img_sequences_ratio
        self.if_cls_token = if_cls_token
        self.use_double_cls_token = use_double_cls_token
        self.use_middle_cls_token = use_middle_cls_token
        self.num_tokens = 1 if if_cls_token else 0

        # pretrain parameters
        self.num_classes = num_classes
        self.d_model = self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models

        # self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, stride=stride, in_chans=channels, embed_dim=embed_dim)
        # num_patches = self.patch_embed.num_patches

        self.CNN_layers = CNN_layers()
        num_patches = 729

        # self.ECG_patch_embedding = ECG_Patch_embedding()
        # num_patches = 1023
        # self.LC = LeadCombiner(lead=6, out_ch=8)

        # self.CNN_layers = D2_CNN_layers()
        # num_patches = 828

        # self.CNN_layers = CNN_layers()
        # num_patches = 775
        
        # if_cls_token: True (in the original script)
        if if_cls_token:
            # use_double_cls_token: False (in the original script)
            if use_double_cls_token:
                self.cls_token_head = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
                self.cls_token_tail = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
                self.num_tokens = 2

            else:
                self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
                # self.num_tokens = 1
        
        # if_abs_pos_embed: True (in the original script)
        if if_abs_pos_embed:
            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, self.embed_dim))
            self.pos_drop = nn.Dropout(p=drop_rate)

        # if_rope: False (in the original script)

        if if_rope:
            half_head_dim = embed_dim // 2
            self.rope = TimeSeriesRotaryEmbeddingFast(dim=embed_dim, seq_len= num_patches + self.num_tokens)
        
        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
        # self.head_LC = nn.Linear(16, num_classes)
        # depth: 24; drop_path_rate: 0.1
        # TODO: release this comment
        # dpr: [0.0, 0.004347826354205608, 0.008695652708411217, ..., 0.09130434691905975, 0.09565217792987823, 0.10000000149011612]

        if drop_path_rate == 0:
            print("This is the drop_path_rate:", drop_path_rate)
            dpr = [x.item() for x in torch.full((depth,), drop_path_rate)]            
        else:
            print("This is the drop_path_rate:", drop_path_rate)
            print("follow the stochastic depth decay rule")
            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule


        # import ipdb;ipdb.set_trace()
        # inter_dpr: [0.0, 0.0, 0.004347826354205608, ...,0.09565217792987823, 0.10000000149011612]
        inter_dpr = [0.0] + dpr
        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()

        # transformer blocks
        # depth: 
        self.layers = nn.ModuleList(
            [
                create_block(
                    embed_dim,
                    ssm_cfg=ssm_cfg,
                    norm_epsilon=norm_epsilon,
                    rms_norm=rms_norm,
                    residual_in_fp32=residual_in_fp32,
                    fused_add_norm=fused_add_norm,
                    layer_idx=i,
                    if_bimamba=if_bimamba,
                    bimamba_type=bimamba_type,
                    drop_path=inter_dpr[i],
                    if_devide_out=if_devide_out,
                    init_layer_scale=init_layer_scale,
                    block_name = block_name,
                    **factory_kwargs,
                )
                for i in range(depth)
            ]
        )
        
        # output head
        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(embed_dim, eps=norm_epsilon, **factory_kwargs)

        # self.pre_logits = nn.Identity()

        # original init
        # self.patch_embed.apply(segm_init_weights)

        self.CNN_layers.apply(segm_init_weights)

        # self.ECG_patch_embedding.apply(segm_init_weights)

        self.head.apply(segm_init_weights)

        # self.head_LC.apply(segm_init_weights)
        # self.LC.apply(segm_init_weights)


        # if_abs_pos_embed: True (in the original script)
        if if_abs_pos_embed:
            trunc_normal_(self.pos_embed, std=.02)
        # if_cls_token: True (in the original script)
        if if_cls_token:
            if use_double_cls_token:
                trunc_normal_(self.cls_token_head, std=.02)
                trunc_normal_(self.cls_token_tail, std=.02)
            # the code is coming here
            else:
                trunc_normal_(self.cls_token, std=.02)

        # mamba init
        self.apply(partial(_init_weights, n_layer=depth, **(initializer_cfg if initializer_cfg is not None else {}),))


    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
        print("the code is going through allocate_inference_cache in the vision mamba")
        return {
            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
            for i, layer in enumerate(self.layers)
        }

    @torch.jit.ignore
    def no_weight_decay(self):
        return {"pos_embed", "cls_token", "dist_token", "cls_token_head", "cls_token_tail"}

    @torch.jit.ignore()
    def load_pretrained(self, checkpoint_path, prefix=""):
        _load_weights(self, checkpoint_path, prefix)

    # x: this is the input 224*224(tensor)
    # inference_params: None
    # if_random_cls_token_position：False
    # if_random_token_rank: False

    def forward_features(self, x, inference_params=None, if_random_cls_token_position=False, if_random_token_rank=False):
        # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
        # with slight modifications to add the dist_token
        
        # x = self.patch_embed(x)

        # x = self.CNN_layers(x)

        x = self.CNN_layers(x)

        # B: batch size 16
        # M: 729 the number of patch,(27*27)
        # D: the hidden state dimension, 384 (small-size variant), this is set by author of Vim, it is 768 in the convential Vit
        # N: SSM dimension, SSM dimension N to 16.
        # L: the number of blocks, we set the number of blocks L to 24

        B, M, _ = x.shape
        
        # if_cls_token: True (in the original script)
        if self.if_cls_token:

            # self.use_double_cls_token: False (in the original script)
            if self.use_double_cls_token:
                cls_token_head = self.cls_token_head.expand(B, -1, -1)
                cls_token_tail = self.cls_token_tail.expand(B, -1, -1)
                token_position = [0, M + 1]
                x = torch.cat((cls_token_head, x, cls_token_tail), dim=1)
                M = x.shape[1]
            else:
                # self.use_middle_cls_token: True(in the original script)
                if self.use_middle_cls_token:
                    cls_token = self.cls_token.expand(B, -1, -1)
                    token_position = M // 2
                    # add cls token in the middle
                    x = torch.cat((x[:, :token_position, :], cls_token, x[:, token_position:, :]), dim=1)
                elif if_random_cls_token_position:
                    cls_token = self.cls_token.expand(B, -1, -1)
                    token_position = random.randint(0, M)
                    x = torch.cat((x[:, :token_position, :], cls_token, x[:, token_position:, :]), dim=1)
                    print("token_position: ", token_position)
                else:
                    cls_token = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
                    token_position = 0
                    x = torch.cat((cls_token, x), dim=1)
                M = x.shape[1]

        # # if_abs_pos_embed: True (in the original script)
        if self.if_abs_pos_embed:
            # if new_grid_size[0] == self.patch_embed.grid_size[0] and new_grid_size[1] == self.patch_embed.grid_size[1]:
            #     x = x + self.pos_embed
            # else:
            #     pos_embed = interpolate_pos_embed_online(
            #                 self.pos_embed, self.patch_embed.grid_size, new_grid_size,0
            #             )
            x = x + self.pos_embed
            x = self.pos_drop(x)



        if_flip_img_sequences = False
        if self.flip_img_sequences_ratio > 0 and (self.flip_img_sequences_ratio - random.random()) > 1e-5:
            x = x.flip([1])
            if_flip_img_sequences = True

        # mamba impl
        # if_bidirectional： false
        # inference_params： None
        residual = None
        hidden_states = x
        if not self.if_bidirectional:
            for layer in self.layers:
                
                # here is false in the original script
                if if_flip_img_sequences and self.if_rope:
                    hidden_states = hidden_states.flip([1])
                    if residual is not None:
                        residual = residual.flip([1])

                # rope about, defaule is false
                if self.if_rope:
                    hidden_states = self.rope(hidden_states)
                    if residual is not None and self.if_rope_residual:
                        residual = self.rope(residual)

                # here is false in the original script
                if if_flip_img_sequences and self.if_rope:
                    hidden_states = hidden_states.flip([1])
                    if residual is not None:
                        residual = residual.flip([1])

                hidden_states, residual = layer(hidden_states, residual, inference_params=inference_params)
                # sys.exit()

        else:
            # get two layers in a single for-loop
            for i in range(len(self.layers) // 2):
                if self.if_rope:
                    hidden_states = self.rope(hidden_states)
                    if residual is not None and self.if_rope_residual:
                        residual = self.rope(residual)

                hidden_states_f, residual_f = self.layers[i * 2](
                    hidden_states, residual, inference_params=inference_params
                )
                hidden_states_b, residual_b = self.layers[i * 2 + 1](
                    hidden_states.flip([1]), None if residual == None else residual.flip([1]), inference_params=inference_params
                )
                hidden_states = hidden_states_f + hidden_states_b.flip([1])
                residual = residual_f + residual_b.flip([1])
        
        # fused_add_norm: True

        if not self.fused_add_norm:
            if residual is None:
                residual = hidden_states
            else:
                residual = residual + self.drop_path(hidden_states)
            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
        else:
            # Set prenorm = False here since we don't need the residual
            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn

            hidden_states = fused_add_norm_fn(
                self.drop_path(hidden_states),
                self.norm_f.weight,
                self.norm_f.bias,
                eps=self.norm_f.eps,
                residual=residual,
                prenorm=False,
                residual_in_fp32=self.residual_in_fp32,
            )

        # return only cls token if it exists
        # if_cls_token: True (in the original script)
        # self.use_middle_cls_token: True
        if self.if_cls_token:
            if self.use_double_cls_token:
                return (hidden_states[:, token_position[0], :] + hidden_states[:, token_position[1], :]) / 2
            else:
                if self.use_middle_cls_token:
                    return hidden_states[:, token_position, :]
                elif if_random_cls_token_position:
                    return hidden_states[:, token_position, :]
                else:
                    return hidden_states[:, token_position, :]

        # self.final_pol_type = 'mean'
        if self.final_pool_type == 'none':
            return hidden_states[:, -1, :]
        elif self.final_pool_type == 'mean':
            return hidden_states.mean(dim=1)
        elif self.final_pool_type == 'max':
            return hidden_states
        elif self.final_pool_type == 'all':
            return hidden_states
        else:
            raise NotImplementedError

    def forward(self, x, return_features=False, inference_params=None, if_random_cls_token_position=False, if_random_token_rank=False):
        x = self.forward_features(x, inference_params, if_random_cls_token_position=if_random_cls_token_position, if_random_token_rank=if_random_token_rank)
        # batch_number = x.shape[0]
        # x = x.view(batch_number, 8, 6, 8)
        # x = self.LC(x)
        # x = self.head_LC(x)
        # print("This is the shape of X (after the fully connected layer):", x.shape)
        # return_features = False 
        # print("This is the return feature:", return_features)
        # if return_features:
        #     return x
        
        # print("This is the shape of X (Before the fully connected layer):", x.shape)
        x = self.head(x)
        
        # final_pool_type = 'mean' in original script
        if self.final_pool_type == 'max':
            x = x.max(dim=1)[0]
        # sys.exit()
        return x



# below is for the vision in mamba
@register_model
def ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, depth=5, fused_add_norm = True, drop_path_rate = 0.1, if_divide_out = True, use_middle_cls_token = True, **kwargs):
    model = VisionMamba(patch_size=16, stride=8, embed_dim=384, depth=depth, rms_norm=True, residual_in_fp32=True, drop_path_rate = drop_path_rate, fused_add_norm = fused_add_norm, final_pool_type='mean', if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type="v2", if_cls_token=True, if_devide_out=if_divide_out, use_middle_cls_token=use_middle_cls_token, **kwargs)
    
    # As a reminder:
    print("This is whether the fused_add_norm:", fused_add_norm)
    print("This is whether the if_divide_out:", if_divide_out)
    print("This is whether the use_middle_cls_token:", use_middle_cls_token)


    model.default_cfg = _cfg()
    
    return model

# below is for the original mamba and 24 blocks
# @register_model
# def ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, **kwargs):
#     model = VisionMamba(patch_size=16, stride=8, embed_dim=384, depth=24, rms_norm=True, residual_in_fp32=True, fused_add_norm=True, final_pool_type='mean', if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type="v2", if_cls_token=True, if_devide_out=True, use_middle_cls_token=True, **kwargs)
#     model.default_cfg = _cfg()
    
#     return model